【PyTorch】继承 nn.Module 创建简单神经网络

在面向对象编程(OOP)中,继承 是一种允许你创建一个新类的机制,新类可以继承已有类的特性(如方法和属性),并且可以对其进行修改或扩展。

nn.Module 是 PyTorch 所有神经网络模块的基类,几乎所有的神经网络层、模型和操作都要继承自它,这是为了确保模型能够正确地与 PyTorch 的自动求导机制、优化器和其他工具一起工作。

1. 方法作用

nn.Module 类提供了几个核心方法:

__init__(self)init () 方法是初始化模型的地方。在这个方法里,你定义所有网络的层 (如 nn.Linear、nn.Conv2d 等)、操作、权重等。也就是说,你在 init () 中定义了网络的结构。这个方法中,你需要调用 super().init() 来初始化父类的属性。

forward(self, *input):在这个方法里,你定义前向传播的过程,也就是输入数据如何经过不同的层处理得到输出,也就是前向传播过程。在 forward() 方法中,你指定了输入数据如何通过各个层和操作,最终得到输出。前向传播是模型的核心计算部分。

2. 简单例子

2.1 定义模型

假设你要自定义一个简单的全连接神经网络(MLP),你需要继承 nn.Module 来实现。这里是一个典型的结构:

复制代码
import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()  # 继承父类 nn.Module 的初始化方法

        # 定义网络中的层
        self.fc1 = nn.Linear(2, 4)  # 第一层: 输入2维,输出4维
        self.fc2 = nn.Linear(4, 1)  # 第二层: 输入4维,输出1维

    def forward(self, x):
        # 定义前向传播过程
        x = torch.relu(self.fc1(x))  # 第一层通过ReLU激活函数
        x = self.fc2(x)  # 第二层没有激活函数,直接输出
        return x

# 实例化模型
model = SimpleNN()
# 查看模型结构
print(model)

解读一下代码:

SimpleNN 类继承了 nn.Module,这是所有 PyTorch 模型的基础类。继承 nn.Module 使得我们能够使用 PyTorch 自动求导、优化器以及模型参数管理等功能。

__init__(self)

init() 中,我们定义了网络的层次结构(即各层的类型和大小)。比如,self.fc1 = nn.Linear(2, 4) 表示一个全连接层(fully connected layer),它的输入维度是 2,输出维度是 4。

注意:你需要调用 super().init() 来初始化父类 nn.Module,这样才能让模型正确地处理参数和梯度。

forward(self, *input)

在 forward() 中,我们定义了数据的流动过程。数据通过 fc1 层(第一层)并通过 ReLU 激活函数,然后再经过 fc2 层(第二层)输出结果。

需要注意的是,forward() 方法是我们定义数据如何从输入经过各层计算到输出的地方。

model = SimpleNN()会创建一个 SimpleNN 类的实例,这时模型的层和参数都已经初始化好了。

因此,我们通过继承 nn.Module 来创建模型类,使得我们能够利用 PyTorch 提供的很多便利功能,如自动求导、模型参数管理等。

2.2 训练网络

一旦你定义了模型,就可以使用 PyTorch 的优化器(如 SGD、Adam 等)进行训练了。以下是一个训练的简单框架:

复制代码
# 定义损失函数和优化器
criterion = nn.MSELoss()  # 使用均方误差损失函数
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 使用SGD优化器,学习率为0.01

# 模拟输入数据和目标输出数据
inputs = torch.tensor([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]])  # 假设输入是2维数据
targets = torch.tensor([[1.0], [2.0], [3.0]])  # 假设目标输出是1维数据

# 训练过程
for epoch in range(100):  # 训练100个epoch
    optimizer.zero_grad()  # 每次迭代时,先将梯度清零
    
    # 前向传播
    outputs = model(inputs)  # 计算模型的预测值
    
    # 计算损失
    loss = criterion(outputs, targets)  # 计算损失值
    
    # 反向传播和优化
    loss.backward()  # 计算梯度
    optimizer.step()  # 更新模型参数

    if (epoch + 1) % 20 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item():.4f}')
相关推荐
beot学AI2 分钟前
机器学习之逻辑回归
人工智能·机器学习·逻辑回归
西猫雷婶17 分钟前
神经网络|(十九)概率论基础知识-伽马函数·下
人工智能·深度学习·神经网络·机器学习·回归·scikit-learn·概率论
三只小黄人17 分钟前
AI驱动开发:颠覆传统编程新范式
人工智能·驱动开发
沐雪架构师1 小时前
Docling将pdf转markdown以及与AI生态集成
人工智能·pdf
kevin 11 小时前
扫描件、PDF、图片都能比对!让文档差异无所遁形
大数据·人工智能·pdf
算家云1 小时前
腾讯最新开源HunyuanVideo-Foley本地部署教程:端到端TV2A框架,REPA策略+MMDiT架构,重新定义视频音效新SOTA!
人工智能·音视频·算家云·hunyuanvideo·模型部署教程·镜像社区
小王爱学人工智能1 小时前
迁移学习的案例
人工智能·机器学习·迁移学习
源雀数智1 小时前
源雀SCRM开源:企微文件防泄密
java·人工智能·企业微信·流量运营
Honeysea_702 小时前
容器的定义及工作原理
人工智能·深度学习·机器学习·docker·ai·持续部署
fantasy_arch2 小时前
SVT-AV1 svt_aom_motion_estimation_kernel 函数分析
人工智能·算法·av1