pytorch 如何在GPU上训练
短信预约 -IT技能 免费直播动态提醒
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上了,所以一直报错
训练网络时指定gpu显卡
查看有哪些可用的gpu
nvidia -smi
实时查看gpu信息1代表每1秒刷新一次
watch -n -1 nvidia -smi
指定使用的gpu
import os
# 使用第一张与第三张GPU卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"
以上为个人经验,希望能给大家一个参考,也希望大家多多支持编程网。
免责声明:
① 本站未注明“稿件来源”的信息均来自网络整理。其文字、图片和音视频稿件的所属权归原作者所有。本站收集整理出于非商业性的教育和科研之目的,并不意味着本站赞同其观点或证实其内容的真实性。仅作为临时的测试数据,供内部测试之用。本站并未授权任何人以任何方式主动获取本站任何信息。
② 本站未注明“稿件来源”的临时测试数据将在测试完成后最终做删除处理。有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341