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.

你可能感兴趣的