pytorch的python API略读--torch.nn.Module

作者：机器视觉全栈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.