服务器多张GPU同时跑代码

训练：

'''
定义device，其中需要注意的是“cuda:0”代表起始的device_id为0，如果直接是“cuda”，
同样默认是从0开始。可以根据实际需要修改起始位置，如“cuda:1”。
'''
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model(input_size, output_size)  # 实例化模型对象

# 多张GPU同时使用
if torch.cuda.device_count() > 1:  # 检查电脑是否有多块GPU
    print(f"Let's use {
      
      torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)  # 将模型对象转变为多GPU并行运算的模型
    # model = nn.DataParallel(model，device_ids=[0,1,2])

model.to(device)  # 把并行的模型移动到GPU上

测试：先并行，再加载模型

net=resnet50(pretrained=True)
in_channel=net.fc.in_features
net.fc = nn.Linear(in_channel, 10)
if torch.cuda.device_count() > 1:
    print(f"Let's use {
      
      torch.cuda.device_count()} GPUs!")
    net = nn.DataParallel(net)
net.load_state_dict(torch.load("resnet50state.pth"), strict=False)

引用：

【PyTorch教程】07-PyTorch如何使用多块GPU训练神经网络模型
pytorch使用多GPU训练