参考来源：
CSDN：pytorch 教程之 nn.Module 类详解——使用 Module 类来自定义模型
CSDN：pytorch 教程之 nn.Module 类详解——使用 Module 类来自定义网络层
知乎：【Pytorch】nn.module() 类解析及冻结特定层参数的方法

一、使用 Module 类来自定义模型

前言：pytorch 中对于一般的序列模型，直接使用 torch.nn.Sequential 类及可以实现，这点类似于keras，但是更多的时候面对复杂的模型，比如：多输入多输出、多分支模型、跨层连接模型、带有自定义层的模型等，就需要自己来定义一个模型了。本文将详细说明如何让使用 **Mudule** 类来自定义一个模型。

1. torch.nn.Module 类概述

个人理解，pytorch 不像 tensorflow 那么底层，也不像 keras 那么高层，这里先比较 keras 和 pytorch 的一些小区别。

  1. keras 更常见的操作是通过继承 Layer 类来实现自定义层，不推荐去继承 Model 类定义模型，详细原因可以参见官方文档
  2. pytorch 中其实一般没有特别明显的 Layer 和 Module 的区别，不管是自定义层、自定义块、自定义模型，都是通过继承 **Module** 类完成的，这一点很重要。其实 Sequential 类也是继承自 Module 类的。

注意：我们当然也可以直接通过继承 torch.autograd.Function 类来自定义一个层，但是这很不推荐，不提倡，至于为什么后面会介绍。
总结：pytorch 里面一切自定义操作基本上都是继承 nn.Module 类来实现的
这里仅仅先讨论使用 Module 来实现自定义模块，自定义层先不做讨论。

2. torch.nn.Module 类的简介

先来简单看一它的定义：

  1. class Module(object):
  2.     def __init__(self):
  3.     def forward(self, *input):
  5.     def add_module(self, name, module):
  6.     def cuda(self, device=None):
  7.     def cpu(self):
  8.     def __call__(self, *input, **kwargs):
  9.     def parameters(self, recurse=True):
  10.     def named_parameters(self, prefix='', recurse=True):
  11.     def children(self):
  12.     def named_children(self):
  13.     def modules(self):  
  14.     def named_modules(self, memo=None, prefix=''):
  15.     def train(self, mode=True):
  16.     def eval(self):
  17.     def zero_grad(self):
  18.     def __repr__(self):
  19.     def __dir__(self):
  20. '''
  21. 有一部分没有完全列出来
  22. '''

我们在定义自已的网络的时候，需要继承 nn.Module 类，并重新实现构造函数 **__init__** 和前向传播 **forward** 这两个方法。但有一些注意技巧：

  1. 一般把网络中具有可学习参数的层（如全连接层、卷积层等）放在构造函数 **__init__()** 中，当然我也可以把不具有参数的层也放在里面。
  2. 一般把不具有可学习参数的层(如 **ReLU**、**dropout**、**BatchNormanation** 层)可放在构造函数中，也可不放在构造函数中，如果不放在构造函数 __init__ 里面，则在 forward 方法里面可以使用 nn.functional 来代替。
  3. **forward** 方法是必须要重写的，它是实现模型的功能，实现各个层之间的连接关系的核心。

下面先看一个简单的例子。

  1. import torch
  3. class MyNet(torch.nn.Module):
  4.     def __init__(self):
  5.         super(MyNet, self).__init__()  # 第一句话，调用父类的构造函数
  6.         self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)
  7.         self.relu1=torch.nn.ReLU()
  8.         self.max_pooling1=torch.nn.MaxPool2d(2,1)
  10.         self.conv2 = torch.nn.Conv2d(3, 32, 3, 1, 1)
  11.         self.relu2=torch.nn.ReLU()
  12.         self.max_pooling2=torch.nn.MaxPool2d(2,1)
  14.         self.dense1 = torch.nn.Linear(32 * 3 * 3, 128)
  15.         self.dense2 = torch.nn.Linear(128, 10)
  17.     def forward(self, x):
  18.         x = self.conv1(x)
  19.         x = self.relu1(x)
  20.         x = self.max_pooling1(x)
  21.         x = self.conv2(x)
  22.         x = self.relu2(x)
  23.         x = self.max_pooling2(x)
  24.         x = self.dense1(x)
  25.         x = self.dense2(x)
  26.         return x
  28. model = MyNet()
  29. print(model)
  32. '''
  33. 运行结果为：
  34. MyNet(
  35.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  36.   (relu1): ReLU()
  37.   (max_pooling1): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
  38.   (conv2): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  39.   (relu2): ReLU()
  40.   (max_pooling2): MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False)
  41.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  42.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  43. )
  44. '''

