pytorch多GPU训练方式 Multi GPU

最新推荐文章于 2023-02-02 15:49:27 发布

Tsingzao-于廷照

最新推荐文章于 2023-02-02 15:49:27 发布

阅读量1.5k

点赞数

CC 4.0 BY-SA版权

分类专栏： Pytorch Linux

本文链接：https://blog.youkuaiyun.com/yutingzhaomeng/article/details/102898147

Pytorch 同时被 2 个专栏收录

20 篇文章

订阅专栏

Linux

16 篇文章

订阅专栏

本文详细介绍如何使用PyTorch的DataParallel进行分布式模型训练，包括模型定义、数据并行处理、GPU资源分配及常见问题解析。

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

1、首先定义模型

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        pass
    def forward(self, x):
        pass

2、利用DataParallel分布式存储数据和模型

gpuID = '0,1,2,3,4,5,6,7'
device_ids = [int(id) for id in gpuID.split(',')]

import os
os.environ['CUDA_VISIBLE_DEVICES'] = gpuID

device = torch.device('cuda: 0')

import torch
model = Net()
model = torch.nn.DataParallel(model, device_ids=device_ids)
model.to(device)
criterion.to(device)

input, output = input.to(device), output.to(device)

output_ = model(input)

loss = criterion(output, output_)

3、查看nvidia-smi命令即可发现模型训练已经自动分配在全部显卡上

4、貌似指定偶数块GPU没问题，实验过程中，如果指定奇数块GPU，如4，5，6，发现模型只占用了4，5号GPU；而指定4，5，6，7时，模型会占用4，5，6，7号GPU。没搞懂原因。搞懂了，跟batchsize有关。