pytorch中正确使用GPU进行训练

【以下为单卡训练过程】

首先看下你的设备有无cuda可用：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Assuming that we are on a CUDA machine, this should print a CUDA device:

print(device)

如果有应该输出：

cuda:0

然后在定义的model/Net后面加上.to(device) ，这样这些方法将递归遍历所有模块，并将其参数和缓冲区转为CUDA tensors：

model = ConvNet(num_classes).to(device)    # 如前面已定义class ConvNet(nn.Module)

记得也要在每一个step在inputs和targets（标签）后面加上.to(device)：

images = images.to(device)
labels = labels.to(device)

作者：MeWTwo
链接：https://www.zhihu.com/question/345418003/answer/926213163