注意：上面的是将所有的层都放在了构造函数 __init__ 里面，但是只是定义了一系列的层，各个层之间到底是什么连接关系并没有，而是在 forward 里面实现所有层的连接关系，当然这里依然是顺序连接的。下面再来看一下一个例子：

  1. import torch
  2. import torch.nn.functional as F
  4. class MyNet(torch.nn.Module):
  5.     def __init__(self):
  6.         super(MyNet, self).__init__()  # 第一句话，调用父类的构造函数
  7.         self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)
  8.         self.conv2 = torch.nn.Conv2d(3, 32, 3, 1, 1)
  10.         self.dense1 = torch.nn.Linear(32 * 3 * 3, 128)
  11.         self.dense2 = torch.nn.Linear(128, 10)
  13.     def forward(self, x):
  14.         x = self.conv1(x)
  15.         x = F.relu(x)
  16.         x = F.max_pool2d(x)
  17.         x = self.conv2(x)
  18.         x = F.relu(x)
  19.         x = F.max_pool2d(x)
  20.         x = self.dense1(x)
  21.         x = self.dense2(x)
  22.         return x
  24. model = MyNet()
  25. print(model)
  28. '''
  29. 运行结果为：
  30. MyNet(
  31.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  32.   (conv2): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  33.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  34.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  35. )
  36. '''

注意：此时，将没有训练参数的层没有放在构造函数里面了，所以这些层就不会出现在 model 里面，但是运行关系是在 forward 里面通过 nn.functional 中的方法实现的。

总结：所有放在构造函数 **__init__** 里面的层的都是这个模型的“固有属性”。

3. torch.nn.Module 类的的多种实现

上面是为了一个简单的演示，但是Module类是非常灵活的，可以有很多灵活的实现方式，下面将一一介绍。

3.1 通过 Sequential 来包装层

即将几个层包装在一起作为一个大的层（块），前面的一篇文章详细介绍了 Sequential 类的使用，包括常见的三种方式，以及每一种方式的优缺点，参见语雀：torch.nn.Sequential 类。

所以这里对层的包装当然也可以通过这三种方式了。

方法一

语雀：最简单的序贯模型

  1. import torch.nn as nn
  2. from collections import OrderedDict
  3. class MyNet(nn.Module):
  4.     def __init__(self):
  5.         super(MyNet, self).__init__()
  6.         self.conv_block = nn.Sequential(
  7.             nn.Conv2d(3, 32, 3, 1, 1),
  8.             nn.ReLU(),
  9.             nn.MaxPool2d(2))
  10.         self.dense_block = nn.Sequential(
  11.             nn.Linear(32 * 3 * 3, 128),
  12.             nn.ReLU(),
  13.             nn.Linear(128, 10)
  14.         )
  15.     # 在这里实现层之间的连接关系，其实就是所谓的前向传播
  16.     def forward(self, x):
  17.         conv_out = self.conv_block(x)
  18.         res = conv_out.view(conv_out.size(0), -1)
  19.         out = self.dense_block(res)
  20.         return out
  22. model = MyNet()
  23. print(model)
  26. '''
  27. 运行结果为：
  28. MyNet(
  29.   (conv_block): Sequential(
  30.     (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  31.     (1): ReLU()
  32.     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  33.   )
  34.   (dense_block): Sequential(
  35.     (0): Linear(in_features=288, out_features=128, bias=True)
  36.     (1): ReLU()
  37.     (2): Linear(in_features=128, out_features=10, bias=True)
  38.   )
  39. )
  40. '''

同前面的文章，这里在每一个包装块里面，各个层是没有名称的，默认按照 0、1、2、3、4 来排名。

