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上了，但是發(fā)現不轉移到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
實時查看gpu信息1代表每1秒刷新一次
指定使用的gpu
import os
# 使用第一張與第三張GPU卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"
以上為個人經驗，希望能給大家一個參考，也希望大家多多支持腳本之家。
您可能感興趣的文章:- pytorch 兩個GPU同時訓練的解決方案
- pytorch使用horovod多gpu訓練的實現
- 解決pytorch多GPU訓練保存的模型,在單GPU環(huán)境下加載出錯問題
- pytorch 指定gpu訓練與多gpu并行訓練示例
- 用Pytorch訓練CNN(數據集MNIST,使用GPU的方法)
- pytorch 使用單個GPU與多個GPU進行訓練與測試的方法
- pytorch使用指定GPU訓練的實例
- 關于pytorch多GPU訓練實例與性能對比分析
