#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))