利用 GPU加速训练

2、1、网络模型、损失函数、优化器、数据(图像,标签)可以.cuda,加速训练。

if torch.cuda.is_avaliable():
    img = img.cuda()

2、用to的方法:

        device = torch.device("cuda:0")

        img = img.to(device)

3、GPU上训练的模型,需要直接进行测试

model = torch.load("model.pth", map_location = torch.device('cpu'))

你可能感兴趣的:(神经网络)