如果对这几种基础模型不太了解,请先参考博客。
首先来看程序文件的组织结构:
├── checkpoints/
├── data/
│ ├── __init__.py
│ ├── dataset.py
│ └── get_data.sh
├── models/
│ ├── __init__.py
│ ├── alexnet.py
│ ├── basic_module.py
│ └── resnet34.py
└── utils/
│ ├── __init__.py
│ └── visualize.py
├── config.py
├── main.py
├── requirements.txt
├── README.md
其中:
checkpoints/
: 用于保存训练好的模型,可使程序在异常退出后仍能重新载入模型,恢复训练data/
:数据相关操作,包括数据预处理、dataset实现等models/
:模型定义,可以有多个模型,例如上面的AlexNet和ResNet34,一个模型对应一个文件utils/
:可能用到的工具函数,在本次实验中主要是封装了可视化工具config.py
:配置文件,所有可配置的变量都集中在此,并提供默认值main.py
:主文件,训练和测试程序的入口,可通过不同的命令来指定不同的操作和参数requirements.txt
:程序依赖的第三方库README.md
:提供程序的必要说明模型的定义主要保存在models/
目录下,其中BasicModule
是对nn.Module
的简易封装,提供快速加载和保存模型的接口。
其它自定义模型一般继承BasicModule
,然后实现自己的模型。其中alexnet.py
实现了alexnet,resnet34.py
实现了resnet34。在models/__init__py
中,代码如下:
from .alexnet import AlexNet
from .resnet34 import ResNet34
from .squeezenet import SqueezeNet
# from torchvision.models import InceptinV3
# from torchvision.models import alexnet as AlexNet
这样在主函数中就可以写成:
from models import AlexNet
或
import models
model = models.AlexNet()
或
import models
model = getattr(models, 'AlexNet')()
其中最后一种写法最为关键,这意味着我们可以通过字符串直接指定使用的模型,而不必使用判断语句,也不必在每次新增加模型后都修改代码。新增模型后只需要在models/__init__.py
中加上from .new_module import new_module
即可。
# coding:utf-8
import torch as t
import time
class BasicModule(t.nn.Module):
"""
简易封装了nn.Module,主要是提供了save和load两个方法
"""
def __init__(self):
super(BasicModule, self).__init__()
self.model_name = str(type(self)) # 模型的默认名字
def load(self, path):
"""
可加载指定路径的模型
"""
self.load_state_dict(t.load(path))
def save(self, name=None):
"""
保存模型,默认使用“模型名字+时间”作为文件名
如:resnet34_05-30_22.29.46.pth
"""
if name is None:
prefix = 'checkpoints/' + self.model_name + '_'
name = time.strftime(prefix + '%m-%d_%H.%M.%S.pth') # 有改动,windows文件名不支持冒号:
t.save(self.state_dict(), name)
return name
def get_optimizer(self, lr, weight_decay):
return t.optim.Adam(self.parameters(), lr=lr, weight_decay=weight_decay)
class Flat(t.nn.Module):
"""
把输入reshape成(batch_size,dim_length)
"""
def __init__(self):
super(Flat, self).__init__()
# self.size = size
def forward(self, x):
return x.view(x.size(0), -1)
# coding:utf-8
from torch import nn
from .basic_module import BasicModule
class AlexNet(BasicModule):
"""
code from torchvision/models/alexnet.py
结构参考
"""
def __init__(self, num_classes=2):
super(AlexNet, self).__init__()
self.model_name = 'alexnet'
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.classifier(x)
return x
# coding:utf-8
from .basic_module import BasicModule
from torch import nn
from torch.nn import functional as F
import time
class ResidualBlock(nn.Module):
"""
实现子module: Residual Block
"""
def __init__(self, inchannel, outchannel, stride=1, shortcut=None):
super(ResidualBlock, self).__init__()
self.left = nn.Sequential(
nn.Conv2d(inchannel, outchannel, 3, stride, 1, bias=False),
nn.BatchNorm2d(outchannel),
nn.ReLU(inplace=True),
nn.Conv2d(outchannel, outchannel, 3, 1, 1, bias=False),
nn.BatchNorm2d(outchannel))
self.right = shortcut
def forward(self, x):
out = self.left(x)
residual = x if self.right is None else self.right(x)
out += residual
return F.relu(out)
class ResNet34(BasicModule):
"""
实现主module:ResNet34
ResNet34包含多个layer,每个layer又包含多个Residual block
用子module来实现Residual block,用_make_layer函数来实现layer
"""
def __init__(self, num_classes=2):
super(ResNet34, self).__init__()
self.model_name = 'resnet34'
# 前几层: 图像转换
self.pre = nn.Sequential(
nn.Conv2d(3, 64, 7, 2, 3, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(3, 2, 1))
# 重复的layer,分别有3,4,6,3个residual block
self.layer1 = self._make_layer(64, 128, 3)
self.layer2 = self._make_layer(128, 256, 4, stride=2)
self.layer3 = self._make_layer(256, 512, 6, stride=2)
self.layer4 = self._make_layer(512, 512, 3, stride=2)
# 分类用的全连接
self.fc = nn.Linear(512, num_classes)
def _make_layer(self, inchannel, outchannel, block_num, stride=1):
"""
构建layer,包含多个residual block
"""
shortcut = nn.Sequential(
nn.Conv2d(inchannel, outchannel, 1, stride, bias=False),
nn.BatchNorm2d(outchannel))
layers = []
layers.append(ResidualBlock(inchannel, outchannel, stride, shortcut))
for i in range(1, block_num):
layers.append(ResidualBlock(outchannel, outchannel))
return nn.Sequential(*layers)
def forward(self, x):
x = self.pre(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = F.avg_pool2d(x, 7)
x = x.view(x.size(0), -1)
return self.fc(x)
# coding:utf-8
from torchvision.models import squeezenet1_1
from models.basic_module import BasicModule
from torch import nn
from torch.optim import Adam
class SqueezeNet(BasicModule):
def __init__(self, num_classes=2):
super(SqueezeNet, self).__init__()
self.model_name = 'squeezenet'
self.model = squeezenet1_1(pretrained=True)
# 修改 原始的num_class: 预训练模型是1000分类
self.model.num_classes = num_classes
self.model.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Conv2d(512, num_classes, 1),
nn.ReLU(inplace=True),
nn.AvgPool2d(13, stride=1)
)
def forward(self,x):
return self.model(x)
def get_optimizer(self, lr, weight_decay):
# 因为使用了预训练模型,我们只需要训练后面的分类
# 前面的特征提取部分可以保持不变
return Adam(self.model.classifier.parameters(), lr, weight_decay=weight_decay)
done~