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

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

• 官网链接：https://pytorch.org/tutorials/beginner/nn_tutorial.html

完整网盘链接: 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

• 官网链接: https://pytorch.org/tutorials/beginner/nn_tutorial.html#mnist-data-setup

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

准备本地环境

from pathlib import Path
import requestsDATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"PATH.mkdir(parents=True, exist_ok=True)

下载MNIST数据

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)

加载数据并编码

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

抽看其中一张图像

from matplotlib import pyplot
import numpy as nppyplot.imshow(x_train[0].reshape((28,28)), cmap='gray')
pyplot.show()
print(x_train.shape)

将数据转换为tensor格式

import torchx_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，因为我们不希望该步骤包含在梯度中。

定义权重与偏置值

import mathweights = 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 代码。

自定义激活函数

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，还包含梯度函数。

batch_size = 64
xb = x_train[0:batch_size]
preds = model(xb)
print(preds[0], preds.shape)

定义损失函数

def nll(input, target):return -input[range(target.shape[0]), target].mean()loss_func = nll

计算一次损失

yb = y_train[0:batch_size]
print(loss_func(preds, yb))

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

def accuracy(out, yb):preds = torch.argmax(out, dim=1)return (preds == yb).float().mean()

计算正确率

accuracy(preds, yb)

执行训练循环

lr = 0.5
epochs = 2for epoch in range(epochs):for i in range((n-1)// batch_size + 1):# 抽取数据start_i = i * batch_sizeend_i = start_i + batch_sizexb = 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 * lrbias -= bias.grad * lrweights.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。

import torch.nn.functional as Floss_func = F.cross_entropydef model(xb):return xb @ weights + biasprint(loss_func(model(xb), yb), accuracy(model(xb), yb))

---

Refactor using nn.Module

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

定义模型

from torch import nnclass 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.biasmodel = Mnist_Logistic()
print(loss_func(model(xb), yb))

执行拟合

def fit():for epoch in range(epochs):for i in range((n-1)//batch_size + 1):start_i = i * batch_sizeend_i = start_i + batch_sizexb = 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 * lrmodel.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 有许多类型的预定义层，可以大大简化代码，而且通常还可以加快速度。

定义模型

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

执行推理与拟合

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

---

Refactor using torch.optim

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

定义模型与优化器

from torch import optimdef get_model():model = Mnist_Logistic()return model, optim.SGD(model.parameters(), lr=lr)model, optimizer = get_model()
print(loss_func(model(xb), yb))

执行拟合

for epoch in range(epochs):for i in range((n-1) // batch_size + 1):start_i = i * batch_sizeend_i = start_i + batch_sizexb = 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包装数据

from torch.utils.data import TensorDatasettrain_ds = TensorDataset(x_train, y_train)

执行拟合

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

from torch.utils.data import DataLoadertrain_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=batch_size)

执行拟合

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

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)

执行拟合

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计算函数

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)

定义拟合函数

import numpy as npdef 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),)

执行拟合

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。最后，执行平均池化。

定义模型

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))

执行拟合

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层

class Lambda(nn.Module):def __init__(self, func):super().__init__()self.func = funcdef forward(self, x):return self.func(x)def preprocess(x):return x.view(-1, 1, 28, 28)

定义模型 & 优化器

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)

执行拟合

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

---

Wrapping DataLoader

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

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

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

def preprocess(x,y):return x.view(-1, 1, 28,28), yclass WrappedDataLoader:def __init__(self, dl, func) -> None:self.dl = dlself.func = funcdef __len__(self):return len(self.dl)def __iter__(self):for b in self.dl:yield (self.func(*b))

对dataloader进行修改

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

定义模型

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)

执行拟合

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

---

Using your Accelerator

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

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu'
print(f"device {device}")

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

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)

将模型移动到加速设备上

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

执行拟合

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：获取任何数据集并创建一个返回批量数据的迭代器；