Pytorch nn.Module的基本使用

nn.Module的基本用法

nn.Module是所有神经网络的基类，所以你的神经网络类也应该要继承这个基类

当使用时，主要需要实现其两个方法：

1. __init__：初始化方法，可以用于定义神经网络的结构
2. forward: 前向传播，用于定义神经网络的前向传播

例如，一个最简单的使用（非神经网络）：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        print("初始化模型")        
        self.model = nn.Linear(1, 1)

    def forward(self, x): # 定义前向传播方式，参数可以任意指定
        print("模型开始进行前向传播")
        return self.model(x)

模型定义好就可以使用了，直接调用构造函数即可：

model = Model()

初始化模型

定义好模型后，直接向model传入参数，就会执行forward方法，并返回forward方法的返回值

model(torch.Tensor([1]))

模型开始进行前向传播





tensor([0.7308], grad_fn=)

nn.Module的其他常用方法

1. cuda(device=None): 如果要使用GPU进行模型运算，需要将“模型参数和缓存”转移到GPU中。需要下载GPU版本的pytorch，并且安装CUDA

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model.cuda(device) # 需要安装cuda
model.to(device) # 推荐使用该方式

Model(
  (model): Linear(in_features=1, out_features=1, bias=True)
)

2. eval()：如果模型训练完毕，需要进行评估，需要显式的告知模型。调用该方法即可。

model.eval()
model.train(False) # 效果一样

Model(
  (model): Linear(in_features=1, out_features=1, bias=True)
)

3. modules(): 获取模型中的所有Module类型的对象，返回一个迭代器

for idx, m in enumerate(model.modules()):
    print(idx, '->', m)

0 -> Model(
  (model): Linear(in_features=1, out_features=1, bias=True)
)
1 -> Linear(in_features=1, out_features=1, bias=True)

4.paramters(): 获取模型的所有参数，返回一个迭代器

for param in model.parameters():
    print(param)

Parameter containing:
tensor([[0.5900]], requires_grad=True)
Parameter containing:
tensor([0.1408], requires_grad=True)

---

参考资料

nn.Module官方文档：https://pytorch.org/docs/stable/generated/torch.nn.Module.html