1.网络模型转移到CUDA上
|
1
2
|
net = AlexNet()net.cuda()#转移到CUDA上 |
2.将loss转移到CUDA上
|
1
2
|
criterion = nn.CrossEntropyLoss()criterion = criterion.cuda() |
这一步不做也可以，因为loss是根据out、label算出来的
|
1
|
loss = criterion(out, label) |
只要out、label在CUDA上，loss自然也在CUDA上了，但是发现不转移到CUDA上准确率竟然降低了1%
3.将数据集转移到CUDA上
这里要解释一下数据集使用方法
|
1
2
3
|
#download the datasettrain_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来训练
|
1
2
3
4
5
6
7
|
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
|
1
|
nvidia -smi |
实时查看gpu信息1代表每1秒刷新一次
|
1
|
watch -n -1 nvidia -smi |
指定使用的gpu
|
1
2
3
|
import os# 使用第一张与第三张GPU卡os.environ["CUDA_VISIBLE_DEVICES"] = "0,3" |
以上为个人经验，希望能给大家一个参考，也希望大家多多支持服务器之家。
原文链接：https://blog.csdn.net/weixin_41680653/article/details/93750326
