网络模型应用实例

import torch
import torchvision
from PIL import Image
from torch import nn

image_path = "D:/test pytorch/images/dog.png"
# 读取图片
image = Image.open(image_path)
image = image.convert('RGB')
# png格式的图片是4通道类型，RGB只有三通道，所以调用image = image.convert('RGB')保留其颜色通道。若本来就是三通道，加上这一步也可以。均适用
print(image)

# 网络要求32*32，但是图片显示201*176，对其resize
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                            torchvision.transforms.ToTensor()])
image = transform(image)
print(image.shape)

 # 加载网络模型
class SUN(nn.Module):
    def __init__(self):
        super(SUN, self).__init__()
        self.model = 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(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x =self.model(x)
        return x
model = torch.load("sun_0.path")

# 实验结果可看，预测结果不准确，实验结果预测是第二个，但实际在模型中，是第五个类别才是狗，原因是，训练次数的不足
# model = torch.load("sun_29.path")
# 报错显示，在GPU上运行的，使用在CPU上，需要加语句;
# model = torch.load("sun_29.path",map_location = torch.device('cpu'))
print(model)

# 报错：Given groups=1, weight of size [32, 3, 5, 5], expected input[1, 4, 32, 32] to have 3 channels, but got 4 channels instead
image = torch.reshape(image,(1,3,32,32))
# 当模型出现dropout或者batchnormal下述2条就需要下述两条语句
model.eval()
with torch.no_grad():
    output = model(image)
output = model(image)
print(output)
print(output.argmax(1))