pytorch ——模型创建与nn.Module

1、网络模型创建步骤

模型模块中分为两个部分，模型创建和权值初始化；

模型创建又分为两部分，构建网络层和拼接网络层；网络层有卷积层，池化层，激活函数等；构建网络层后，需要进行网络层的拼接，拼接成LeNet，AlexNet和ResNet等。

创建好模型后，需要对模型进行权值初始化，pytorch提供了丰富的初始化方法，Xavier，Kaiming，均匀分布，正态分布等。

以上一切都会基于nn.Module进行，nn.Module是整个模块的根基，下面会详细介绍nn.Module。

1.1 模型创建步骤

以LeNet模型进行讲解，LeNet由很多网络层构成，由两个卷积层，两个池化层和三个全连接层组成。

在创建LeNet的时候会首先构建子模块，构建子模块之后按照一定的顺序进行连接，然后包装起来就可以得到LeNet。

LeNet是由很多子模块构成的，而这些子模块是基于nn.Module构成的。下面形象表示一下创建LeNet的过程（计算图）是怎么的。

上图可以看做一个计算图，计算图有两个主要的概念，一个是节点一个是边，节点就是张量数据，边就是运算，在图中就是箭头。

构建模型有两要素，第一是构建子模块，比如LeNet是由很多网络层构成的，所以首先得构建子模块中的网络层。构建好网络层后，第二是拼接子模块，按照一定拓扑结构拼接子模块就可以得到模型。

下面以代码阐述模型构建的步骤，首先是初始化部分，构建子模块是在__init__()中进行的：

import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module):
    def __init__(self, classes):
        super(LeNet, self).__init__()    # 继承父类nn.Module的初始化
        self.conv1 = nn.Conv2d(3, 6, 5)    # 卷积层，卷积核为5*5，输入通道为3，输出通道为6
        self.conv2 = nn.Conv2d(6, 16, 5)    # 卷积层
        self.fc1 = nn.Linear(16*5*5, 120)    # 全连接层
        self.fc2 = nn.Linear(120, 84)	# 全连接层
        self.fc3 = nn.Linear(84, classes)	# 全连接层
	    
	def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out

    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                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):
                nn.init.normal_(m.weight.data, 0, 0.1)
                m.bias.data.zero_()

构建完子模块之后，进行模型的拼接是在另一部分代码中进行的：

net = LeNet(classes=2)
outputs = net(inputs)

当将输入传递给LeNet类对象net的时候，进入net(inputs)代码内部看具体情况，点击step into会进入module.py文件当中的__call__()函数，因为LeNet是继承module类的，module类中有__call__()函数表示一个实例是可以被调用的，我们进入__call__()当中的forward()函数进行查看（代码中第11行）：

    def __call__(self, *input, **kwargs):
        for hook in self._forward_pre_hooks.values():
            result = hook(self, input)
            if result is not None:
                if not isinstance(result, tuple):
                    result = (result,)
                input = result
        if torch._C._get_tracing_state():
            result = self._slow_forward(*input, **kwargs)
        else:
            result = self.forward(*input, **kwargs)
        for hook in self._forward_hooks.values():
            hook_result = hook(self, input, result)
            if hook_result is not None:
                result = hook_result
        if len(self._backward_hooks) > 0:
            var = result
            while not isinstance(var, torch.Tensor):
                if isinstance(var, dict):
                    var = next((v for v in var.values() if isinstance(v, torch.Tensor)))
                else:
                    var = var[0]
            grad_fn = var.grad_fn
            if grad_fn is not None:
                for hook in self._backward_hooks.values():
                    wrapper = functools.partial(hook, self)
                    functools.update_wrapper(wrapper, hook)
                    grad_fn.register_hook(wrapper)
        return result

对代码中的forward()函数点击step into，然后就进入了上面第一段代码中LeNet类的forward(self, x)函数当中：

def forward(self, x):
    out = F.relu(self.conv1(x))  # import torch.nn.functional as F
    out = F.max_pool2d(out, 2)
    out = F.relu(self.conv2(out))
    out = F.max_pool2d(out, 2)
    out = out.view(out.size(0), -1)
    out = F.relu(self.fc1(out))
    out = F.relu(self.fc2(out))
    out = self.fc3(out)
    return out

