创建几个网络 ，用于训练和预测。lib.nets.py:
import torch.nn as nn
from torchvision import models
from script.setting import *
import torch.nn.functional as F
# 定义ResNet网络
class ResNet18(nn.Module):
def __init__(self, num_classes=classes, pretrained=is_pretrained):
super(ResNet18, self).__init__()
self.model = models.resnet18(pretrained=pretrained) # 调用torchvision.models中的resnet18
self.num_ftrs = self.model.fc.in_features # 获取全连接层的输入特征数
self.model.fc = nn.Linear(self.num_ftrs, num_classes) # 修改全连接层
def forward(self, x): # 前向传播
out = self.model(x)
return out
def resnet18_model(): # 定义网络结构的应用函数
return ResNet18()
# 定义VGG网络
class VGG_simple(nn.Module):
def __init__(self, num_classes=classes):
super(VGG_simple, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), # 输入通道为3，输出通道为64，卷积核大小为3x3，步长为1，填充为1
nn.BatchNorm2d(64), # 批量归一化
nn.ReLU() # 激活函数
)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化，池化核大小为2x2，步长为2
# 经过第一层卷积后，图像的尺寸为 16x16，通道数为64
self.conv2_1 = nn.Sequential( # 定义第2个卷积层的第一次卷积
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), # 输入通道为64，输出通道为128，卷积核大小为3x3，步长为1，填充为1
nn.BatchNorm2d(128), # 批量归一化
nn.ReLU() # 激活函数
)
self.conv2_2 = nn.Sequential( # 定义第2个卷积层的第二次卷积
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1), # 输入通道为64，输出通道为128，卷积核大小为3x3，步长为1，填充为1
nn.BatchNorm2d(128), # 批量归一化
nn.ReLU() # 激活函数
)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化，池化核大小为2x2，步长为2
# 经过第二层卷积后，图像的尺寸为 8x8，通道数为128
self.conv3_1 = nn.Sequential( # 定义第3个卷积层的第一次卷积
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), # 输入通道为64，输出通道为128，卷积核大小为3x3，步长为1，填充为1
nn.BatchNorm2d(256), # 批量归一化
nn.ReLU() # 激活函数
)
self.conv3_2 = nn.Sequential( # 定义第3个卷积层的第二次卷积
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(256), # 批量归一化
nn.ReLU() # 激活函数
)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
# 经过第三层卷积后，图像的尺寸为 4x4，通道数为256
self.conv4_1 = nn.Sequential( # 定义第4个卷积层的第一次卷积
nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(512), # 批量归一化
nn.ReLU() # 激活函数
)
self.conv4_2 = nn.Sequential( # 定义第4个卷积层的第二次卷积
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(512), # 批量归一化
nn.ReLU() # 激活函数
)
self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化，池化核大小为2x2，步长为2
# 经过第四层卷积后，图像的尺寸为 2x2，通道数为512
self.fc = nn.Linear(512 * 4, num_classes) # 全连接层，输入为512*4，输出为10
def forward(self, x): # 定义前向传播过程
batch_size = x.size(0) # 获取输入的batch_size
x = self.conv1(x) # 第1个卷积层
x = self.pool1(x) # 第1个池化层
x = self.conv2_1(x) # 第2个卷积层
x = self.conv2_2(x) # 第2个卷积层
x = self.pool2(x) # 第2个池化层
x = self.conv3_1(x) # 第3个卷积层
x = self.conv3_2(x) # 第3个卷积层
x = self.pool3(x) # 第3个池化层
x = self.conv4_1(x) # 第4个卷积层
x = self.conv4_2(x) # 第4个卷积层
x = self.pool4(x) # 第4个池化层
x = x.view(batch_size, -1) # 将图片展开成一行，-1表示自动计算这一维的大小
x = self.fc(x) # 全连接层
output = F.log_softmax(x, dim=1) # 对输出进行log_softmax处理
return output
def vgg_model(): # 定义网络结构的应用函数
return VGG_simple()