PyTorch 提供的基本张量存储及运算功能,就可以实现多种神经网络模型,但是这种实现方式不但难度高,而且容易出错。因此,PyTorch将常用的神经网络模型封装到了 torch.nn 包内,以而可以方便灵活地加以调用。
在PyTorch 中构建一个自定义神经网络模型非常简单,就是从torch.nn 中的Module 类派生一个子类,并实现构造函数和 forward 函数。其中,构造函数定义了模型所需的成员对象,如构成该模型的各层,并对其中的参数进行初始化等。而forward 函数用来实现该模块的前向过程,即对输人进行逐层的处理,从而得到最终的输出结果。下面以多层感知器模型为例,介绍如何自定义 一个神经网络模型,
import torch
from torch import nn
from torch.nn import functional as F
class MLP(nn.Module):
def __init__(self, input_dim, hidden_dim, num_class):
super(MLP, self).__init__()
# 线性变换:输入层->隐含层
self.linear1 = nn.Linear(input_dim, hidden_dim)
# 使用ReLU激活函数
self.activate = F.relu
# 线性变换:隐含层->输出层
self.linear2 = nn.Linear(hidden_dim, num_class)
def forward(self, inputs):
hidden = self.linear1(inputs)
activation = self.activate(hidden)
outputs = self.linear2(activation)
probs = F.softmax(outputs, dim=1) # 获得每个输入属于某一类别的概率
return probs
mlp = MLP(input_dim=4, hidden_dim=5, num_class=2)
inputs = torch.rand(3, 4) # 输入形状为(3, 4)的张量,其中3表示有3个输入,4表示每个输入的维度
probs = mlp(inputs) # 自动调用forward函数
print(probs) # 输出3个输入对应输出的概率
输出
tensor([[0.4173, 0.5827],
[0.4212, 0.5788],
[0.4291, 0.5709]], grad_fn=<SoftmaxBackward0>)