这段代码具体实现了前向传播，实现了每一层的计算，最终得到分类结果out，然后out会返回forward函数中的result = self.forward(*input, **kwargs)，这样就得到了outputs = net(inputs)。

综上，构建模型的子模块是在__init__()函数中实现的，拼接模块是在forward()函数中实现的，这就是模型构建的两个要素，构建子模块(init())和拼接子模块(forward())。

2、nn.Module

在构建模型模块的过程中，有一个非常重要的概念是nn.Module，所有的网络层都是继承于这个类的，现在了解一下nn.Module这个类。

介绍一下与nn.Module相关的几个模块，第一个是torch.nn，这是pytorch的一个神经网络模块，在torch.nn中有很多子模块，这里介绍四个：

• nn.Parameter：张量子类，表示可学习参数，如weight，bias；
• nn.Module：所有网络层基类，管理网络属性；
• nn.functional：函数具体实现，如卷积，池化，激活函数等；
• nn.init：参数初始化方法

这里主要介绍nn.Module，nn.module中有八个重要的属性用于管理整个模型。这里主要关注其中的parameters和modules两个属性。

• parameters：管理存储属于nn.Parameter类的属性，例如权值或者偏置参数；
• modules：用来存储管理nn.Module类，例如在LeNet中会构建子模块，modules就会存储创建的卷积层等；
• buffers：存储管理缓冲的属性，如训练过程中BN的均值，或者是方差都会存储在buffers
• ***_hooks：存储管理钩子函数（5个，暂时不去了解）
  

现在主要了解nn.Module的创建以及对属性的管理机制。以上面介绍的LeNet模型了解nn.Module的创建。

class LeNet(nn.Module):
    def __init__(self, classes):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, classes)

代码中，LeNet是继承nn.Module类的（class LeNet(nn.Module)）。super(LeNet, self).init()的作用是实现父类的函数调用功能，LeNet的父类是nn.Module，所以调用nn.Module的__init__()函数，进入super(LeNet, self).init()查看init()函数实现了什么操作。

def __init__(self):
    self._construct()
    # initialize self.training separately from the rest of the internal
    # state, as it is managed differently by nn.Module and ScriptModule
    self.training = True

module.py代码中的__init__()函数只有两行，第一行是self._construct()，进入construct()函数。

    def _construct(self):
        """
        Initializes internal Module state, shared by both nn.Module and ScriptModule.
        """
        torch._C._log_api_usage_once("python.nn_module")
        self._backend = thnn_backend
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._backward_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()

在construct()函数中实现了上面介绍的八个有序字典的初始化，这里主要关注其中的self._parameters和self._modules。

接着返回nn.Module()的__init__()函数中，代码中的self.training = True表示模型的训练状态，这样就构建好了一个module的基本属性。

从图中可以看到，LeNet就有了八个有序字典，可以看到字典都是空的。接着代码就开始构建子模块，第一个网络层是Conv2d的卷积层（self.conv1 = nn.Conv2d(3, 6, 5)），现在进入这个卷积层看看，进入Conv2d这个类中。

class Conv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(Conv2d, self).__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation,
            False, _pair(0), groups, bias, padding_mode)

Conv2d是继承于ConvNd这个类的，看一下init()函数实现什么操作。我们进入代码中的super(Conv2d, self).init(in_channels, out_channels, kernel_size, stride, padding, dilation,False, _pair(0), groups, bias, padding_mode)看看具体实现了什么操作。

