Pytorch学习笔记（十一）Learning PyTorch - What is torch.nn really

nenchoumi31192025-03-28 16:35

这篇博客瞄准的是 pytorch 官方教程中 Learning PyTorch 章节的 What is torch.nn really? 部分。主要是教你如何一步一步将最原始的代码进行重构至pytorch标准的代码，如果你已经熟悉了如何使用原始代码以及pytorch标准形式构建模型，可以跳过这一篇。

txt 复制代码 
完整网盘链接: https://pan.baidu.com/s/1L9PVZ-KRDGVER-AJnXOvlQ?pwd=aa2m 提取码: aa2m

What is torch.nn really?

PyTorch 提供了模块和类 torch.nn、torch.optim、Dataset 和 DataLoader，以创建和训练神经网络。为了充分利用它们的强大功能并针对问题对其进行自定义，需要真正了解它们的作用。首先在 MNIST 数据集上训练基本神经网络；最初将仅使用最基本的 PyTorch Tensor功能。然后逐步从 torch.nn、torch.optim、Dataset 或 DataLoader 中添加一个特征，准确展示每个部分的作用以及它如何工作以使代码更简洁或更灵活。

MNIST data setup

这里使用经典的 MNIST 数据集，该数据集由手绘数字（0 到 9 之间）的黑白图像组成。使用 pathlib 来处理路径，并使用请求下载数据集。

准备本地环境

python 复制代码 
from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

下载MNIST数据

python 复制代码 
URL = "https://github.com/pytorch/tutorials/raw/main/_static/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)

加载数据并编码

python 复制代码 
import pickle
import gzip

with gzip.open((PATH/FILENAME).as_posix(), "rb") as f:
    ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')

抽看其中一张图像

python 复制代码 
from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28,28)), cmap='gray')
pyplot.show()
print(x_train.shape)

将数据转换为tensor格式

python 复制代码 
import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)

n, c = x_train.shape
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

Neural net from scratch (without torch.nn)

首先，只使用 PyTorch Tensor操作创建一个模型。PyTorch 提供了创建随机或零填充Tensor的方法，使用这些方法为简单的线性模型创建权重和偏差，告诉 PyTorch 它们需要梯度，PyTorch 会记录对Tensor执行的所有操作，以便它可以在反向传播期间自动计算梯度。对于权重，在初始化后设置了 require_grad，因为我们不希望该步骤包含在梯度中。

定义权重与偏置值

python 复制代码 
import math

weights = torch.randn(784, 10) / math.sqrt(784)
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

由于 PyTorch 能够自动计算梯度，可以使用任何标准 Python 函数作为模型，只需编写一个简单的矩阵乘法和广播加法即可创建一个简单的线性模型。还需要编写一个 log_softmax 激活函数。但尽管 PyTorch 提供了许多预先编写的损失函数、激活函数等，但仍然可以使用普通的 Python 编写自己的函数。PyTorch 甚至会自动为您的函数创建快速加速器或矢量化 CPU 代码。

自定义激活函数

python 复制代码 
def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

执行一次前向传播，前向传播计算得到的preds Tensor不仅包含value，还包含梯度函数。

python 复制代码 
batch_size = 64
xb = x_train[0:batch_size]
preds = model(xb)
print(preds[0], preds.shape)

定义损失函数

python 复制代码 
def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll

计算一次损失

python 复制代码 
yb = y_train[0:batch_size]
print(loss_func(preds, yb))

定义用于计算模型正确率函数

python 复制代码 
def accuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

计算正确率

python 复制代码 
accuracy(preds, yb)

执行训练循环

python 复制代码 
lr = 0.5
epochs = 2

for epoch in range(epochs):
    for i in range((n-1)// batch_size + 1):
        # 抽取数据
        start_i = i * batch_size
        end_i = start_i + batch_size
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        # 执行推理
        pred = model(xb)
        loss = loss_func(pred, yb)
        # 反向传播
        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()
            
    print(loss_func(model(xb), yb), accuracy(model(xb), yb))

Using torch.nn.functional

现在将重构代码使其与之前的功能相同，开始利用 PyTorch 的 nn 类使其更简洁、更灵活。

第一步将手写的激活和损失函数替换为来自 torch.nn. functional的函数来缩短代码，此模块包含 torch.nn 库中的所有函数。除了各种损失和激活函数外，还可以在这里找到一些用于创建神经网络的便捷函数，例如pooling函数。 Pytorch 还提供了一个将负对数似然损失和对数 softmax 激活两者结合起来的单一函数 F.cross_entropy。

python 复制代码 
import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

Refactor using nn.Module

接下来使用 nn.Module 和 nn.Parameter，以实现更清晰、更简洁的训练循环。创建一个类来保存权重、偏差和前向传播函数的方法。nn.Module 有许多将要使用的属性和方法（例如 .parameters() 和 .zero_grad()）。

定义模型

python 复制代码 
from torch import nn

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(784, 10) / math.sqrt(784))
        self.bias = nn.Parameter(torch.zeros(10))
        
    def forward(self, xb):
        return xb @ self.weights + self.bias

