pytorch深度学习框架–gpu和cpu的选择

基于pytorch框架，最近实现了一个简单的手写数字识别的程序，我安装的pytorch是gpu版(你也可以安装cpu版本的，根据个人需要)，这里我介绍pytorch的gpu版本和cpu版本的安装以及训练手写数字识别时gpu和cpu之间的切换。
1、pytorch的安装
1.1 pytorch(带有gpu)安装
首先进入pytorch官网，选择自己所需要的版本，这里我选择的版本如下图所示。
然后打开anaconda Prompt，首先输入：conda activate py3激活py3（解释一下为什么是py3，因为我之前装的是python3.6,创建的名字为py3），然后输入：conda install pytorch torchvision cudatoolkit=9.0 -c pytorch安装pytorch，等待安装就好，如下图所示。


1.2 pytorch(无gpu)安装
这时CUDA选择none即可

打开anaconda终端，首先激活py3,然后输入这个命令：conda install pytorch-cpu torchvision-cpu -c pytorch，等待安装就好，如下图所示，

1.3 测试是否安装成功
首先cmd打开终端，输入python即可查看当前安装的python的版本，然后import torch 等待几秒出现如下图所示，这样就成功安装了
2、选择cpu进行网络的训练（推荐下载带有gpu的）
2.1新建一个model.py模块

from torch import nnclass CNN(nn.Module):def __init__(self):super(CNN, self).__init__()# 使用序列工具快速构建self.conv1 = nn.Sequential(nn.Conv2d(1, 16, kernel_size=5, padding=2),  # 2？nn.BatchNorm2d(16),nn.ReLU(),nn.MaxPool2d(2))self.conv2 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=5, padding=2),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2))self.fc = nn.Linear(7 * 7 * 32, 10)  # ?def forward(self, x):out = self.conv1(x)out = self.conv2(out)out = out.view(out.size(0), -1)  # reshapeout = self.fc(out)return out

2.2 新建一个train.py模块（使用cpu训练的）

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import os
batch_size = 64
learning_rate = 0.001# 将数据处理成Variable, 如果有GPU, 可以转成cuda形式
def get_variable(x):x = Variable(x)return x.cpu() if torch.cuda.is_available() else xtrain_dataset = datasets.MNIST(root='./mnist/',train=True,transform=transforms.ToTensor(),download=True)transforms = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])train_loader = torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)images, labels = next(iter(train_loader))
img = torchvision.utils.make_grid(images)
img = img.numpy().transpose(1, 2, 0)
print(labels)
plt.imshow(img)
plt.show()# 两层卷积
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()# 使用序列工具快速构建self.conv1 = nn.Sequential(nn.Conv2d(1, 16, kernel_size=5, padding=2),  nn.BatchNorm2d(16),nn.ReLU(),nn.MaxPool2d(2))self.conv2 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=5, padding=2),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2))self.fc = nn.Linear(7 * 7 * 32, 10)  # ?def forward(self, x):out = self.conv1(x)out = self.conv2(out)out = out.view(out.size(0), -1)  # reshapeout = self.fc(out)return outcnn = CNN()
if torch.cuda.is_available():cnn = cnn.cpu()# 选择损失函数和优化方法
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
print(cnn)num_epochs = 2for epoch in range(num_epochs):for i, (images, labels) in enumerate(train_loader):images = get_variable(images)labels = get_variable(labels)  ##?outputs = cnn(images)optimizer.zero_grad()  ##loss = loss_func(outputs, labels)loss.backward()  # 反向传播，自动计算每个节点的锑度至optimizer.step()if (i + 1) % 100 == 0:print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f'% (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size, loss.item()))
torch.save(cnn.state_dict(), 'cnn.pkl')

2.3 新建一个test.py模块

import torch
import torchvision
import matplotlib.pyplot as plt
import torchvision.datasets as normal_datasets
import torchvision.transforms as transforms
from torch.autograd import Variable
from mymodel import CNN# 见数据加载器和batch
test_dataset = normal_datasets.MNIST(root='./mnist/',train=False,transform=transforms.ToTensor())data_loader_test=torch.utils.data.DataLoader(dataset=test_dataset,batch_size=4,shuffle=True)model = CNN()
model.load_state_dict(torch.load('cnn.pkl'))X_test, y_test = next(iter(data_loader_test))
inputs = Variable(X_test)
pred = model(inputs)
_, pred = torch.max(pred, 1)print("Predict Label is:", [i for i in pred.data])
print("Real Label is :", [i for i in y_test])img = torchvision.utils.make_grid(X_test)
img = img.numpy().transpose(1, 2, 0)plt.imshow(img)
plt.show()

3、选择gpu进行网络的训练
3.1 model.py模块不变，可以参考以上2.1
3.2 train.py模块（gpu训练）代码如下，可以对比以上2.2

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
batch_size = 64
learning_rate = 0.001# 将数据处理成Variable, 如果有GPU, 可以转成cuda形式
def get_variable(x):x = Variable(x)return x.cuda() if torch.cuda.is_available() else x# return nn.DataParallel(x, device_ids=[0])if torch.cuda.device_count() > 1 else x#如果有多个gpu时可以选择上面的语句，例如上面写的时设备0train_dataset = datasets.MNIST(root='./mnist/',train=True,transform=transforms.ToTensor(),download=True)transforms = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])train_loader = torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)images, labels = next(iter(train_loader))
img = torchvision.utils.make_grid(images)
img = img.numpy().transpose(1, 2, 0)
print(labels)
plt.imshow(img)
plt.show()# 两层卷积
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()# 使用序列工具快速构建self.conv1 = nn.Sequential(nn.Conv2d(1, 16, kernel_size=5, padding=2),  nn.BatchNorm2d(16),nn.ReLU(),nn.MaxPool2d(2))self.conv2 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=5, padding=2),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2))self.fc = nn.Linear(7 * 7 * 32, 10)  def forward(self, x):out = self.conv1(x)out = self.conv2(out)out = out.view(out.size(0), -1)  # reshapeout = self.fc(out)return outcnn = CNN()
if torch.cuda.is_available():cnn = cnn.cuda()
# if torch.cuda.device_count() > 1:
#     cnn = nn.DataParallel(cnn, device_ids=[0])
#如果多个gpu时，需要修改如上# 选择损失函数和优化方法
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
print(cnn)num_epochs = 2for epoch in range(num_epochs):for i, (images, labels) in enumerate(train_loader):images = get_variable(images)labels = get_variable(labels)  # print(labels)outputs = cnn(images)optimizer.zero_grad()  loss = loss_func(outputs, labels)loss.backward()  # 反向传播，自动计算每个节点的锑度至optimizer.step()if (i + 1) % 100 == 0:print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f'% (epoch + 1, num_epochs, i + 1, len(train_dataset) // batch_size, loss.item()))torch.save(cnn.state_dict(), 'cnn.pkl')

3.3 test.py模块不变，可以参考以上2.3

4、最后附上测试结果