前言
到这一章节为止,依据小土堆课程的PyTorch深度学习笔记基础部分已经完结了,接下来将依据李沐动手学深度学习课程进行PyTorch深度学习笔记的进阶部分
预测图片
完整的模型验证(测试,demo)套路,利用已经训练好的模型,然后给它提供输入。
输入狗的图片,并打开
python
image_path = "imgs/dog.png"
image = Image.open(image_path)
4通道的RGBA转为3通道的RGB图片
python
image = image.convert("RGB")
转换图像格式并设定网络
python
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,1,2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self, x):
x = self.model1(x)
return x
GPU上训练的东西映射到CPU上
python
model = torch.load("model/tudui_29.pth",map_location=torch.device('cpu'))
转为四维,符合网络输入需求
python
image = torch.reshape(image,(1,3,32,32))
将模型转为测试类型
python
model.eval()
不进行梯度计算,减少内存计算
python
with torch.no_grad():
output = model(image)
概率最大类别的输出
python
print(output.argmax(1))
完整代码
python
import torchvision
from PIL import Image
from torch import nn
import torch
image_path = "imgs/dog.png"
image = Image.open(image_path)
image = image.convert("RGB")
print(image)
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,1,2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self, x):
x = self.model1(x)
return x
model = torch.load("model/tudui_29.pth",map_location=torch.device('cpu'))
print(model)
image = torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():
output = model(image)
output = model(image)
print(output)
print(output.argmax(1))