Pytorch 使用GPU加速

Pytorch中使用torch.device()选取并返回抽象出的设备,然后在定义的网络模块或者Tensor后面加上.to(device变量)就可以将它们搬到设备上了。

以上一篇代码为例,使用GPU设备:

device = torch.device('cuda:0') # 使用第一张显卡

需要将如下部分搬移到GPU上:

1. 定义的网络

net = MLP().to(device)

2. 损失函数

criteon = nn.CrossEntropyLoss()

3.每次取出的训练集和验证集的batch

data, target = data.to(device), target.to(device)

 

你可能感兴趣的:(深度学习,pytorch)