PyTorch学习笔记（9）使用GPU训练模型并测试自己的模型

对之前的训练再优化，如果自己的电脑没有 GPU，也可以去网上找一些服务，如谷歌、阿里云等。

---

使用GPU训练

只需要将网络模型、损失函数、数据放入 cuda() 即可。

例如：

loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()

imgs,targets = data
imgs = imgs.cuda()
targets = targets.cuda()

但是没有 GPU 的时候会报异常，所以最好加判断：

loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
	loss_fn = loss_fn.cuda()

显然，这样写，代码修改起来很麻烦，因此还有更简便的方法。

# 在最开头加上此句
# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 之后网络模型、损失函数、数据如下格式修改即可
imgs,targets = data
imgs = imgs.to(device)
targets = targets.to(device)

测试自己的模型

因为上文是对 CIFAR10 进行的训练，因此测试图片应去找十个类别之一的物体，这里随便找了一张猫的图片。

值得注意的是，png 格式是四通道的，除了 RGB 三通道以外，还有一个透明度通道，因此需要 img = img.convert("RGB") 来保留颜色通道。（这样就能适应各种格式的图片了）

最后的 output 为该图片属于 10 个类别的各自概率，再用 output.argmax(1) 取最大索引，一般就认为这个最大的概率是正确的。测试中所对应的标签刚好是 cat，意味着本次实验运气不错，猜对了。

import torch
import torchvision
from PIL import Image
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear

# 随便在网上找一张动物的图片
img_path = "../imgs/cat.png"
img = Image.open(img_path)
print(img)

#只保留颜色三通道   .png格式多了一个透明度通道
img = img.convert("RGB")

transform = torchvision.transforms.Compose([ torchvision.transforms.Resize((32,32)),
                                           torchvision.transforms.ToTensor() ])
img = transform(img)
print(img.shape)


# 搭建神经网络
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.moudle = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x =  self.moudle(x)
        return x

model = torch.load("./moudels/moudel_20.pth",map_location=torch.device("cpu")) #方式一：将GPU上的模型映射到CPU上
# model = model.cuda()        #方式二：用GPU训练的模型必须用GPU来测试使用！！！！！！
# print(model)
img = torch.reshape(img,(1,3,32,32))
# img = img.cuda()
output = model(img)
print(output)

# 最大预测概率对应的标签索引
print(output.argmax(1))
# tensor([3])
# 猜测为 cat

# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']