使用GPU进行模型训练和预测的两种方式

需要转移到CUDA的有：

• 网络模型
• 损失函数
• 数据集

方法一：cuda()

#网络模型转移到cuda
mymodel = MyModel()
if torch.cuda.is_available():
    mymodel = mymodel.cuda()
 
#损失函数转移到cuda
loss = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss.cuda()

#数据集转移到cuda
imgs, targets = data
if torch.cuda.is_available():
	imgs = imgs.cuda()
    targets=targets.cuda()

方法二：to(device)

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

#网络模型转移到cuda
mymodel = MyModel()
mymodel.to(device)
 
#损失函数转移到cuda
loss = nn.CrossEntropyLoss()
loss.to(device)

#数据集转移到cuda
imgs, targets = data
imgs = imgsto(device)
targets=targets.to(device)