Pytorch使用DataParallel后仍只使用一个GPU

问题

服务器类型：slurm管理下的超算服务器
系统：linux
python版本：3.8
虚拟环境：virtualenv
pytorch版本：1.10
问题描述：在服务器上用上卡跑模型时，使用

model = Model()
model = torch.nn.DataParallel(model).cuda()

后，模型和数据仍然只占用第一个gpu，然后内存爆掉。

解决过程

查阅https://github.com/openai/CLIP/issues/111后得知，对于clip类模型，即在模型训练过程中会调用子模块的情况，比如下面代码。调用子模块的部分不会在多gpu上运行。

model = Model()
model = torch.nn.DataParallel(model).cuda()
img = batch['img']
prompt = batch['prompt']
img_feat = model.module.image_encoder(img)
prompt_feat = model.module.prompt_encoder(prompt)

解决方案给子模块包上nn.module，然后再加上dataparallel，如下示例

import torch
import torch.nn as nn
import torch.optim as optim
import os 
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

class BigModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layer1 = nn.Conv2d(3,1,1,1,0)
        self.layer2 = nn.Conv2d(1,1,1,1,0)
    def forward(self,img):
        out = self.layer1(img)
        out = self.layer2(out)
        return out


class SplitModel(nn.Module):
    def __init__(self,model) -> None:
        super().__init__()
        self.model = model
    def forward(self,img):
        return self.model.layer1(img)


model = BigModel().cuda()
split_model = SplitModel(model)
split_model = nn.DataParallel(split_model).cuda()
x = torch.randn(6,3,64,64).float().cuda() # batch['img'].float().cuda()
target = torch.ones(6,1,64,64).float().cuda() # batch['label'].float().cuda()

loss_func = nn.BCELoss()

model.train()
pred = split_model(x)
pred = torch.sigmoid(pred)
loss = loss_func(pred,target)

此处只想使用BigModel类的layer2，所以单独创建了一个SplitModel类。