GPU accelerated

[155****7220]

可以使用该代码来查看当前环境是否支持CUDA

import torch
print(torch.cuda.is_available())
# 返回True代表支持，False代表不支持

具体在神经网络中，该这样使用

import torch
import torch.nn as nn
import torch.optim as optim

device = torch.device('cuda:0')
# 'cuda:0'当中的0为想要使用显卡的编号
# 这里的0表示使用的是第一张显卡
net = MLP().to(device)
# 使用.to函数将神经网络模块搬到MLP上进行运算
optimizer = optim.SGD(net.parameters(), lr=1e-3)
criteon = nn.CrossEntropyLoss().to(device)
# 将loss部分的计算转移到GPU上去

同样地，数据部分也可以转移到GPU上去

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