方法二

语雀：给每一个层添加名称

  1. import torch.nn as nn
  2. from collections import OrderedDict
  3. class MyNet(nn.Module):
  4.     def __init__(self):
  5.         super(MyNet, self).__init__()
  6.         self.conv_block = nn.Sequential(
  7.             OrderedDict(
  8.                 [
  9.                     ("conv1", nn.Conv2d(3, 32, 3, 1, 1)),
  10.                     ("relu1", nn.ReLU()),
  11.                     ("pool", nn.MaxPool2d(2))
  12.                 ]
  13.             ))
  15.         self.dense_block = nn.Sequential(
  16.             OrderedDict([
  17.                 ("dense1", nn.Linear(32 * 3 * 3, 128)),
  18.                 ("relu2", nn.ReLU()),
  19.                 ("dense2", nn.Linear(128, 10))
  20.             ])
  21.         )
  23.     def forward(self, x):
  24.         conv_out = self.conv_block(x)
  25.         res = conv_out.view(conv_out.size(0), -1)
  26.         out = self.dense_block(res)
  27.         return out
  29. model = MyNet()
  30. print(model)
  33. '''
  34. 运行结果为：
  35. MyNet(
  36.   (conv_block): Sequential(
  37.     (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  38.     (relu1): ReLU()
  39.     (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  40.   )
  41.   (dense_block): Sequential(
  42.     (dense1): Linear(in_features=288, out_features=128, bias=True)
  43.     (relu2): ReLU()
  44.     (dense2): Linear(in_features=128, out_features=10, bias=True)
  45.   )
  46. )
  47. '''

方法三

