当前位置：首页 > news >正文

一个简单的图像分类项目（五）编写脚本：创建网络

news 2024/10/28 6:06:11

创建几个网络，用于训练和预测。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()

查看全文

http://www.mrgr.cn/news/60067.html