Python 笔记 PyTorch nn_seq && nn_save && test

#nn_seq#

from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential
from torch import nn
import torch
from torch.utils.tensorboard import SummaryWriter


class Test(nn.Module):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 利用Sequential使代码更简洁, 更易管理
        self.module1 = 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.module1(x)
        return x


test = Test()
print(test)

input = torch.ones((64, 3, 32, 32))  # 测试网络是否正确搭建
output = test(input)
print(output)

writer = SummaryWriter("D:/dev/python/pyWork/Season2/Stage1/logs")
writer.add_graph(test, input)
writer.close()

#nn_save#

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrianed=False)

# 保存方式1: 模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")

# 保存方式2: 模型参数(官方推荐)
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

# 加载模型
model = torch.load("vgg16_method1.pth")

#test#

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

"""
    验证模型
"""


class Test(nn.Module):

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 利用Sequential使代码更简洁, 更易管理
        self.module1 = Sequential(
            Conv2d(3, 32, 5, 1, 2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, 1, 2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, 1, 2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

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


# 数据处理
image_path = "D:/dev/python/pyWork/Season2/Stage1/data/myimg/prime.jpg"
image = Image.open(image_path)

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

image = transform(image)
image = torch.reshape(image, (1, 3, 32, 32))
print(image.shape)

# 调用模型
model = torch.load("test_30.pth", map_location=torch.device("cpu"))  # 模型是用GPU训练的, 测试需要转换到CPU
model.eval()
with torch.no_grad():  # 节约性能
    output = model(image)

print(output.argmax(1))