model = Mnist_Logistic()
print(loss_func(model(xb), yb))

执行拟合

python 复制代码 
def fit():
    for epoch in range(epochs):
        for i in range((n-1)//batch_size + 1):
            start_i = i * batch_size
            end_i = start_i + batch_size
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            pred = model(xb)
            loss = loss_func(pred, yb)
            
            loss.backward()
            with torch.no_grad():
                for p in model.parameters():
                    p -= p.grad * lr
                model.zero_grad()

fit()
loss_func(model(xb), yb)

Refactor using nn.Linear

继续重构代码，不再手动定义和初始化 self.weights 和 self.bias，也不再计算 xb @ self.weights + self.bias，而是使用 Pytorch 类 nn.Linear 作为线性层。Pytorch 有许多类型的预定义层，可以大大简化代码，而且通常还可以加快速度。

定义模型

python 复制代码 
class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)
        
    def forward(self, xb):
        return self.lin(xb)

执行推理与拟合

python 复制代码 
model = Mnist_Logistic()
loss_func(model(xb), yb)

fit()
loss_func(model(xb), yb)

Refactor using torch.optim

Pytorch 还有一个包含各种优化算法的包，torch.optim。使用优化器中的 step 方法来实现自动参数更新。

定义模型与优化器

python 复制代码 
from torch import optim

def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, optimizer = get_model()
print(loss_func(model(xb), yb))

执行拟合

python 复制代码 
for epoch in range(epochs):
    for i in range((n-1) // batch_size + 1):
        start_i = i * batch_size
        end_i = start_i + batch_size
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
print(loss_func(model(xb), yb))

Refactor using Dataset

PyTorch 有一个抽象的 Dataset 类。Dataset 可以是任何具有 __len__ 和 __getitem__ 函数作为索引方式的对象。这部分介绍如何创建自定义 FacialLandmarkDataset 类作为 Dataset 的子类。

PyTorch 的 TensorDataset 是一个包装Tensor的 Dataset。通过定义长度和索引方式，提供了一种沿Tensor的第一维进行迭代、索引和切片的方法。这在训练时更容易在同一行中访问独立变量和因变量。

用Dataset包装数据

python 复制代码 
from torch.utils.data import TensorDataset

train_ds = TensorDataset(x_train, y_train)

执行拟合

python 复制代码 
model, optimizer = get_model()

for epoch in range(epochs):
    for i in range((n-1) // batch_size + 1):
        xb, yb = train_ds[i*batch_size: i*batch_size+batch_size]
        pred = model(xb)
        loss = loss_func(pred, yb)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
print(loss_func(model(xb), yb))

Refactor using DataLoader

PyTorch 的 DataLoader 负责管理batch，可以从任何数据集创建 DataLoader。DataLoader 使迭代变得更容易，无需使用 train_ds[i*bs : i*bs+bs]，DataLoader 会自动提供每个小bacth。

定义Loader

python 复制代码 
from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=batch_size)

执行拟合

python 复制代码 
model, optimizer = get_model()

for epoch in range(epochs):
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
print(loss_func(model(xb), yb))

Add validation

在实际的训练过程中终应该有一个验证集，以便确定是否过度拟合。打乱训练数据对于防止batch之间的相关性和过度拟合非常重要。另一方面，无论是否打乱验证集，验证损失都将相同。由于打乱需要额外的时间，因此打乱验证数据是没有意义的。

将使用比训练集大两倍的验证集batch size，因为验证集不需要反向传播，因此占用的内存更少（它不需要存储梯度）。

准备训练集、验证集的loader

python 复制代码 
train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=batch_size)

valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=batch_size*2)

执行拟合

python 复制代码 
model, optimizer = get_model()

for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
    model.eval()
    with torch.no_grad():
        valid_loss = sum(loss_func(model(xb), yb) for xb,yb in valid_dl)
    print(epoch, valid_loss / len(valid_dl))

Create fit() and get_data()

继续重构为训练集传递一个优化器，并使用它来执行反向传播。

定义一个batch的loss计算函数

python 复制代码 
def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)
    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()
    return loss.item(), len(xb)

定义拟合函数

python 复制代码 
import numpy as np

def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        # 模型训练模式
        model.train()
        for xb,yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)
        # 模型验证模式
        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb, yb) for xb,yb in valid_dl]
            )
            valid_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)
            print(epoch, valid_loss)
            
def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

执行拟合

python 复制代码 
train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)
model, optimizer = get_model()
fit(epochs, model, loss_func, optimizer, train_dl, valid_dl)

Switch to CNN

现在用三个卷积层构建神经网络，使用 PyTorch 的预定义 Conv2d 类作为卷积层。定义一个具有 3 个卷积层的 CNN，每个卷积后跟一个 ReLU。最后，执行平均池化。

定义模型

