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不能访问内存.第三个问题:这么做有啥意义?(⁄ ⁄•⁄ω⁄•⁄ ⁄)

### DINO模型概述 DINO(Vision Transformers with Self-Supervised Learning)是一种基于PyTorch的开源项目,专注于通过自监督学习方法训练视觉变换器模型[^1]。该模型无需依赖大量标注数据即可完成高效训练,并能应用于种计算机视觉任务。 #### 原理与实现方式 DINO的核心思想在于利用对比学习和动量更新机制来优化Transformer架构中的特征表示能力。具体而言: - **无监督预训练**:DINO采用了一种称为“teacher-student”的框架来进行无监督学习。其中,“student”网络负责提取输入图像的不同增强视图的特征向量;而“teacher”网络则通过对齐这些特征向量之间的相似度矩阵来指导“student”网络的学习过程。 - **动量更新策略**:为了稳定训练过程并减少过拟合风险,DINO引入了一个指数移动平均(EMA)技术对“teacher”参数进行渐进式的调整[^2]。 以下是简化版DINO算法的主要组成部分: ```python import torch.nn as nn class DINOLoss(nn.Module): def __init__(self, out_dim, ncrops, warmup_teacher_temp, teacher_temp, student_temp=0.1, center_momentum=0.9): super().__init__() self.student_temp = student_temp self.center_momentum = center_momentum # 初始化温度变量和其他超参... def forward(self, student_output, teacher_output): # 对齐学生和教师输出间的分布差异... def train_one_epoch(model, data_loader, optimizer, loss_fn, epoch): model.train() total_loss = 0 for images in data_loader: outputs = model(images) loss = loss_fn(outputs['student'], outputs['teacher']) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / len(data_loader) return avg_loss ``` #### 应用场景 由于其强大的泛化能力和灵活性,DINO可以广泛应用于以下领域: - 图像分类:得益于出色的特征表达能力,经过微调后的DINO能够在ImageNet等标准测试集上取得优异成绩。 - 图像检索:凭借全局上下文建模的优势,它能够更精准地匹配查询图片与其数据库中最接近的结果。 - 视频分割:当扩展到时间维度时,DINO同样表现出色,可有效捕捉动态物体的变化轨迹。 #### 源码获取与教程推荐 官方仓库地址位于GitHub平台之上,提供了详尽文档以及示例脚本帮助开发者快速入门。对于初学者来说,可以从阅读README文件开始熟悉整个项目的目录结构及其功能模块划分情况。此外,还有许第三方博客文章针对特定环节做了深入剖析,比如如何配置环境、运行实验或者调试常见错误等等[^4]。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值