|
- #!/usr/bin/env python
- # -*- coding: utf-8 -*-
- # @Time : 2023/2/5 下午7:32
- # @File : model
- # ----------------------------------------------
- # ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆
- # >>> Author : kevin
- # >>> QQ : 565479588
- # >>> Mail : lovecode@gmail.com
- # >>> Github : https://github.com/lovecode100
- # >>> Blog : https://www.cnblogs.com/lovecode
- # ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆ ☆
- import torch.nn as nn
- import torch
-
- # official pretrain weights
- model_urls = {
- 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
- 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
- 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
- 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'
- }
-
- #定义VGG的类,继承于nn.moudel这个父类
- class VGG(nn.Module):
- # 初始化函数(self,通过make_features(cfg: list)函数生成的提取特征网络结构,所需要分类的类别个数,是否对网络进行权重初始化)
- def __init__(self, features, num_classes=1000, init_weights=False):
- super(VGG, self).__init__()
- # 卷积层提取特征
- self.features = features
-
- # 全连接层进行分类
- self.classifier = nn.Sequential( #生成分类网络结构
- nn.Linear(512*7*7, 4096),
- nn.ReLU(True),
- nn.Dropout(p=0.5), #50%的比例进行随机失活
- nn.Linear(4096, 4096),
- nn.ReLU(True),
- nn.Dropout(p=0.5),
- nn.Linear(4096, num_classes) #num_classes输出节点个数就是分类类别的个数
- )
- if init_weights: #是否需要对网络进行参数的初始化,若init_weights为True则进行初始化权重函数
- self._initialize_weights()
-
- #正向传播的过程
- def forward(self, x):
- # N x 3 x 224 x 224
- x = self.features(x) # 将数据传入features结构得到输出x
- # N x 512 x 7 x 7
- x = torch.flatten(x, start_dim=1) # 将输出进行展平处理,start_dim=1表示从哪个维度进行展平处理
- # N x 512*7*7
- x = self.classifier(x) # 将特征矩阵输出的提前定义好的分类网络结构classifier函数中,得到输出
- return x
-
- #权重初始化函数
- def _initialize_weights(self):
- for m in self.modules(): # 进行遍历网络的每一个字模块(每一层)
- if isinstance(m, nn.Conv2d): # 如果遍历的当前层是卷积层
- # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
- nn.init.xavier_uniform_(m.weight) # 使用函数初始化卷积核的权重
- if m.bias is not None: # 如果卷积核使用了偏置的话
- nn.init.constant_(m.bias, 0) # 则把偏置默认初始化为0
- elif isinstance(m, nn.Linear): # 如果遍历的当前层是全连接层
- nn.init.xavier_uniform_(m.weight) # 使用函数初始化全连接层的权重
- # nn.init.normal_(m.weight, 0, 0.01)
- nn.init.constant_(m.bias, 0) # 则把偏置默认初始化为0
-
- #提取生成特征网络
- def make_features(cfg: list): #传入的是配置变量,是列表类型
- layers = [] # 定义空列表,用来存放每一层的结构
- in_channels = 3 # 由于输入的是彩色图片,定义输入通道为3
- for v in cfg: # 通过for循环来遍历配置列表
- # 最大池化层
- if v == "M": # 如果当前的配置元素是一个M字符的话
- layers += [nn.MaxPool2d(kernel_size=2, stride=2)] # 表示该层是最大池化层,则创建一个最大池化下采样层
- # 卷积层
- else:
- conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) # 创建一个卷积操作,stride默认是1,不需要写进去
- layers += [conv2d, nn.ReLU(True)]
- in_channels = v
- return nn.Sequential(*layers) # 利用nn.Sequential()函数通过非关键字的形式传进去。这单星号(*)将参数以元组(tuple)的形式导入
-
- # vgg网络模型配置列表,数字表示卷积核个数,'M'表示最大池化层
- cfgs = {
- 'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
- 'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
- 'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
- 'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
- }
-
- #定义vgg()函数,实例化给定的配置模型
- def vgg(model_name="vgg16", **kwargs): #model_name表示需要实例化哪个配置。这双星号(**)将参数以字典的形式导入
- assert model_name in cfgs, "Warning: model number {} not in cfgs dict!".format(model_name)
- cfg = cfgs[model_name] # 将配置传入字典,得到需要的配置列表
-
- model = VGG(make_features(cfg), **kwargs) # 通过VGG类进行实例化网络,**kwargs可变长度的字典变量
- return model
|