class _ConvNd(Module):

    __constants__ = ['stride', 'padding', 'dilation', 'groups', 'bias',
                     'padding_mode', 'output_padding', 'in_channels',
                     'out_channels', 'kernel_size']

    def __init__(self, in_channels, out_channels, kernel_size, stride,
                 padding, dilation, transposed, output_padding,
                 groups, bias, padding_mode):
        super(_ConvNd, self).__init__()
        if in_channels % groups != 0:
            raise ValueError('in_channels must be divisible by groups')
        if out_channels % groups != 0:
            raise ValueError('out_channels must be divisible by groups')
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.transposed = transposed
        self.output_padding = output_padding
        self.groups = groups
        self.padding_mode = padding_mode
        if transposed:
            self.weight = Parameter(torch.Tensor(
                in_channels, out_channels // groups, *kernel_size))
        else:
            self.weight = Parameter(torch.Tensor(
                out_channels, in_channels // groups, *kernel_size))
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

进入ConvNd模块中，可以看到，ConvNd是继承于Module，代码中还是用了super(_ConvNd, self).init()调用module类的init()函数，这里就是为了构建八个有序字典，这样一个网络层就构建完毕，返回LeNet()函数中。因为nn.Conv2d是一个module，所以会存储在module字典中。如下图所示：

这样就实现了一个子网络层的构建。下面构建第二个网络层来观察module是如何将子module存储到modules字典中的，观察一下这个机制。

观察self.conv2 = nn.Conv2d(6, 16, 5)，具体的构建过程和第一个卷积层是一样的，当代码构建完成返回到代码self.conv2 = nn.Conv2d(6, 16, 5)时，这时候其实是还没有对self.conv2进行赋值的。这里只是实现了nn.Conv2d(6, 16, 5)的实例化，下一步才是赋值到属性self.conv2中，这里并不能直接赋值给self.conv2。因为在module中有一个机制，会拦截所有的类属性赋值操作，在赋值之前会跳转到module.py中的__setattr__()函数中，如下代码所示。

    def __setattr__(self, name, value):
        def remove_from(*dicts):
            for d in dicts:
                if name in d:
                    del d[name]

        params = self.__dict__.get('_parameters')
        if isinstance(value, Parameter):
            if params is None:
                raise AttributeError(
                    "cannot assign parameters before Module.__init__() call")
            remove_from(self.__dict__, self._buffers, self._modules)
            self.register_parameter(name, value)
        elif params is not None and name in params:
            if value is not None:
                raise TypeError("cannot assign '{}' as parameter '{}' "
                                "(torch.nn.Parameter or None expected)"
                                .format(torch.typename(value), name))
            self.register_parameter(name, value)
        else:
            modules = self.__dict__.get('_modules')
            if isinstance(value, Module):
                if modules is None:
                    raise AttributeError(
                        "cannot assign module before Module.__init__() call")
                remove_from(self.__dict__, self._parameters, self._buffers)
                modules[name] = value
            elif modules is not None and name in modules:
                if value is not None:
                    raise TypeError("cannot assign '{}' as child module '{}' "
                                    "(torch.nn.Module or None expected)"
                                    .format(torch.typename(value), name))
                modules[name] = value
            else:
                buffers = self.__dict__.get('_buffers')
                if buffers is not None and name in buffers:
                    if value is not None and not isinstance(value, torch.Tensor):
                        raise TypeError("cannot assign '{}' as buffer '{}' "
                                        "(torch.Tensor or None expected)"
                                        .format(torch.typename(value), name))
                    buffers[name] = value
                else:
                    object.__setattr__(self, name, value)

这个函数的功能是会拦截所有类属性的赋值，观察上面的代码，代码会对参数name和value的数据类型进行判断（if isinstance(value, Parameter):）。如果value是parameter的话，就会存储到self.register_parameter(name, value)中。因为我们要赋值的Conv2d是一个类，不是parameter类型。所以代码会往下继续运行，进入if isinstance(value, Module):判断，判断数据是不是一个module，因为要赋值的数据是一个module，所以会把数据存储到modules[name] = value当中，这样就完成了对LeNet当中module字典的更新。

其余的网络层的效果也一样，这样就完成了对LeNet的构建。

nn.Module的属性构建会在module类中进行属性赋值的时候会被setattr()函数拦截，在这个函数当中会判断即将要赋值的数据类型是否是nn.parameters类，如果是的话就会存储到parameters字典中；如果是module类就会存储到modul字典中。

3、nn.Module总结

• 一个module可以包含多个子module；例如LeNet包含很多个子module，例如卷积层，池化层等。
• 一个module相当于一个运算，必须实现forward()函数；
• 每个module都有8个字典管理它的属性；