DINO的模型并行

最新推荐文章于 2026-01-10 15:05:49 发布

原创

最新推荐文章于 2026-01-10 15:05:49 发布 · 460 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #ubuntu #人工智能

参考链接

使用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,
        )