torch.nn到底是什么？（精简版）

此文为《torch.nn到底是什么？》的总结版。

首先创建基本的神经网络，然后逐步添加torch.nn，torch.optim，Dateset 和 DataLoader的功能，以显示每一部分的具体作用。

1、设置MNIST数据

使用经典的 MNIST 数据集，该数据集由手写数字（0-9）的黑白图像组成。

使用 pathlib 来处理路径（Python3标准库的一部分），用 requests 下载数据。

from pathlib import Path
import requests

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

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

URL = "http://deeplearning.net/data/mnist/"
FILENAME = "mnist.pkl.gz"

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

该数据集的格式为NumPy array，使用 pickle 存储。

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

每个图片大小为28x28，并存储为长度为784（=28x28）的扁平行。

查看其中的一个图片：

from matplotlib import pyplot
import numpy as np

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

输出为：

(50000, 784)

PyTorch使用 tensor 而不是 NumPy array，所以我们需要将其转换。

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
x_train, x_train.shape, y_train.min(), y_train.max()
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

输出：

tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]]) tensor([5, 0, 4,  ..., 8, 4, 8])
torch.Size([50000, 784])
tensor(0) tensor(9)

2、从头构建神经网络（不使用 torch.nn）

首先只使用PyTorch tensor 操作创建一个模型。

#initializing the weights with Xavier initialisation (by multiplying with 1/sqrt(n)).

import math

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

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

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

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

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

bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
yb = y_train[0:bs]

preds = model(xb)  # predictions

print(preds[0], preds.shape)
print(loss_func(preds, yb))
print(accuracy(preds, yb))

输出：

tensor([-1.7022, -3.0342, -2.4138, -2.6452, -2.7764, -2.0892, -2.2945, -2.5480,
        -2.3732, -1.8915], grad_fn=<SelectBackward>) torch.Size([64, 10])

tensor(2.3783, grad_fn=<NegBackward>)
tensor(0.0938)

现在我们可以进行训练。对于每次迭代，将会做以下几件事：

• 选择一批数据（mini-batch）
• 使用模型进行预测
• 计算损失
• loss.backward() 更新模型的梯度，即权重和偏置

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
       #set_trace()
        start_i = i * bs
        end_i = start_i + bs
        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))

输出：

tensor(0.0806, grad_fn=<NegBackward>) tensor(1.)

3、使用 torch.nn.functional

如果使用了负对数似然损失函数和 log softnax 激活函数，那么Pytorch提供的F.cross_entropy 结合了两者。所以我们甚至可以从我们的模型中移除激活函数。

import torch.nn.functional as F

loss_func = F.cross_entropy

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

注意，在 model 函数中我们不再需要调用 log_softmax。让我们确认一下，损失和精确度与前边计算的一样：

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

输出：

tensor(0.0806, grad_fn=<NllLossBackward>) tensor(1.)

4、使用 nn.Module 重构

继承 nn.Module（它本身是一个类并且能够跟踪状态）建立子类，并实例化模型：

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

输出：

tensor(2.3558, grad_fn=<NllLossBackward>)

将训练循环包装到一个 fit 函数中，以便我们以后运行。

def fit():
    for epoch in range(epochs):
        for i in range((n - 1) // bs + 1):
            start_i = i * bs
            end_i = start_i + bs
            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()

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

输出：

tensor(0.0826, grad_fn=<NllLossBackward>)

5、使用 nn.Linear 重构

使用PyTorch 的 nn.Linear 类建立一个线性层，以替代手动定义和初始化 self.weights 和 self.bias、计算 xb @ self.weights + self.bias 等工作。

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()
print(loss_func(model(xb), yb))

输出：

tensor(2.3156, grad_fn=<NllLossBackward>)

我们仍然能够像之前那样使用 fit 方法

fit()

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

输出：

tensor(0.0809, grad_fn=<NllLossBackward>)

6、使用 optim 重构

定义一个函数来创建模型和优化器，以便将来可以重用它。

from torch import optim

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

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

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

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

输出：

tensor(2.2861, grad_fn=<NllLossBackward>)
tensor(0.0815, grad_fn=<NllLossBackward>)

7、使用 Dataset 重构

from torch.utils.data import TensorDataset

train_ds = TensorDataset(x_train, y_train)
model, opt = get_model()

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        xb, yb = train_ds[i * bs: i * bs + bs]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

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

输出：

tensor(0.0800, grad_fn=<NllLossBackward>)

8、使用 DataLoader 重构

from torch.utils.data import DataLoader

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

for epoch in range(epochs):
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

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

输出：

tensor(0.0821, grad_fn=<NllLossBackward>)

9、增加验证

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)

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

我们将在每个epoch结束时计算和打印验证损失。（注意，我们总是在训练之前调用model.train()，在推理之前调用 model.eval()，因为这些由诸如 nn.BatchNorm2d 和nn.Dropout 等层使用，以确保这些不同阶段的适当行为。）

model, opt = 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()
        opt.step()
        opt.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))

输出：

0 tensor(0.2981)
1 tensor(0.3033)

10、创建 fit() 和 get_data()

loss_batch 函数计算每个批次的损失。

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)

fit 运行必要的操作来训练我们的模型并计算每个epoch的训练和验证损失。

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]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)

get_data 为训练集合验证集返回 DataLoader。

def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

现在，我们获取 DataLoader 和拟合模型的整个过程可以在3行代码中运行：

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
model, opt = get_model()
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

输出：

0 0.3055081913471222
1 0.31777948439121245

11、总结

我们现在有一个通用数据流水线和训练循环，你可以使用它来训练多种类型PyTorch模型。 各部分的功能总结如下：

• torch.nn
  • Module：创建一个可调用的对象，其行为类似于一个函数，但也可以包含状态（例如神经网络层权重）。 它知道它包含哪些参数，并且可以将所有梯度归零，循环遍历它们更新权重等。
  • Parameter：tensor 的包装器（wrapper），它告诉 Module 它具有在反向传播期间需要更新的权重。 只更新具有 requires_grad 属性的 tensor。
  • functional：一个模块（通常按惯例导入到F命名空间中），它包含激活函数，损失函数等，以及非状态（non-stateful）版本的层，如卷积层和线性层。
• torch.optim：包含 SGD 等优化器，可在后向传播步骤中更新 Parameter 的权重。
• Dataset：带有 __len__ 和 __getitem__ 的对象的抽象接口，包括 PyTorch 提供的类，如TensorDataset。
• DataLoader：获取任何 Dataset 并创建一个返回批量数据的迭代器。