Pytorch多GPU的计算和SyncBatchNorm_nn.SyncBatchNorm - CSDN文库

32 浏览量更新于2023-05-04 评论 1 收藏 84KB PDF 举报

身份认证购VIP最低享 7 折!

领优惠券(最高得80元）

资源详情

资源评论

资源推荐

Pytorch多多GPU的计算和的计算和Sync BatchNorm

nn.DataParallel

pytorch中使用GPU非常方便和简单：

import torch

import torch.nn as nn

input_size = 5

output_size = 2

class Model(nn.Module):

def __init__(self, input_size, output_size):

super(Model, self).__init__()

self.fc = nn.Linear(input_size, output_size)

def forward(self, input):

output = self.fc(input)

print("[In Model]: device",torch.cuda.current_device() ," input size", input.size()," output size", output.size())

return output

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

model = Model(input_size, output_size)

model.to(device)

x = torch.Tensor(2,5)

x = x.to(device)

y = model(x)

这里需要注意的是，仅仅调用Tensor.to()只会在GPU上返回一个新的copy，并不会对原来的引用造成变化，因此需要通过赋值

rewrite。

上述只是对单个GPU的使用方法，对于多个GPU，pytorch也提供了封装好的接口——DataParallel,只需要将model 对象放入容器

中即可：

model = Model(input_size, output_size)

print("Let's use", torch.cuda.device_count(), "GPUs!")

model = nn.DataParallel(model)

model.to(device)

print(model)

# output

Let's use 2 GPUs!

DataParallel(

(module): Model(

(fc): Linear(in_features=5, out_features=2, bias=True)

)

)

看到这次输出的model外面还有一层DataParallel，但这里并没有体现出存在多个GPU。

接下来构造一个Dummy DataSet，来跑一下模型：

from torch.utils.data import Dataset, DataLoader

batch_size = 30

data_size = 100

class RandomDataset(Dataset):

def __init__(self, size, length):

self.len = length

self.data = torch.randn(length, size) # 有length个样本，每个样本是size长度的向量

本内容试读结束，登录后可阅读更多

下载后可阅读完整内容，剩余3页未读，立即下载

评论0

weixin_38628552

粉丝: 3
资源: 907

会员权益专享

图片转文字

全年可省5，000元立即开通

最新资源

资源上传下载、课程学习等过程中有任何疑问或建议，欢迎提出宝贵意见哦~我们会及时处理！点击此处反馈