torch.nn.Module

torch.nn.Module
- 函数使用

torch.nn.Module

所有神经网络模块的基类。我们的模型也应该继承这个类。模块还可以包含其他模块，允许将它们嵌套在树结构中。我们可以将子模块分配为常规属性：

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

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Note:根据上面的示例，必须在对子类进行赋值之前对父类进行 init() 调用。

函数使用

函数add_module(name,module):

作用是在原有网络结构中加入一层网络

参数：name，即为模块的名字
参数：module，需要加入网络里的模块
用法：

from torch import nn
from torchsummary import summary

class Net_test(nn.Module):
    def __init__(self):
        super(Net_test,self).__init__()
        self.conv_1 = nn.Conv2d(3,6,3)
        self.add_module('conv_2', nn.Conv2d(6,12,3))
        self.conv_3 = nn.Conv2d(12,24,3)
        
    def forward(self,x):
        x = self.conv_1(x)
        x = self.conv_2(x)
        x = self.conv_3(x)
        return x

这个用法是从一篇博客中看来的，感觉有些鸡肋，于是自己做了一个新的测试：

import torch
from torch import nn
import torch.nn.functional as F
#%%
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(8, 8, kernel_size=3,padding=1)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3,padding=1,stride=2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

model = Model()
model.add_module("flatten",nn.Flatten())
model.add_module("linear1",nn.Linear(8*224*224,2))
#%%
print(model)
#这是输出
Model(
  (conv1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(8, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear1): Linear(in_features=401408, out_features=2, bias=True)
)

#%%
x = torch.rand(1,8,224,224)
out = model(x)
print(out.shape)
#输出
torch.Size([1, 16, 112, 112])

很明显看到，后面加入的模块并没有加入计算，原因是forward函数没有使用到后面加入的模块这里先mark一下，看看之后有没有办法修改forward函数

函数apply(fn)

作用是递归地将 fn 应用于每个子模块（由 .children() 返回）以及 self。典型用途包括初始化模型的参数（另见 torch.nn.init）。

参数：fn，fn是一个要应用到每个子模块的函数
返回值是网络自身，返回类型module

官网的例子

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)