python 复制代码 
class Mnist_CNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)
        
    def forward(self, xb):
        xb = xb.view(-1, 1, 28, 28)
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4)
        return xb.view(-1, xb.size(1))

执行拟合

python 复制代码 
model = Mnist_CNN()
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

Using nn.Sequential

torch.nn 还有另一个方便的类，可以使用它来简化代码：Sequential。Sequential 对象以顺序方式运行其中包含的每个模块。使用Lambda 将创建一个view层，然后用 Sequential 定义网络时使用它。

定义view层

python 复制代码 
class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func
        
    def forward(self, x):
        return self.func(x)
 
def preprocess(x):
    return x.view(-1, 1, 28, 28)

定义模型 & 优化器

python 复制代码 
model = torch.nn.Sequential(
    Lambda(preprocess),
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1), 
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AvgPool2d(4),
    Lambda(lambda x: x.view(x.size(0), -1))
)
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

执行拟合

python 复制代码 
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

Wrapping DataLoader

上面的 CNN 相当简洁，但它只适用于 MNIST，因为：

  • 假设输入是一个 28*28 长的向量；
  • 假设最终的 CNN 网格大小为 4*4（因为这是我们使用的平均池化内核大小）

这里要做的是让模型适用于任何 2d 单通道图像。通过将数据预处理移到生成器中来删除初始 Lambda 层：

python 复制代码 
def preprocess(x,y):
    return x.view(-1, 1, 28,28), y

class WrappedDataLoader:
    def __init__(self, dl, func) -> None:
        self.dl = dl
        self.func = func
    
    def __len__(self):
        return len(self.dl)

    def __iter__(self):
        for b in self.dl:
            yield (self.func(*b))

对dataloader进行修改

python 复制代码 
train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)

train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

定义模型

python 复制代码 
model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

执行拟合

python 复制代码 
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

Using your Accelerator

检查当前设备是否支持计算加速

python 复制代码 
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu'
print(f"device {device}")

在预处理阶段将数据移动到加速设备上

python 复制代码 
def preprocess(x,y):
    return x.view(-1,1,28,28).to(device), y.to(device)

train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

将模型移动到加速设备上

python 复制代码 
model.to(device)
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

执行拟合

python 复制代码 
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

Closing thoughts

总结一下这篇教程中的内容：

  • torch.nn:
    • Module：创建一个被调用函数，可以包含状态（例如神经网络层权重）。并且可以将其所有梯度归零，循环遍历它们以更新权重等；
    • Parameter：Tensor的包装器，它告诉模块它具有在反向传播期间需要更新的权重，只有设置了 require_grad 属性的Tensor才会更新；
    • functional：包含激活函数、损失函数等组件的子模块，以及卷积层和线性层等非状态版本的layer；
  • torch.optim：包含优化器，在反向传播期间更新参数的权重；
  • Dataset：具有 __len__ 和 __getitem__ 的对象的抽象接口，包括 Pytorch 提供的类，例如 TensorDataset；
  • DataLoader：获取任何数据集并创建一个返回批量数据的迭代器；
相关推荐
MarkHD30 分钟前
智能体在车联网中的应用：第51天 模仿学习与离线强化学习：破解数据效率与安全困局的双刃剑
学习·安全
Drawing stars3 小时前
JAVA后端 前端 大模型应用 学习路线
java·前端·学习
崇山峻岭之间3 小时前
Matlab学习记录33
开发语言·学习·matlab
玄〤4 小时前
黑马点评中 VoucherOrderServiceImpl 实现类中的一人一单实现解析（单机部署）
java·数据库·redis·笔记·后端·mybatis·springboot
科技林总4 小时前
【系统分析师】3.5 多处理机系统
学习
知乎的哥廷根数学学派5 小时前
基于多模态特征融合和可解释性深度学习的工业压缩机异常分类与预测性维护智能诊断（Python）
网络·人工智能·pytorch·python·深度学习·机器学习·分类
芯思路6 小时前
STM32开发学习笔记之三【按键】
笔记·stm32·学习
Lips6116 小时前
2026.1.11力扣刷题笔记
笔记·算法·leetcode
charlie1145141916 小时前
从 0 开始的机器学习——NumPy 线性代数部分
开发语言·人工智能·学习·线性代数·算法·机器学习·numpy
咚咚王者6 小时前
人工智能之核心基础 机器学习 第十二章 半监督学习
人工智能·学习·机器学习
热门推荐
01GitHub 镜像站点02网站改了域名，如何查找？03Linux下V2Ray安装配置指南04Labelme从安装到标注：零基础完整指南05AI 规范驱动开发“三剑客”深度对比：Spec-Kit、Kiro 与 OpenSpec 实战指南062025-04-03 Latex学习1——本地配置Latex + VScode环境07安娜的档案(Anna’s Archive) 镜像网站/国内最新可访问入口（持续更新）08百度网盘偷偷给电脑“降频”？09【踩坑笔记】50系显卡适配的 PyTorch 安装10Overleaf编译超时，超出免费计划编译时限（已解决）