[model.add_module(self, name, module)](https://www.yuque.com/yuque-qsztn/va7nxh/znemg4#y1Jzx)

  1. import torch.nn as nn
  2. from collections import OrderedDict
  3. class MyNet(nn.Module):
  4.     def __init__(self):
  5.         super(MyNet, self).__init__()
  6.         self.conv_block=torch.nn.Sequential()
  7.         self.conv_block.add_module("conv1",torch.nn.Conv2d(3, 32, 3, 1, 1))
  8.         self.conv_block.add_module("relu1",torch.nn.ReLU())
  9.         self.conv_block.add_module("pool1",torch.nn.MaxPool2d(2))
  11.         self.dense_block = torch.nn.Sequential()
  12.         self.dense_block.add_module("dense1",torch.nn.Linear(32 * 3 * 3, 128))
  13.         self.dense_block.add_module("relu2",torch.nn.ReLU())
  14.         self.dense_block.add_module("dense2",torch.nn.Linear(128, 10))
  16.     def forward(self, x):
  17.         conv_out = self.conv_block(x)
  18.         res = conv_out.view(conv_out.size(0), -1)
  19.         out = self.dense_block(res)
  20.         return out
  22. model = MyNet()
  23. print(model)
  26. '''
  27. 运行结果为：
  28. MyNet(
  29.   (conv_block): Sequential(
  30.     (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  31.     (relu1): ReLU()
  32.     (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  33.   )
  34.   (dense_block): Sequential(
  35.     (dense1): Linear(in_features=288, out_features=128, bias=True)
  36.     (relu2): ReLU()
  37.     (dense2): Linear(in_features=128, out_features=10, bias=True)
  38.   )
  39. )
  40. '''

上面的方式二和方式三，在每一个包装块里面，每个层都是有名称的。

3.2 Module 类的几个常见方法使用

特别注意：Sequential 类虽然继承自 Module 类，二者有相似部分，但是也有很多不同的部分，集中体现在：
**Sequenrial** 类实现了整数索引，故而可以使用 **model[index]** 这样的方式获取一个曾，但是 **Module** 类并没有实现整数索引，不能够通过整数索引来获得层，那该怎么办呢？它提供了几个主要的方法，如下：

  1. def children(self):
  3. def named_children(self):
  5. def modules(self):
  7. def named_modules(self, memo=None, prefix=''):
  9. '''
  10. 注意：这几个方法返回的都是一个Iterator迭代器，故而通过for循环访问，当然也可以通过next
  11. '''

1. model.children() 和 model.named_children() 方法

  1. model.children() 和 model.named_children() 方法返回的是迭代器 **iterator** ；
  2. **model.children()**：每一次迭代返回的每一个元素实际上是 **Sequential** 类型，而 Sequential 类型又可以使用下标 index 索引来获取每一个 Sequenrial 里面的具体层，比如 conv 层、dense 层等；
  3. **model.named_children()**：每一次迭代返回的每一个元素实际上是 一个元组类型，元组的第一个元素是名称，第二个元素就是对应的层或者是 Sequential 。

2. model.children() 方法

  1. import torch
  2. import torch.nn as nn
  3. from collections import OrderedDict
  6. class MyNet(nn.Module):
  7.     def __init__(self):
  8.         super(MyNet, self).__init__()
  9.         self.conv_block = torch.nn.Sequential()
  10.         self.conv_block.add_module("conv1", torch.nn.Conv2d(3, 32, 3, 1, 1))
  11.         self.conv_block.add_module("relu1", torch.nn.ReLU())
  12.         self.conv_block.add_module("pool1", torch.nn.MaxPool2d(2))
  14.         self.dense_block = torch.nn.Sequential()
  15.         self.dense_block.add_module("dense1", torch.nn.Linear(32 * 3 * 3, 128))
  16.         self.dense_block.add_module("relu2", torch.nn.ReLU())
  17.         self.dense_block.add_module("dense2", torch.nn.Linear(128, 10))
  19.     def forward(self, x):
  20.         conv_out = self.conv_block(x)
  21.         res = conv_out.view(conv_out.size(0), -1)
  22.         out = self.dense_block(res)
  23.         return out
  26. model = MyNet()
  28. for i in model.children():
  29.     print(i)
  30.     print(type(i)) # 查看每一次迭代的元素到底是什么类型，实际上是 Sequential 类型,所以有可以使用下标 index 索引来获取每一个 Sequenrial 里面的具体层
  33. '''
  34. 运行结果为：
  35. Sequential(
  36.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  37.   (relu1): ReLU()
  38.   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  39. )
  40. <class 'torch.nn.modules.container.Sequential'>
  41. Sequential(
  42.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  43.   (relu2): ReLU()
  44.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  45. )
  46. <class 'torch.nn.modules.container.Sequential'>
  47. '''

3. model.named_children() 方法

  1. import torch
  2. import torch.nn as nn
  3. from collections import OrderedDict
  6. class MyNet(nn.Module):
  7.     def __init__(self):
  8.         super(MyNet, self).__init__()
  9.         self.conv_block = torch.nn.Sequential()
  10.         self.conv_block.add_module("conv1", torch.nn.Conv2d(3, 32, 3, 1, 1))
  11.         self.conv_block.add_module("relu1", torch.nn.ReLU())
  12.         self.conv_block.add_module("pool1", torch.nn.MaxPool2d(2))
  14.         self.dense_block = torch.nn.Sequential()
  15.         self.dense_block.add_module("dense1", torch.nn.Linear(32 * 3 * 3, 128))
  16.         self.dense_block.add_module("relu2", torch.nn.ReLU())
  17.         self.dense_block.add_module("dense2", torch.nn.Linear(128, 10))
  19.     def forward(self, x):
  20.         conv_out = self.conv_block(x)
  21.         res = conv_out.view(conv_out.size(0), -1)
  22.         out = self.dense_block(res)
  23.         return out
  26. model = MyNet()
  28. for i in model.named_children():
  29.     print(i)
  30.     print(type(i)) # 查看每一次迭代的元素到底是什么类型，实际上是 返回一个 tuple,tuple 的第一个元素是名称
  33. '''
  34. 运行结果为：
  35. ('conv_block', Sequential(
  36.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  37.   (relu1): ReLU()
  38.   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  39. ))
  40. <class 'tuple'>
  41. ('dense_block', Sequential(
  42.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  43.   (relu2): ReLU()
  44.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  45. ))
  46. <class 'tuple'>
  47. '''

4. model.modules() 和 model.named_modules()

  1. model.modules() 和 model.named_modules() 方法返回的是迭代器 iterator；
  2. model 的 modules() 方法和 named_modules() 方法都会将整个模型的所有构成（包括包装层、单独的层、自定义层等）由浅入深依次遍历出来，只不过 modules() 返回的每一个元素是直接返回的层对象本身，而 named_modules() 返回的每一个元素是一个元组，第一个元素是名称，第二个元素才是层对象本身。
  3. 如何理解 children和 modules 之间的这种差异性。注意 pytorch 里面不管是模型、层、激活函数、损失函数都可以当成是 Module 的拓展，所以 modules 和 named_modules 会层层迭代，由浅入深，将每一个自定义块 block 、然后 block 里面的每一个层都当成是 module 来迭代。而 children 就比较直观，就表示的是所谓的“孩子”，所以没有层层迭代深入。

5. model.modules() 方法

  1. import torch
  2. import torch.nn as nn
  3. from collections import OrderedDict
  6. class MyNet(nn.Module):
  7.     def __init__(self):
  8.         super(MyNet, self).__init__()
  9.         self.conv_block = torch.nn.Sequential()
  10.         self.conv_block.add_module("conv1", torch.nn.Conv2d(3, 32, 3, 1, 1))
  11.         self.conv_block.add_module("relu1", torch.nn.ReLU())
  12.         self.conv_block.add_module("pool1", torch.nn.MaxPool2d(2))
  14.         self.dense_block = torch.nn.Sequential()
  15.         self.dense_block.add_module("dense1", torch.nn.Linear(32 * 3 * 3, 128))
  16.         self.dense_block.add_module("relu2", torch.nn.ReLU())
  17.         self.dense_block.add_module("dense2", torch.nn.Linear(128, 10))
  19.     def forward(self, x):
  20.         conv_out = self.conv_block(x)
  21.         res = conv_out.view(conv_out.size(0), -1)
  22.         out = self.dense_block(res)
  23.         return out
  26. model = MyNet()
  28. for i in model.modules():
  29.     print(i)
  30.     print("==================================================")
  33. '''
  34. 运行结果为：
  35. MyNet(
  36.   (conv_block): Sequential(
  37.     (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  38.     (relu1): ReLU()
  39.     (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  40.   )
  41.   (dense_block): Sequential(
  42.     (dense1): Linear(in_features=288, out_features=128, bias=True)
  43.     (relu2): ReLU()
  44.     (dense2): Linear(in_features=128, out_features=10, bias=True)
  45.   )
  46. )
  47. ==================================================
  48. Sequential(
  49.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  50.   (relu1): ReLU()
  51.   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  52. )
  53. ==================================================
  54. Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  55. ==================================================
  56. ReLU()
  57. ==================================================
  58. MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  59. ==================================================
  60. Sequential(
  61.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  62.   (relu2): ReLU()
  63.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  64. )
  65. ==================================================
  66. Linear(in_features=288, out_features=128, bias=True)
  67. ==================================================
  68. ReLU()
  69. ==================================================
  70. Linear(in_features=128, out_features=10, bias=True)
  71. ==================================================
  72. '''

6. model.named_modules() 方法

  1. import torch
  2. import torch.nn as nn
  3. from collections import OrderedDict
  6. class MyNet(nn.Module):
  7.     def __init__(self):
  8.         super(MyNet, self).__init__()
  9.         self.conv_block = torch.nn.Sequential()
  10.         self.conv_block.add_module("conv1", torch.nn.Conv2d(3, 32, 3, 1, 1))
  11.         self.conv_block.add_module("relu1", torch.nn.ReLU())
  12.         self.conv_block.add_module("pool1", torch.nn.MaxPool2d(2))
  14.         self.dense_block = torch.nn.Sequential()
  15.         self.dense_block.add_module("dense1", torch.nn.Linear(32 * 3 * 3, 128))
  16.         self.dense_block.add_module("relu2", torch.nn.ReLU())
  17.         self.dense_block.add_module("dense2", torch.nn.Linear(128, 10))
  19.     def forward(self, x):
  20.         conv_out = self.conv_block(x)
  21.         res = conv_out.view(conv_out.size(0), -1)
  22.         out = self.dense_block(res)
  23.         return out
  26. model = MyNet()
  28. for i in model.named_modules():
  29.     print(i)
  30.     print("==================================================")
  33. '''
  34. 运行结果是：
  35. ('', MyNet(
  36.   (conv_block): Sequential(
  37.     (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  38.     (relu1): ReLU()
  39.     (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  40.   )
  41.   (dense_block): Sequential(
  42.     (dense1): Linear(in_features=288, out_features=128, bias=True)
  43.     (relu2): ReLU()
  44.     (dense2): Linear(in_features=128, out_features=10, bias=True)
  45.   )
  46. ))
  47. ==================================================
  48. ('conv_block', Sequential(
  49.   (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  50.   (relu1): ReLU()
  51.   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  52. ))
  53. ==================================================
  54. ('conv_block.conv1', Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))
  55. ==================================================
  56. ('conv_block.relu1', ReLU())
  57. ==================================================
  58. ('conv_block.pool1', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))
  59. ==================================================
  60. ('dense_block', Sequential(
  61.   (dense1): Linear(in_features=288, out_features=128, bias=True)
  62.   (relu2): ReLU()
  63.   (dense2): Linear(in_features=128, out_features=10, bias=True)
  64. ))
  65. ==================================================
  66. ('dense_block.dense1', Linear(in_features=288, out_features=128, bias=True))
  67. ==================================================
  68. ('dense_block.relu2', ReLU())
  69. ==================================================
  70. ('dense_block.dense2', Linear(in_features=128, out_features=10, bias=True))
  71. ==================================================
  72. '''

注意：上面这四个方法是以层包装为例来说明的，如果没有层的包装，我们依然可以使用这四个方法，其实结果也是类似的这样去推，这里就不再列出来了。

二、使用 Module 类来自定义网络层

前言：前面介绍了如何自定义一个模型——通过继承 nn.Module 类来实现，在 __init__ 构造函数中申明各个层的定义，在 forward 中实现层之间的连接关系，实际上就是前向传播的过程。
事实上，在 pytorch 里面自定义层也是通过继承自 nn.Module 类来实现的，我前面说过，**pytorch** 里面一般是没有层的概念，层也是当成一个模型来处理的，这里和 keras 是不一样的。前面介绍过，我们当然也可以直接通过继承 torch.autograd.Function 类来自定义一个层，但是这很不推荐，不提倡，至于为什么后面会介绍。记住一句话，keras 更加注重的是层 Layer、pytorch 更加注重的是模型 Module 。
所以本文就专门来介绍如何通过 **nn.Module** 类来实现自定义层。

1. 从系统预定义的层说起

1.1 Linear 层的代码

  1. import math
  2. import torch
  3. from torch.nn.parameter import Parameter
  4. from .. import functional as F
  5. from .. import init
  6. from .module import Module
  7. from ..._jit_internal import weak_module, weak_script_method
  9. class Linear(Module):
  10.     __constants__ = ['bias']
  12.     def __init__(self, in_features, out_features, bias=True):
  13.         super(Linear, self).__init__()
  14.         self.in_features = in_features
  15.         self.out_features = out_features
  16.         self.weight = Parameter(torch.Tensor(out_features, in_features))
  17.         if bias:
  18.             self.bias = Parameter(torch.Tensor(out_features))
  19.         else:
  20.             self.register_parameter('bias', None)
  21.         self.reset_parameters()
  23.     def reset_parameters(self):
  24.         init.kaiming_uniform_(self.weight, a=math.sqrt(5))
  25.         if self.bias is not None:
  26.             fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
  27.             bound = 1 / math.sqrt(fan_in)
  28.             init.uniform_(self.bias, -bound, bound)
  30.     @weak_script_method
  31.     def forward(self, input):
  32.         return F.linear(input, self.weight, self.bias)
  34.     def extra_repr(self):
  35.         return 'in_features={}, out_features={}, bias={}'.format(
  36.             self.in_features, self.out_features, self.bias is not None
  37.         )

1.2 Conv2d 类的实现

  1. class Conv2d(_ConvNd):
  2.     def __init__(self, in_channels, out_channels, kernel_size, stride=1,
  3.                  padding=0, dilation=1, groups=1,
  4.                  bias=True, padding_mode='zeros'):
  5.         kernel_size = _pair(kernel_size)
  6.         stride = _pair(stride)
  7.         padding = _pair(padding)
  8.         dilation = _pair(dilation)
  9.         super(Conv2d, self).__init__(
  10.             in_channels, out_channels, kernel_size, stride, padding, dilation,
  11.             False, _pair(0), groups, bias, padding_mode)
  13.     @weak_script_method
  14.     def forward(self, input):
  15.         if self.padding_mode == 'circular':
  16.             expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
  17.                                 (self.padding[0] + 1) // 2, self.padding[0] // 2)
  18.             return F.conv2d(F.pad(input, expanded_padding, mode='circular'),
  19.                             self.weight, self.bias, self.stride,
  20.                             _pair(0), self.dilation, self.groups)
  21.         return F.conv2d(input, self.weight, self.bias, self.stride,
  22.                         self.padding, self.dilation, self.groups)

1.3 初步总结

我在前面的文章里面说过，torch 里面实现神经网络有两种方式

  1. 高层 **API** 方法：使用 torch.nn.**** 来实现；
  2. 低层 **API** 方法：使用低层函数方法，torch.nn.functional.**** 来实现；

其中，我们推荐使用高层 **API** 的方法，原因如下：
高层 API 是使用类的形式来包装的，既然是类就可以存储参数，比如全连接层的权值矩阵、偏置矩阵等都可以作为类的属性存储着，但是低层 API 仅仅是实现函数的运算功能，没办法保存这些信息，会丢失参数信息，但是高层 API 是依赖于低层 API 的计算函数的，比如上面的两个层：

  • **Linear** 高级层——>低层 **F.linear()** 函数。
  • **Conv2d** 高级层——>低层 **F.conv2d()** 函数。

1.4 自定义层的步骤

要实现一个自定义层大致分以下几个主要的步骤：

  1. 自定义一个类，继承自 **Module** 类，并且一定要实现两个基本的函数，第一是构造函数 __init__ ，第二个是层的逻辑运算函数，即所谓的前向计算函数 forward 函数。
  2. 在构造函数 init_ 中实现层的参数定义。比如 Linear 层的权重和偏置，Conv2d 层的 in_channels、out_channels、kernel_size、stride=1、padding=0、dilation=1、groups=1、bias=True、padding_mode='zeros' 这一系列参数；
  3. 在前向传播 **forward** 函数里面实现前向运算。这一般都是通过 torch.nn.functional.*** 函数来实现，当然很多时候我们也需要自定义自己的运算方式。如果该层含有权重，那么权重必须是 nn.Parameter 类型，关于Tensor 和 Variable（0.3版本之前）与 Parameter 的区别请参阅相关的文档。简单说就是 Parameter 默认需要求导，其他两个类型则不会。另外一般情况下，可能的话，为自己定义的新层提供默认的参数初始化，以防使用过程中忘记初始化操作。
  4. 补充：一般情况下，我们定义的参数是可以求导的，但是自定义操作如不可导，需要实现 backward 函数。

总结：这里其实和定义一个自定义模型是一样的，核心都是实现最基本的构造函数 __init__ 和前向运算函数 forward 函数，可以参考上面的 语雀：一、使用 Module 类来自定义模型 。

2. 自定义层的简单例子

比如我要实现一个简单的层，这个层的功能是 y=w*sqrt(x2+bias)，即输入 X 的平方再加上一个偏执项，再开根号，然后再乘以权值矩阵 w，那要怎么做呢，按照上面的定义过程，我们先定义一个这样的层（即一个类），代码如下：

2.1 定义一个自定义层 MyLayer

  1. # 定义一个 my_layer.py
  2. import torch
  4. class MyLayer(torch.nn.Module):
  5.     '''
  6.     因为这个层实现的功能是：y=weights*sqrt(x**2+bias),所以有两个参数：
  7.     权值矩阵weights
  8.     偏置矩阵bias
  9.     输入 x 的维度是（in_features,)
  10.     输出 y 的维度是（out_features,) 故而
  11.     bias 的维度是（in_fearures,)，注意这里为什么是in_features,而不是out_features，注意体会这里和Linear层的区别所在
  12.     weights 的维度是（in_features, out_features）注意这里为什么是（in_features, out_features）,而不是（out_features, in_features），注意体会这里和Linear层的区别所在
  13.     '''
  14.     def __init__(self, in_features, out_features, bias=True):
  15.         super(MyLayer, self).__init__()  # 和自定义模型一样，第一句话就是调用父类的构造函数
  16.         self.in_features = in_features
  17.         self.out_features = out_features
  18.         self.weight = torch.nn.Parameter(torch.Tensor(in_features, out_features)) # 由于weights是可以训练的，所以使用Parameter来定义
  19.         if bias:
  20.             self.bias = torch.nn.Parameter(torch.Tensor(in_features))             # 由于bias是可以训练的，所以使用Parameter来定义
  21.         else:
  22.             self.register_parameter('bias', None)
  24.     def forward(self, input):
  25.         input_=torch.pow(input,2)+self.bias
  26.         y=torch.matmul(input_,self.weight)
  27.         return y

2.2 自定义模型并且训练

  1. import torch
  2. from my_layer import MyLayer # 自定义层
  5. N, D_in, D_out = 10, 5, 3  # 一共10组样本，输入特征为5，输出特征为3 
  7. # 先定义一个模型
  8. class MyNet(torch.nn.Module):
  9.     def __init__(self):
  10.         super(MyNet, self).__init__()  # 第一句话，调用父类的构造函数
  11.         self.mylayer1 = MyLayer(D_in,D_out)
  13.     def forward(self, x):
  14.         x = self.mylayer1(x)
  16.         return x
  18. model = MyNet()
  19. print(model)
  22. '''
  23. 运行结果为：
  24. MyNet(
  25.   (mylayer1): MyLayer()   # 这就是自己定义的一个层
  26. )
  27. '''

下面开始训练

  1. # 创建输入、输出数据
  2. x = torch.randn(N, D_in)  #（10，5）
  3. y = torch.randn(N, D_out) #（10，3）
  6. #定义损失函数
  7. loss_fn = torch.nn.MSELoss(reduction='sum')
  9. learning_rate = 1e-4
  10. #构造一个optimizer对象
  11. optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
  13. for t in range(10): # 
  15.     # 第一步：数据的前向传播，计算预测值p_pred
  16.     y_pred = model(x)
  18.     # 第二步：计算计算预测值p_pred与真实值的误差
  19.     loss = loss_fn(y_pred, y)
  20.     print(f"第 {t} 个epoch, 损失是 {loss.item()}")
  22.     # 在反向传播之前，将模型的梯度归零，这
  23.     optimizer.zero_grad()
  25.     # 第三步：反向传播误差
  26.     loss.backward()
  28.     # 直接通过梯度一步到位，更新完整个网络的训练参数
  29.     optimizer.step()

程序的运行结果为：

  1. 第 0 个epoch, 损失是 29.241456985473633
  2. 第 1 个epoch, 损失是 29.223047256469727
  3. 第 2 个epoch, 损失是 29.20465850830078
  4. 第 3 个epoch, 损失是 29.186279296875
  5. 第 4 个epoch, 损失是 29.167924880981445
  6. 第 5 个epoch, 损失是 29.14959716796875
  7. 第 6 个epoch, 损失是 29.131284713745117
  8. 第 7 个epoch, 损失是 29.112987518310547
  9. 第 8 个epoch, 损失是 29.094717025756836
  10. 第 9 个epoch, 损失是 29.076465606689453

总结：

本文的实践说明了如何使用 **Module** 父类来拓展实现自定义模型、自定义层，我们发现二者有异曲同工之处，这也是 pytorch 如此受欢迎的原因之一了。后面还会继续讲解通过 Function 来自定义一个层。需要注意的是：
**Function** 与 **Module** 都可以对 **pytorch** 进行自定义拓展，使其满足网络的需求，但这两者还是有十分重要的不同，具体的不同后面再说。