pytorch文档阅读(四)如何在GPU上训练

2022-10-03 12:27:51

1.网络模型转移到CUDA上

net = AlexNet()
net.cuda()#转移到CUDA上

2.将loss转移到CUDA上

criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()

这一步不做也可以,因为loss是根据out、label算出来的

loss = criterion(out, label)

只要out、label在CUDA上,loss自然也在CUDA上了,但是发现不转移到CUDA上准确率竟然降低了1%

3.将数据集转移到CUDA上

这里要解释一下数据集使用方法

#download the dataset
train_set = CIFAR10("./data_cifar10", train=True, transform=data_tf, download=True)
train_data = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

dataset是把所有的input,label都制作成了一个大的多维数组

dataloader是在这个大的多维数组里采样制作成batch,用这些batch来训练

    for im, label in train_data:
        i = i + 1
        im = im.cuda()#把数据迁移到CUDA上
        im = Variable(im)#把数据放到Variable里
        label = label.cuda()
        label =Variable(label)
        out = net(im)#the output should have the size of (N,10)

遍历batch的时候,首先要把拿出来的Image、label都转移到CUDA上,这样接下来的计算都是在CUDA上了

开始的时候只在转成Variable以后才迁移到CUDA上,这样在网络传播过程中就数据不是在CUDA上了,所以一直报错

参考链接:

https://blog.csdn.net/teeyohuang/article/details/79219368

https://www.cnblogs.com/zle1992/p/9047905.html

  • 作者:陌生的天花板
  • 原文链接:https://blog.csdn.net/weixin_41680653/article/details/93750326
    更新时间:2022-10-03 12:27:51