Pytorch 多块GPU的使用详解
作者:R.X.Geng 时间:2021-01-21 09:19:09
注:本文针对单个服务器上多块GPU的使用,不是多服务器多GPU的使用。
在一些实验中,由于Batch_size的限制或者希望提高训练速度等原因,我们需要使用多块GPU。本文针对Pytorch中多块GPU的使用进行说明。
1. 设置需要使用的GPU编号
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,4"
ids = [0,1]
比如我们需要使用第0和第4块GPU,只用上述三行代码即可。
其中第二行指程序只能看到第1块和第4块GPU;
第三行的0即为第二行中编号为0的GPU;1即为编号为4的GPU。
2.更改网络,可以理解为将网络放入GPU
class CNN(nn.Module):
def __init__(self):
super(CNN,self).__init__()
self.conv1 = nn.Sequential(
......
)
......
self.out = nn.Linear(Liner_input,2)
......
def forward(self,x):
x = self.conv1(x)
......
output = self.out(x)
return output,x
cnn = CNN()
# 更改,.cuda()表示将本存储到CPU的网络及其参数存储到GPU!
cnn.cuda()
3. 更改输出数据(如向量/矩阵/张量):
for epoch in range(EPOCH):
epoch_loss = 0.
for i, data in enumerate(train_loader2):
image = data['image'] # data是字典,我们需要改的是其中的image
#############更改!!!##################
image = Variable(image).float().cuda()
############################################
label = inputs['label']
#############更改!!!##################
label = Variable(label).type(torch.LongTensor).cuda()
############################################
label = label.resize(BATCH_SIZE)
output = cnn(image)[0]
loss = loss_func(output, label) # cross entropy loss
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # backpropagation, compute gradients
optimizer.step()
... ...
4. 更改其他CPU与GPU冲突的地方
有些函数必要在GPU上完成,例如将Tensor转换为Numpy,就要使用data.cpu().numpy(),其中data是GPU上的Tensor。
若直接使用data.numpy()则会报错。除此之外,plot等也需要在CPU中完成。如果不是很清楚哪里要改的话可以先不改,等到程序报错了,再哪里错了改哪里,效率会更高。例如:
... ...
#################################################
pred_y = torch.max(test_train_output, 1)[1].data.cpu().numpy()
accuracy = float((pred_y == label.cpu().numpy()).astype(int).sum()) / float(len(label.cpu().numpy()))
假如不加.cpu()便会报错,此时再改即可。
5. 更改前向传播函数,从而使用多块GPU
以VGG为例:
class VGG(nn.Module):
def __init__(self, features, num_classes=2, init_weights=True):
super(VGG, self).__init__()
... ...
def forward(self, x):
#x = self.features(x)
#################Multi GPUS#############################
x = nn.parallel.data_parallel(self.features,x,ids)
x = x.view(x.size(0), -1)
# x = self.classifier(x)
x = nn.parallel.data_parallel(self.classifier,x,ids)
return x
... ...
然后就可以看运行结果啦,nvidia-smi查看GPU使用情况:
可以看到0和4都被使用啦
来源:https://blog.csdn.net/qazwsxrx/article/details/89672578
标签:Pytorch,GPU
![](/images/zang.png)
![](/images/jiucuo.png)
猜你喜欢
php checkbox复选框值的获取与checkbox默认值输出方法
2023-11-14 13:41:12
Python实现给qq邮箱发送邮件的方法
2022-11-16 21:01:21
Python语言生成水仙花数代码示例
2022-11-16 18:17:48
![](https://img.aspxhome.com/file/2023/9/110439_0s.png)
基于Python中单例模式的几种实现方式及优化详解
2022-10-24 14:20:45
python神经网络学习使用Keras进行回归运算
2023-01-30 14:08:45
![](https://img.aspxhome.com/file/2023/9/90629_0s.png)
Langchain集成管理prompt功能详解
2022-12-13 22:56:31
在Python中表示一个对象的方法
2023-12-25 20:04:22
opencv读取视频并保存图像的方法
2023-02-21 17:30:16
![](https://img.aspxhome.com/file/2023/6/105296_0s.jpg)
python中的Reportlab模块详解最新推荐
2023-04-09 21:33:46
![](https://img.aspxhome.com/file/2023/8/66898_0s.png)
双屏显示提升前端开发10%工作效率
2009-03-16 18:22:00
![](https://img.aspxhome.com/file/UploadPic/20093/16/screen-34s.jpg)
Python Django框架url反向解析实现动态生成对应的url链接示例
2021-08-23 04:30:38
Python创建或生成列表的操作方法
2022-05-18 10:21:23
![](https://img.aspxhome.com/file/2023/2/112462_0s.jpg)
tensorflow之读取jpg图像长和宽实例
2023-10-12 11:57:41
![](https://img.aspxhome.com/file/2023/2/128582_0s.jpg)
Python 200行代码实现一个滑动验证码过程详解
2022-08-24 18:24:35
python中的turtle库函数简单使用教程
2022-08-06 23:22:18
![](https://img.aspxhome.com/file/2023/4/92054_0s.jpg)
Python实现的下载8000首儿歌的代码分享
2021-02-03 05:41:51
Python解析命令行读取参数之argparse模块
2021-12-17 22:15:32
Pandas中DataFrame交换列顺序的方法实现
2023-01-21 10:21:08
Python入门教程(一)Python简单介绍
2023-10-25 03:19:16
用python写个颜值评分器筛选最美主播
2022-11-25 22:13:08
![](https://img.aspxhome.com/file/2023/1/80731_0s.png)