DINO的模型并行

参考链接

使用DINO训练自己的数据集
DINO源码
resnet50模型并行

问题介绍

一张超高分辨率图片送入DINO进行检测,会导致爆显存,只要分别把模型放入不同gpu,再让图片分别经过这几个gpu就能解决,这就是模型并行

解决办法

化繁为简:先将resnet50拆开实验,来源于第三个链接

"""
运行resnet50,5卡,要拆到具体结构了.
"""
from torchvision.models.resnet import ResNet, Bottleneck
import torch.nn as nn
import torch
import torch.optim as optim
from tqdm import tqdm

class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=2, *args, **kwargs)

        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
        ).to('cuda:0') # 移动到forward里面也是ok

        self.seq2 = nn.Sequential(
            self.layer1,
        ).to('cuda:1')

        self.seq3 = nn.Sequential(
            self.layer2,
        ).to('cuda:2')

        self.seq4 = nn.Sequential(
            self.layer3,
        ).to('cuda:3')

        self.seq5 = nn.Sequential(
            self.layer4,
            self.avgpool,
        ).to('cuda:4')

        # 替换全连接层以适应新的输出类别数
        self.fc.to('cuda:4')

    def forward(self, x):
        x = self.seq1(x.to('cuda:0'))  # 第一层在cuda:0上计算
        x = x.to('cuda:1')  # 将结果转移到cuda:1上
        x = self.seq2(x)  # 第二层在cuda:1上计算

        x = x.to('cuda:2')  # 将结果转移到cuda:2上
        x = self.seq3(x)  # 第三层在cuda:2上计算

        x = x.to('cuda:3')  # 将结果转移到cuda:3上
        x = self.seq4(x)  # 第四层在cuda:3上计算

        x = x.to('cuda:4')  # 将结果转移到cuda:3上
        x = self.seq5(x)  # 第四层在cuda:3上计算

        x = self.fc(x.view(x.size(0), -1))  # 全连接层在cuda:3上计算
        return x

def print_gpu_memory_usage():
    # 获取当前GPU的显存占用信息
    torch.cuda.synchronize()  # 确保所有之前的操作都已完成
    allocated = torch.cuda.memory_allocated()  # 已分配的显存
    cached = torch.cuda.memory_reserved()  # 保留的显存（包括缓存）

    print(f"Allocated: {allocated / 1024 ** 3:.2f} GB, Cached: {cached / 1024 ** 3:.2f} GB")
    
def train(model, epochs, batch_size, num_classes, imgsz):
    model.train()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    for _ in tqdm(range(epochs), desc='Training Epochs'):
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, imgsz, imgsz) # 8192
        labels = torch.randint(0, num_classes, (batch_size,))

        # run forward pass
        optimizer.zero_grad()
        outputs = model(inputs)

        # run backward pass
        loss = loss_fn(outputs, labels.to('cuda:4')) # output在cuda:1上
        loss.backward()
        optimizer.step()

        # 显存占用
        print_gpu_memory_usage()

if __name__ == '__main__':
    
    num_classes = 2
    epochs = 100
    batch_size = 1
    imgsz = 4480 # 5卡 224,1100? 2240,13.3G? 1120,4.1G?
    modelParallelResNet50 = ModelParallelResNet50()
    train(modelParallelResNet50, epochs, batch_size, num_classes, imgsz)

分而治之(感觉更像并行,但百度释义是串行?):跑通DINO(参考第一个链接即可,其中yolo转coco在文章评论里面),修改图片分辨率(非重点,修改coco_transformer.py即可,拆分网络结构(本文主要工作)
首先,注释掉main.py的model.to(device),把model放入单一的设备是不利于后续的拆分,同时把命令行参数补齐,便于vscode调试.
然后就是修改涉及模型模块的具体代码,dino.py,utils.py,deformable_transformer.py.

        self.transformer = transformer.to('cuda:0') # 单独模块
        self.label_enc = nn.Embedding(dn_labelbook_size + 1, hidden_dim).to('cuda:0') # 单独模块
        self.input_proj = nn.ModuleList(input_proj_list).to('cuda:0') # 单独模块
        self.backbone = backbone.to('cuda:0') # 单独模块
        self.class_embed.to('cuda:0') # 模块放入cuda和前向传播区别的差异性?

        self.layers.to('cuda:0') # 增加和修改区别不大,一个便于查看,一个便于修改

        self.enc_out_class_embed.to('cuda:0') # 实在找不到定义的位置

最后,重点关注诸如dn_components.py里面的.cuda部分,这里其实不应该这么写,因为直接相当于指定了cuda,应该写成device比较好,能够动态变化,或者就不需要device什么事,全部一个个指定就好.

m = known_labels_expaned.long().to('cuda') #  整个项目类似地方还有很多

结语

具体替换的代码很简单,我放在博客里面了,有兴趣的直接替换就好了,DINO模型并行.
大家有兴趣的可以一起讨论,其实还存在很多问题,第一个问题:比如拆分作用有待商榷?因为21G+20G?单卡也才23G呀.第2个问题:在执行完1个epoch时,执行test时,GPU不能访问内存.第三个问题:这么做有啥意义?(⁄ ⁄•⁄ω⁄•⁄ ⁄)