Pytorch学习笔记（三）：模型定义

本篇笔记主要记录关于网络模型的一系列内容，包括模型的定义，模型参数初始化方法，模型的保存和加载，模型的 finetune (本质上还是模型权值初始化)，以及使用陈云所述的代码进行猫狗大战数据集的实例分析。

部分学习资源链接：

• Pytorch官方在优达学城公布的教学视频：
  https://classroom.udacity.com/courses/ud188
• 与其配套的代码和练习答案：
  https://github.com/udacity/deep-learning-v2-pytorch
• 陈云《深度学习框架Pytorch入门与实践》：
  https://github.com/chenyuntc/pytorch-book
• 《Pytorch模型训练实用教程》：
  https://github.com/tensor-yu/PyTorch_Tutorial

之前的相关文章：

• 《PYTORCH学习笔记（一）：代码结构布局》
• 《PYTORCH学习笔记（二）：载入数据》

模型定义的三要素

首先，必须继承 nn.Module 这个类，要让PyTorch知道这个类是一个 Module。

其次，在 __init__(self) 中设置好需要的“组件"(如conv、pooling、Linear、BatchNorm等)。

最后，在 forward(self, x) 中用定义好的“组件”进行组装，就像搭积木，把网络结构搭建出来，这样一个模型就定义好了。

这里选取笔记（一）中的代码进行示例。

首先可以看到先是定义了一个类 Net(nn.Module) :

def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

第一行是初始化，往后定义了一系列组件，如由 Conv2d 构成的 conv1 ，有 MaxPool2d 构成的 pool1 ，这些操作均由 torch.nn 提供，torch.nn 中的操作可查看官方文档： https://pytorch.org/docs/stable/nn.html

当这些组件定义好之后，就可以定义forward()函数，用来搭建网络结构：

def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

x 为模型的输入，第一行表示，x 经过 conv1，然后经过激活函数 relu，再经过 pool1 操作； 第二行于第一行一样；

第三行，表示将 x 进行 reshape，为了后面做为全连接层的输入；

第四，第五行的操作都一样，先经过全连接层 fc ，然后经过 relu ；

第六行，模型的最终输出是 fc3 输出。

至此，一个模型定义完毕，接着就可以在后面进行使用。

例如，实例化一个模型 net = Net() ，然后把输入 inputs 扔进去，outputs = net(inputs)，就可以得到输出 outputs 。

复杂模型代码举例 —— ResNet34

from .basic_module import BasicModule
from torch import nn
from torch.nn import functional as F


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)

还是从三要素出发看看是怎么定义Resnet34的。

首先，继承 nn.Module ；

其次，看 __init__() 函数，在 __init()__ 中，定义了这些组件，self.pre，self.layer1-4 , self.fc ；

最后，看 forward() ，分别用了在 __init()__ 中定义的一系列组件，并且用了 torch.nn.functional.avg_pool2d 这个操作

接着来详细分析一下具体代码内容。

nn.Sequetial

官方文档解释： https://pytorch.org/docs/stable/nn.html?highlight=sequential#torch.nn.Sequential

torch.nn.Sequential 其实就是 Sequential 容器，该容器将一系列操作按先后顺序给包起来，方便重复使用，

例如 Resnet 中有很多重复的 block ，就可以用 Sequential 容器把重复的地方包起来。

# Example of using Sequential
model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

# Example of using Sequential with OrderedDict
model = nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(1,20,5)),
          ('relu1', nn.ReLU()),
          ('conv2', nn.Conv2d(20,64,5)),
          ('relu2', nn.ReLU())
        ]))

可以看出模型的定义就是先继承，再构建组件，最后组装。

其中基本组件可从 torch.nn 中获取，或者从 torch.nn.functional 中获取，同时为了方便重复使用组件，可以使用 Sequential 容器将一系列组件包起来，最后在 forward() 函数中将这些组件组装成你的模型。

权值初始化

模型定义完成后，通常还需要对权值进行初始化，才能开始训练，并且初始化方法会直接影响到模型的收敛与否。

权值初始化流程

