一个简单的图像分类项目(五)编写脚本:创建网络
创建几个网络 ,用于训练和预测。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中的resnet18self.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 outdef resnet18_model(): # 定义网络结构的应用函数return ResNet18()# 定义VGG网络
class VggCustom(nn.Module):def __init__(self, num_classes=classes):super(VggCustom, self).__init__()self.conv1 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), # 输入通道为3,输出通道为64,卷积核大小为3x3,步长为1,填充为1nn.BatchNorm2d(64), # 批量归一化nn.ReLU() # 激活函数)self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化,池化核大小为2x2,步长为2# 经过第一层卷积后,图像的尺寸为 16x16,通道数为64self.conv2_1 = nn.Sequential( # 定义第2个卷积层的第一次卷积nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), # 输入通道为64,输出通道为128,卷积核大小为3x3,步长为1,填充为1nn.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,填充为1nn.BatchNorm2d(128), # 批量归一化nn.ReLU() # 激活函数)self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化,池化核大小为2x2,步长为2# 经过第二层卷积后,图像的尺寸为 8x8,通道数为128self.conv3_1 = nn.Sequential( # 定义第3个卷积层的第一次卷积nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), # 输入通道为64,输出通道为128,卷积核大小为3x3,步长为1,填充为1nn.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,通道数为256self.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,通道数为512self.fc = nn.Linear(512 * 4, num_classes) # 全连接层,输入为512*4,输出为10def forward(self, x): # 定义前向传播过程batch_size = x.size(0) # 获取输入的batch_sizex = 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 outputdef vgg_model(): # 定义网络结构的应用函数return VggCustom()