适合入门深度学习的项目！！！搭建LeNet-5网络模型——3 测试模型搭建

在上一节中，我们保存了训练效果最好的模型"save_model/best.pth"，接下来对它进行测试，观察它的测试效果。详细代码如下：

import torch
from net import MyLeNet
from torch.autograd import Variable
from torchvision import datasets,transforms
from torchvision.transforms import ToPILImage

data_transform = transforms.ToTensor()

# 加载预训练的数据集
train_dataset = datasets.MNIST(root='./data',train=True,transform=data_transform,download=True)
train_dataloader = torch.utils.data.DataLoader(dataset = train_dataset,batch_size = 16, shuffle = True)

# 加载测集的数据
test_dataset = datasets.MNIST(root='./data',train=False,transform=data_transform,download=True)
test_dataloader = torch.utils.data.DataLoader(dataset = test_dataset,batch_size = 16, shuffle = True)

# 选择设备
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

model = MyLeNet().to(device)

#加载模型
model.load_state_dict(torch.load("save_model/best.pth"))

# 获取结果
classes = [
    "0",
    "1",
    "2",
    "3",
    "5",
    "5",
    "6",
    "7",
    "8",
    "9",
]
# 把tensor转为图片方便可视化
ShowImg = ToPILImage()

# 进入验证
for i in range(20):
    X,y = test_dataset[i][0],test_dataset[i][1]
    ShowImg(X).show()
    # 将张量扩展为4维
    # 张量可以看做多维数组，有时需要对维度进行相关操作
    # 比如，使用 PyTorch 对二维图像进行处理，需要扩展为四维的张量，即需要增加一个维度表示通道。
    # torch.squeeze() 和 torch.unsqueeze() 便能实现这样的操作
    # squeeze 即“挤压”的意思，用于维度压缩，unsqueeze 即解压缩的过程，用于增加维度
    X = Variable(torch.unsqueeze(X,dim=0).float()).to(device)
    with torch.no_grad():
        pred = model(X)
        predicted, actual = classes[torch.argmax(pred[0])],classes[y]
        print(f"predicted: \n Train Accuracy: {predicted}, actual: {actual} \n")

对应的图片都展示出来了：

发现模型效果很好：