nn.Module的基本用法
nn.Module是所有神经网络的基类,所以你的神经网络类也应该要继承这个基类
当使用时,主要需要实现其两个方法:
__init__
:初始化方法,可以用于定义神经网络的结构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=<AddBackward0>)
nn.Module的其他常用方法
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)
)
eval()
:如果模型训练完毕,需要进行评估,需要显式的告知模型。调用该方法即可。
model.eval()
model.train(False) # 效果一样
Model(
(model): Linear(in_features=1, out_features=1, bias=True)
)
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