用之前的神经网络训练的模型进行验证,即
# 定义神经网络
class deep_nn(nn.Module):
def __init__(self):
super().__init__()
# 第一个卷积层输入通道为3,输出通道为32,卷积核5×5,计算得padding=2
self.conv1 = Conv2d(3, 32, 5, padding=2)
# 引入激活函数
self.relu = ReLU()
# 池化层不改变通道数,默认不填充
self.maxpool1 = MaxPool2d(2)
self.conv2 = Conv2d(32, 32, 5, padding=2)
self.maxpool2 = MaxPool2d(2)
self.conv3 = Conv2d(32, 64, 5, padding=2)
self.maxpool3 = MaxPool2d(2)
self.flatten = Flatten()
self.liner1 = Linear(1024, 64)
self.liner2 = Linear(64, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = self.maxpool1(x)
x = self.conv2(x)
x = self.relu(x)
x = self.maxpool2(x)
x = self.conv3(x)
x = self.relu(x)
x = self.maxpool3(x)
x = self.flatten(x)
x = self.liner1(x)
x = self.liner2(x)
return x
该网络是一个分类模型,导入一张狗的图片
验证代码
from torch import nn
import torch
import torchvision
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, ReLU
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from PIL import Image
from module import *
img_path = "./image/dog.png"
# 转化为PIL格式
image = Image.open(img_path)
image = image.convert('RGB')
# 转化为tensor数据类型,并resize成32×32
trans = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)), torchvision.transforms.ToTensor()])
image = trans(image)
# 加载网络模型
model = torch.load("epoch_10.pt")
print(image.shape)
image = torch.reshape(image, [1, 3, 32, 32])
image = image.cuda()
# 测试开始
model.eval()
with torch.no_grad():
output = model(image)
print(output.argmax(1))
输出
该网络模型的10个类编号
可以看出,正确的预测了dog类