pytorch的代码 CPU改为GPU

pytorch的代码 CPU改为GPU

找一下几个地方改动

• 网络模型
• 数据（输入、标记）
• 损失函数
• cuda()

一、网络模型

mymodule = MyModule()
mymodule = mymodule.cuda()

二、数据（输入、标记）

        images,targets = data
        images = images.cuda()
        targets = targets.cuda() #训练数据和测试数据集都要进行此操作

三、损失函数

# 损失函数
loss = nn.CrossEntropyLoss()
loss = loss.cuda()

上面的写法，在无GPU时会报错。更好的写法是有GPU用GPU，无GPU用CPU

mymodule = MyModule()
if torch.cuda.is_available():
    mymodule = mymodule.cuda()