GPU accelerated
-
可以使用该代码来查看当前环境是否支持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)