model.py

from torch import nn
import torchvision

# 搭建自定义神经网络
class classify_model(nn.Module):
    def __init__(self):
        super(classify_model, 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
    
# 利用现有的网络，并修改
class vgg_16(nn.Module):
    def __init__(self):
        super(vgg_16, self).__init__()
        self.model = torchvision.models.vgg16(weights=None)
        self.model.classifier.add_module('add_Linear', nn.Linear(1000, 10))

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