Pytorch第五课：package-torch.nn详解（1）之Containers（容器）

1 class torch.nn.Module

torch.nn.Module是所有网络的基类。我们创建的任何模型都应该继承这个类。

Modules也可以包含其它Modules,允许使用树结构嵌入他们。你可以将子模块赋值给模型属性。

1.1 构建模型

现在，让我们来继承Module来建立一个超简单的模型吧。

import torch.nn as nn
import torch.nn.functional as F  # funcation类再之后会详细讲到

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 建立2个卷积子模型
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)
        
    def forward(self, x):
        # 构建前向计算过程
        x = F.relu(self.conv1(x))
        y = F.relu(self.conv2(x))
        
        return y

以上，就构建了一个模型Model， 这个模型中有两个子模型submodule：conv1，conv2

通过上面方式赋值的submodule会被注册。当调用 .cuda() 的时候，submodule的参数也会转换为cuda Tensor。

1.2 添加子模型：add_module(name, module)

将一个 child module 添加到当前 modle。 被添加的module可以通过 name属性来获取。

例如：

import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 通过add_module添加子模型
        self.add_module("conv1", nn.Conv2d(10, 20, 4))
        # 也可以用上面一节中的方式添加，两者等价
        self.conv2 = nn.Conv2d(10, 20, 4)
        
model = Model()
print(model)

Model(
  (conv1): Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))
  (conv2): Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))
)

可以直接通过名字取出子模型：

print(model.conv1)

Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))

print(model.conv2)

Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))

1.3 获取子模型：children()

返回当前模型 子模块的迭代器。

首先还是建立一个模型类：

import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.add_module("conv", nn.Conv2d(10, 20, 4))
        self.add_module("conv1", nn.Conv2d(20 ,10, 4