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：获取任何数据集并创建一个返回批量数据的迭代器；
相关推荐
xinxiyinhe14 分钟前
GitHub上免费学习工具的精选汇总
学习·github
大桶矿泉水22 分钟前
RK3588使用笔记：设置程序/服务开机自启
笔记·linux程序自启·linux服务自启·嵌入式服务自启·嵌入式程序自启
【云轩】22 分钟前
《孟婆汤的SHA-256加密》
笔记·嵌入式硬件
Y1nhl1 小时前
搜广推校招面经六十二
人工智能·pytorch·python·算法·机器学习·推荐算法·搜索算法
热爱学习地派大星1 小时前
FPGA调试笔记
笔记·fpga开发
傍晚冰川1 小时前
【STM32】WDG看门狗（学习笔记）
c语言·笔记·科技·stm32·单片机·嵌入式硬件·学习
Moonnnn.1 小时前
运算放大器（三）运算放大器的典型应用
笔记·学习·硬件工程
技术小齐2 小时前
网络运维学习笔记（DeepSeek优化版）026 OSPF vlink（Virtual Link，虚链路)配置详解
运维·网络·学习
云梦之上2 小时前
视觉风格提示词：Visual Style Prompting with Swapping Self-Attention（风格迁移）
pytorch·python·计算机视觉·ai作画·prompt
溯源0063 小时前
pytorch中不同的mask方法：masked_fill, masked_select, masked_scatter
人工智能·pytorch·python
热门推荐
01HTML项目一键打包工具：HTML2EXE 最新版02RAG 实践- Ollama+RagFlow 部署本地知识库03DeepSeek各版本说明与优缺点分析04PyTorch量化技术教程：第四章 PyTorch在量化交易中的应用05本地化部署AI知识库：基于Ollama+DeepSeek+AnythingLLM保姆级教程06从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑07RAG生成中的多文档动态融合及去重加权策略探讨08苍穹外卖面试总结09如何在WPS和Word/Excel中直接使用DeepSeek功能10前端使用WPS WebOffice 做在线文档预览与编辑