Pytorch 多块GPU的使用

本文详细介绍了如何在单服务器多GPU环境下使用Pytorch进行深度学习模型的并行训练,包括设置GPU编号、迁移网络及数据至GPU、解决CPU与GPU冲突问题以及在多GPU上执行前向传播等关键步骤。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

注:本文针对单个服务器上多块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都被使用啦

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

付费内容不再更新,请勿购买!!

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值