pytorch的代码 CPU改为GPU
找一下几个地方改动
一、网络模型
- 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()
-
湘公网安备 43102202000103号