总共两步，

第一步，先设定什么层用什么初始化方法，初始化方法在 torch.nn.init 中给出；

第二步，实例化一个模型之后，执行该函数，即可完成初始化。

依旧拿笔记（一）中的代码举例：

def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.xavier_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                torch.nn.init.normal_(m.weight.data, 0, 0.01)
                m.bias.data.zero_()

这段代码基本流程是这样，先从 self.modules() 中遍历每一层，然后判断各层属于什么类型，例如，是否是 nn.Conv2d 、nn.BatchNorm2d 、nn.Linear 等，然后根据不同类型的层，设定不同的权值初始化方法，例如，Xavier ， kaiming ，normal_ ，uniform_ 等。

Ps: kaiming 也称之为MSRA初始化，当年何恺明还在微软亚洲研究院，因而得名。

来看看第一行代码中的 self.modules() ，

 def modules(self):
        for name, module in self.named_modules():
            yield module

功能是：Returns an iterator over all modules in the network. 能依次返回模型中的各层。例如：

接着，判断 m 的类型，属于什么类型，可以看到当前 m 属于 Conv2d 类型，则进行如下初始化：

torch.nn.init.xavier_normal(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()

以上代码表示采用 torch.nn.init.xavier_normal 方法对该层的 weight 进行初始化，并判断是否存在偏置 (bias) ，若存在，将 bias 初始化为全0。

这样，该层就初始化完毕，参照以上流程，不断遍历模型的每一层，最终完成模型的初始化。

权值初始化的方法

PyTorch在 torch.nn.init 中提供了常用的初始化方法函数，阅读官方文档可以详细查明。

torch.nn.init: https://pytorch.org/docs/stable/nn.init.html

Xavier 初始化方法，论文在《Understanding the difficulty of training deep feedforward neural networks》，公式推导是从“方差一致性”出发，初始化的分布有均匀分布和正态分布两种。

从代码中发现，即使不进行初始化，我们模型的权值也不为空，而是有值的，其实，在创建网络实例的过程中, 一旦调用 nn.Conv2d 的时候就会有对权值进行初始化，初始化过程是在 Conv2d 的基类 ConvNd 中进行的，来看看关系：

class Conv2d(_ConvNd):
--> 在_ConvNd中：
--> self.reset_parameters()
---> def reset_parameters(self)
---> self.weight.data.uniform(-stdv, stdv)

模型Finetune

我们知道一个良好的权值初始化，可以使收敛速度加快，甚至可以获得更好的精度。而在实际应用中，我们通常采用一个已经训练模型的模型的权值参数作为我们模型的初始化参数，也称之为 Finetune ，更宽泛的称之为迁移学习。迁移学习中的 Finetune 技术，本质上就是让我们新构建的模型，拥有一个较好的权值初始值。

finetune 权值初始化三步曲，finetune 就相当于给模型进行初始化，其流程共用三步：

• 第一步：保存模型，拥有一个预训练模型；
• 第二步：加载模型，把预训练模型中的权值取出来；
• 第三步：初始化，将权值对应的“放”到新模型中

Finetune之权值初始化

在进行 finetune 之前我们需要拥有一个模型或者是模型参数，因此需要了解如何保存模型。官方文档中介绍了两种保存模型的方法，一种是保存整个模型，另外一种是仅保存模型参数（官方推荐用这种方法），这里采用官方推荐的方法。

第一步：保存模型参数

若拥有模型参数，可跳过这一步。

假设创建了一个 net = Net()，并且经过训练，通过以下方式保存：

torch.save(net.state_dict(), 'net_params.pkl')

第二步：加载模型

进行三步曲中的第二步，加载模型，这里只是加载模型的参数：

pretrained_dict = torch.load('net_params.pkl')

第三步：初始化

进行三步曲中的第三步，将取到的权值，对应的放到新模型中：

首先我们创建新模型，并且获取新模型的参数字典 net_state_dict ：

net = Net() # 创建net<br>
net_state_dict = net.state_dict() # 获取已创建net的state_dict

接着将 pretrained_dict 里不属于 net_state_dict 的键剔除掉：

pretrained_dict_1 = {k: v for k, v in pretrained_dict.items() if k in net_state_dict}

然后，用预训练模型的参数字典 对 新模型的参数字典 net_state_dict 进行更新：

net_state_dict.update(pretrained_dict_1)

最后，将更新了参数的字典 “放”回到网络中：

net.load_state_dict(net_state_dict)

这样，利用预训练模型参数对新模型的权值进行初始化过程就做完了。

采用 finetune 的训练过程中，有时候希望前面层的学习率低一些，改变不要太大，而后面的全连接层的学习率相对大一些。这时就需要对不同的层设置不同的学习率，下面就介绍如何为不同层配置不同的学习率。

不同层设置不同的学习率

在利用 pre-trained model 的参数做初始化之后，我们可能想让fc层更新相对快一些，而希望前面的权值更新小一些，这就可以通过为不同的层设置不同的学习率来达到此目的。

为不同层设置不同的学习率，主要通过优化器对多个参数组进行设置不同的参数。所以，只需要将原始的参数组，划分成两个，甚至更多的参数组，然后分别进行设置学习率。

这里将原始参数“切分”成fc3层参数和其余参数，为fc3层设置更大的学习率。

ignored_params = list(map(id, net.fc3.parameters())) # 返回的是parameters的 内存地址
base_params = filter(lambda p: id(p) not in ignored_params, net.parameters())
optimizer = optim.SGD([
{'params': base_params},
{'params': net.fc3.parameters(), 'lr': 0.001*10}], 0.001, momentum=0.9, weight_decay=1e-4)

第一行+ 第二行的意思就是，将 fc3 层的参数 net.fc3.parameters() 从原始参数 net.parameters() 中剥离出来，base_params 就是剥离了 fc3 层的参数的其余参数，然后在优化器中为 fc3 层的参数单独设定学习率。

optimizer = optim.SGD(……) 这里的意思就是 base_params 中的层，用 0.001 , momentum=0.9 , weight_decay=1e-4

fc3 层设定学习率为： 0.001*10

补充：

挑选出特定的层的机制是利用内存地址作为过滤条件，将需要单独设定的那部分参数，从总的参数中剔除。base_params 是一个 list ，每个元素是一个 Parameter 类， net.fc3.parameters() 是一个 <generator object parameters>

ignored_params = list(map(id, net.fc3.parameters()))

net.fc3.parameters() 是一个 <generator object parameters at 0x11b63bf00>，所以迭代的返回其中的 parameter ，这里有 weight 和 bias ，最终返回 weight 和 bias 所在内存的地址。

完整代码如下：

# coding: utf-8

import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import numpy as np
import os
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import sys
sys.path.append("..")
from utils.utils import MyDataset, validate, show_confMat
from datetime import datetime

train_txt_path = os.path.join("..", "..", "Data", "train.txt")
valid_txt_path = os.path.join("..", "..", "Data", "valid.txt")

classes_name = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

train_bs = 16
valid_bs = 16
lr_init = 0.001
max_epoch = 1

# log
result_dir = os.path.join("..", "..", "Result")

now_time = datetime.now()
time_str = datetime.strftime(now_time, '%m-%d_%H-%M-%S')

log_dir = os.path.join(result_dir, time_str)
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

# -------------------------------------------- step 1/5 : 加载数据 -------------------------------------------

# 数据预处理设置
normMean = [0.4948052, 0.48568845, 0.44682974]
normStd = [0.24580306, 0.24236229, 0.2603115]
normTransform = transforms.Normalize(normMean, normStd)
trainTransform = transforms.Compose([
    transforms.Resize(32),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    normTransform
])

validTransform = transforms.Compose([
    transforms.ToTensor(),
    normTransform
])

# 构建MyDataset实例
train_data = MyDataset(txt_path=train_txt_path, transform=trainTransform)
valid_data = MyDataset(txt_path=valid_txt_path, transform=validTransform)

# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=train_bs, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=valid_bs)

# ------------------------------------ step 2/5 : 定义网络 ------------------------------------


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    # 定义权值初始化
    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.xavier_normal_(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                torch.nn.init.normal_(m.weight.data, 0, 0.01)
                m.bias.data.zero_()


net = Net()     # 创建一个网络

# ================================ #
#        finetune 权值初始化
# ================================ #

# load params
pretrained_dict = torch.load('net_params.pkl')

# 获取当前网络的dict
net_state_dict = net.state_dict()

# 剔除不匹配的权值参数
pretrained_dict_1 = {k: v for k, v in pretrained_dict.items() if k in net_state_dict}

# 更新新模型参数字典
net_state_dict.update(pretrained_dict_1)

# 将包含预训练模型参数的字典"放"到新模型中
net.load_state_dict(net_state_dict)

# ------------------------------------ step 3/5 : 定义损失函数和优化器 ------------------------------------
# ================================= #
#         按需设置学习率
# ================================= #

# 将fc3层的参数从原始网络参数中剔除
ignored_params = list(map(id, net.fc3.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, net.parameters())

# 为fc3层设置需要的学习率
optimizer = optim.SGD([
    {'params': base_params},
    {'params': net.fc3.parameters(), 'lr': lr_init*10}],  lr_init, momentum=0.9, weight_decay=1e-4)

criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)     # 设置学习率下降策略

# ------------------------------------ step 4/5 : 训练 --------------------------------------------------

for epoch in range(max_epoch):

    loss_sigma = 0.0    # 记录一个epoch的loss之和
    correct = 0.0
    total = 0.0
    scheduler.step()  # 更新学习率

    for i, data in enumerate(train_loader):
        # 获取图片和标签
        inputs, labels = data
        inputs, labels = Variable(inputs), Variable(labels)

        # forward, backward, update weights
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 统计预测信息
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).squeeze().sum().numpy()
        loss_sigma += loss.item()

        # 每10个iteration 打印一次训练信息，loss为10个iteration的平均
        if i % 10 == 9:
            loss_avg = loss_sigma / 10
            loss_sigma = 0.0
            print("Training: Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch + 1, max_epoch, i + 1, len(train_loader), loss_avg, correct / total))
            print('参数组1的学习率:{}, 参数组2的学习率:{}'.format(scheduler.get_lr()[0], scheduler.get_lr()[1]))
    # ------------------------------------ 观察模型在验证集上的表现 ------------------------------------
    loss_sigma = 0.0
    cls_num = len(classes_name)
    conf_mat = np.zeros([cls_num, cls_num])  # 混淆矩阵
    net.eval()
    for i, data in enumerate(valid_loader):

        # 获取图片和标签
        images, labels = data
        images, labels = Variable(images), Variable(labels)

        # forward
        outputs = net(images)
        outputs.detach_()

        # 计算loss
        loss = criterion(outputs, labels)
        loss_sigma += loss.item()

        # 统计
        _, predicted = torch.max(outputs.data, 1)
        # labels = labels.data    # Variable --> tensor

        # 统计混淆矩阵
        for j in range(len(labels)):
            cate_i = labels[j].numpy()
            pre_i = predicted[j].numpy()
            conf_mat[cate_i, pre_i] += 1.0

    print('{} set Accuracy:{:.2%}'.format('Valid', conf_mat.trace() / conf_mat.sum()))
print('Finished Training')

# ------------------------------------ step5: 绘制混淆矩阵图 ------------------------------------

conf_mat_train, train_acc = validate(net, train_loader, 'train', classes_name)
conf_mat_valid, valid_acc = validate(net, valid_loader, 'valid', classes_name)

show_confMat(conf_mat_train, classes_name, 'train', log_dir)
show_confMat(conf_mat_valid, classes_name, 'valid', log_dir)

实战解析—— Kaggle 上的经典比赛：Dogs vs. Cats 的模型定义

这里使用的是陈云给出的自己的例子，例子中他的模型定义时和平常所见到的并不一样，使用的是 BasicModule ，这个是对 nn.Module 的简易封装，提供快速加载和保存模型的接口。

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):
        """
        保存模型，默认使用“模型名字+时间”作为文件名，
        如AlexNet_0710_23:57:29.pth
        """
        if name is None:
            #存储到文件夹checkpoints下面
            prefix = 'checkpoints/' + self.model_name + '_'
            name = time.strftime(prefix + '%m%d_%H:%M:%S.pth')
        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) #得到的是批数据的大小

