Pytorch 快速搭建网络搭积木方法

前言

研究lightnet源代码时，看到这种技巧，惊为天人，于是单独摘出来。

感谢作者EAVISE，lightnet传送门。

 

一、 使用OrderedDict([ ])

import torch
import torch.nn as nn
from collections import OrderedDict


layer_list = [
            # Sequence 1 : 
            OrderedDict([
                ('1_conv2d',  nn.Conv2d(512, 64, 1, 1)),
                ('2_Relu',    nn.ReLU(inplace=True) ),
                ]),

            # Sequence 2 : 
            OrderedDict([
                ('3_conv2d',  nn.Conv2d((4*64)+1024, 1024, 3, 1)),
                ('4_bn',      nn.BatchNorm2d(1024, 16, 1, 1, 0)),
                ]),
            ]

sequence_list=[nn.Sequential(layer_dict) for layer_dict in layer_list]

是不是很简单，是不是很轻松写意？

我们看看原理。

 

这个是利用了nn.Sequential的构造函数，只需要传入一个OrderedDict对象，就能把里面的nn.module全部加进自己的seuqential序列中来。

详见Sequential的__init__方法：

class Sequential(Module):
    """
        # 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())
                ]))
    """

    def __init__(self, *args):
        super(Sequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)

 

于是，我们通过这句代码：

sequence_list=[nn.Sequential(layer_dict) for layer_dict in layer_list]

得到sequence_list ，其中任意一个元素，都是一个nn.Sequential类的实例对象。

 

更进一步的，我们可以把 list of nn.Sequential 包装为 nn.ModuleList

 module_list = nn.ModuleList(sequence_list)

这步也是利用到了nn.ModuleList的构造方法，传入一个list of nn.sequential，合并这个list。

#nn.ModuleList构造函数
def __init__(self, modules=None):
    super(ModuleList, self).__init__()
    if modules is not None:
        self += modules

可以看到直接用了self += xxx ，

这是因为 nn.ModuleList 类可以当成list来使用。

一般场景中没有太大区别，例如：

#nn.ModuleList可以当list来使用
def myforward(self,x):
    for _ , module in enumerate(module_list):
        x = module(x)

    return x

 

需要注意的是，我们构造OrderedDict()时，传入一个list of tuple

在二元组tuple内，有格式：('name', module)

 

对第一个元素 'name' 的要求是，

1.必须为str类型，判断方法是isinstance(name,torch._six.string_classes) ；

2.不能与已有的属性重名。判断方法是if hasattr(self, name) and name not in self._modules ；

3.不能含有点字符'.'，以免引起解释器歧义。判断方法是 if '.'  in name ；

5.不能为空。 判断方法是if name=='' ；

通过之后，就会在字典冲增加一条：

self._modules[name] = module

 

二、 与直接使用nn.Sequential()的区别

 

我们来看看两种构造方式的区别。

一般情况下，如果用不指定name的方法来构造Sequential，

默认的name会按照顺序依次为0,1,2,3,4......

model = nn.Sequential(
                  nn.Conv2d(1,20,5),
                  nn.ReLU(),
                  nn.Conv2d(20,64,5),
                  nn.ReLU()
                )

print(model.modules)

#output:
<bound method Module.modules of Sequential(
  (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (1): ReLU()
  (2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
  (3): ReLU()
)>

 

在方法一中，由于我们约定了name，output会变成如下形式，

<bound method Module.modules of Sequential(
  (1_conv2d): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1))
  (2_Relu): ReLU(inplace)
)>

 

 

三、用自定义layer充当积木

 

前面我们是用nn.Conv2d和nn.ReLU等公用积木搭网络。

那么能不能自己设计一些自定义的layer呢？

 

比如我们平时Conv2d之后又要batchnorm，又要LeakyReLU。

能不能直接设计一个layer，叫做ConvBNLeaky()，一块积木顶三块呢？

 

lightnet的作者真是奇才，当真做到了这一点。

我怎么就想不到呢。

 

看代码:

class Conv2dBatchLeaky(nn.Module):
    """ This convenience layer groups a 2D convolution, a batchnorm and a leaky ReLU.
    They are executed in a sequential manner.

    Args:
        in_channels (int): Number of input channels
        out_channels (int): Number of output channels
        kernel_size (int or tuple): Size of the kernel of the convolution
        stride (int or tuple): Stride of the convolution
        padding (int or tuple): padding of the convolution
        leaky_slope (number, optional): Controls the angle of the negative slope of the leaky ReLU; Default **0.1**
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, leaky_slope=0.1):
        super(Conv2dBatchLeaky, self).__init__()

        # Parameters
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        if isinstance(kernel_size, (list, tuple)):
            self.padding = [int(ii/2) for ii in kernel_size]
        else:
            self.padding = int(kernel_size/2)
        self.leaky_slope = leaky_slope

        # Layer
        self.layers = nn.Sequential(
            nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, self.stride, self.padding, bias=False),
            nn.BatchNorm2d(self.out_channels), #, eps=1e-6, momentum=0.01),
            nn.LeakyReLU(self.leaky_slope, inplace=True)
        )

    def __repr__(self):
        s = '{name} ({in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}, padding={padding}, negative_slope={leaky_slope})'
        return s.format(name=self.__class__.__name__, **self.__dict__)

    def forward(self, x):
        x = self.layers(x)
        return x

 

这样就可以了！！！

这样一来，Conv2dBatchLeaky就变成跟nn.Conv2d同等级的类了，因为也是继承自nn.Module。

 

之后可以用方法一来直接构造网络。

from xxx import Conv2dBatchLeaky

layer_1 = OrderedDict([
                ('1_conv2d',  Conv2dBatchLeaky(512, 64, 1, 1)),
                ])

sequence = nn.Sequential(layer_1)

一行顶三行，简单无脑易操作！