作者：机器视觉全栈er
网站：www.cvtutorials.com
2.3 torch.nn
本节我们看下torch.nn，这个模块是基础模块。
2.3.1 容器
torch.nn.Module：这个类是所有网络模块的基类，甚至你写普通的数学函数的时候，也可以用这个作为基类。内置的模块和你自定义的模块都可以相互嵌套，构成一个树状的模块结构。下面我们来看下这个Module类常用案例：
import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 3, 3)
self.conv2 = nn.Conv2d(3, 3, 3)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
my_module = MyModule()
cvtutorials = torch.randn(3, 1, 5, 5)
print("=====================================")
print(cvtutorials)
print("=====================================")
print(my_module(cvtutorials))
除此以外，该类有方法eval()和train()，运行两种不同的方法会分别进入不同的模式，在训练的时候要进入train模式，在测试的时候要进入eval模式。两者常见区别在于Dropout和BatchNorm，train模式下Dropout和BatchNorm都会起作用，而在eval模式下，Dropout和BatchNorm会失效。
Module类里面有个重要的方法是forward，当自定义一个模块的时候，必须重写这个方法。
在Module类中，还有几个有意思的方法，这里统称为Hook函数吧，他们分别是register_forward_hook(hook)、register_forward_pre_hook(hook)、register_full_backward_hook(hook)（register_backward_hook(hook)这个函数会逐步淘汰）。register是注册的意思，可以将hook函数注册到module或者module的子模块里面，这些hook的函数会在特定时间被调用，这些调用时间指的是forward函数运行前、forward函数运行后、backward函数运行后。这里用个简单的例子解释下：
import torch
import torch.nn as nn
import torch.nn.functional as F
def my_pre_hook(module, data_input):
print("I am function before forward function.")
def my_after_hook(module, data_input, data_output):
print("I am function after forward function.")
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 3, 3)
self.conv2 = nn.Conv2d(3, 3, 3)
def forward(self, x):
x = F.relu(self.conv1(x))
# print("hello, I am forward function.")
x = F.relu(self.conv2(x))
return x
my_module = MyModule()
my_module.conv1.register_forward_pre_hook(my_pre_hook)
my_module.conv2.register_forward_hook(my_after_hook)
cvtutorials = torch.randn(3, 1, 5, 5)
print("=================")
my_module(cvtutorials)
上面的函数运行结果如下：
=================
I am function before forward function.
I am function after forward function.