在实际使用中，直接调用 model.save() 及 model.load(opt.load_path) 即可。

其它自定义模型一般继承 BasicModule ，然后实现自己的模型。其中 alexNet.py 实现了 AlexNet ，resNet34 实现了 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

alexnet.py 为：

# coding:utf8
from torch import nn
from .basic_module import BasicModule


class AlexNet(BasicModule):
    """
    code from torchvision/models/alexnet.py
    结构参考 <https://arxiv.org/abs/1404.5997>
    """

    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), # inplace-选择是否进行覆盖运算
            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

resnet34.py 为：

# coding:utf8
from .basic_module import BasicModule
from torch import nn
from torch.nn import functional as F


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)

squeezenet.py 为：

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)

这样在主函数就可以写成：

from models import AlexNet
#或
import models
model = models.AlexNet()
#或
import models
model = getattr('models', 'AlexNet')()

其中最后一种写法最为关键，这意味着我们可以通过字符串直接指定使用的模型，而不必使用判断语句，也不必在每次新增加模型后都修改代码。新增模型后只需要在 models/__init__.py 中加上 from .new_module import new_module 即可。

此部分陈云自己也说明代码存在不少争议，主要是因为不同的人数学代码的习惯也不同， 有不少人喜欢将模型的训练过程集成于模型的定义之中，代码结构如下所示：

