PyTorch学习笔记：多GPU训练

 方法一：

# 选择GPU执行 or CPU执行，方法一
if torch.cuda.is_available():
    AEmodel = AutoEncoder()
    AEmodel = torch.nn.DataParallel(AEmodel).cuda()
else:
    AEmodel = AutoEncoder()

方法二：

# 选择GPU执行 or CPU执行，方法二
# 定义device，其中需要注意的是“cuda:0”代表起始的device_id为0，如果直接是“cuda”，同样默认是从0开始。可以根据实际需要修改起始位置，如“cuda:1”
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
AEmodel = AutoEncoder()
if torch.cuda.device_count() > 1:
    AEmodel = nn.DataParallel(AEmodel, device_ids=[0, 1, 2])

AEmodel.to(device)

 查看GPU是否可用及输出网络结构：

print("torch.cuda.is_available：", torch.cuda.is_available())
print(AEmodel)