class MyModel(nn.Module):
  	
      def __init__(self,opt):
          self.dataloader = Dataloader(opt)
          self.optimizer  = optim.Adam(self.parameters(),lr=0.001)
          self.lr = opt.lr
          self.model = make_model()
      
      def forward(self,input):
          pass
      
      def train_(self):
          # 训练模型
          for epoch in range(opt.max_epoch)
          	for ii,data in enumerate(self.dataloader):
              	train_epoch()
              
          	model.save()
  	
      def train_epoch(self):
          pass

抑或是专门设计一个 Trainer 对象，形如：

  '''
  code simplified from:
  https://github.com/pytorch/pytorch/blob/master/torch/utils/trainer/trainer.py
  '''
  import heapq
  from torch.autograd import Variable

  class Trainer(object):

      def __init__(self, model=None, criterion=None, optimizer=None, dataset=None):
          self.model = model
          self.criterion = criterion
          self.optimizer = optimizer
          self.dataset = dataset
          self.iterations = 0

      def run(self, epochs=1):
          for i in range(1, epochs + 1):
              self.train()

      def train(self):
          for i, data in enumerate(self.dataset, self.iterations + 1):
              batch_input, batch_target = data
              self.call_plugins('batch', i, batch_input, batch_target)
              input_var = Variable(batch_input)
              target_var = Variable(batch_target)
    
              plugin_data = [None, None]
    
              def closure():
                  batch_output = self.model(input_var)
                  loss = self.criterion(batch_output, target_var)
                  loss.backward()
                  if plugin_data[0] is None:
                      plugin_data[0] = batch_output.data
                      plugin_data[1] = loss.data
                  return loss
    
              self.optimizer.zero_grad()
              self.optimizer.step(closure)
    
          self.iterations += i

还有一些人喜欢模仿 keras 和 scikit-learn 的设计，设计一个 fit 接口。对读者来说，这些处理方式很难说哪个更好或更差，找到最适合自己的方法才是最好的。

BasicModule 的封装，可多可少。训练过程中的很多操作都可以移到BasicModule之中，比如 get_optimizer 方法用来获取优化器，比如 train_step 用来执行单歩训练。对于不同的模型，如果对应的优化器定义不一样，或者是训练方法不一样，可以复写这些函数自定义相应的方法，取决于自己的喜好和项目的实际需求。