Skip to main content

Pytorch nn AND nn.Functional

🚀 PyTorch API 指南：从核心构建到高级操作

你好，PyTorch 探索者！无论你是刚入门的新手还是寻求进阶的开发者，掌握 PyTorch 丰富的 API 都是提升你模型构建能力的关键。

本指南将一些强大而常用的 PyTorch API 分为五大类，从模型的基础骨架到高效的数据操作，为你提供一份清晰、易懂的参考。让我们开始吧！

第一部分：神经网络的核心构建块 🏗️

这些是构建任何神经网络模型的基石。

1. nn.Module：万物之基

nn.Module 是 PyTorch 中所有神经网络模块的父类。当你构建自定义模型或层时，都应该继承它。

核心功能:

  • 自动参数跟踪: 任何被定义为 nn.Parameter 的属性都会被自动添加到模型的参数列表中。
  • 子模块管理: 它可以像容器一样持有其他 nn.Module，形成一个嵌套的树状结构。
  • 便捷功能: 提供了如 .to(device)（移动模型到 CPU/GPU）、.train()（设为训练模式）、.eval()（设为评估模式）等实用方法。

用法示例:

import torch
import torch.nn as nn

# 创建一个简单的自定义模型
class MySimpleNet(nn.Module):
    def __init__(self, input_size, output_size):
        super(MySimpleNet, self).__init__()
        # 将 nn.Linear 层作为子模块
        self.linear = nn.Linear(input_size, output_size)

    def forward(self, x):
        return self.linear(x)

# 实例化模型
model = MySimpleNet(input_size=10, output_size=2)
print(model)
# MySimpleNet(
#   (linear): Linear(in_features=10, out_features=2, bias=True)
# )

2. nn.Parameter：模型的可学习参数

nn.Parameter 是 torch.Tensor 的一个特殊包装类。当它被赋值为 nn.Module 的属性时，它就自动被注册为模型的可学习参数，这意味着优化器在 optimizer.step() 时会自动更新它。

核心功能:

  • 向 nn.Module 注册一个需要计算梯度的张量。
  • 可以通过 model.parameters() 访问到。

用法示例:

class MyCustomLayer(nn.Module):
    def __init__(self, input_features, output_features):
        super(MyCustomLayer, self).__init__()
        # 手动定义一个可学习的权重矩阵
        self.weight = nn.Parameter(torch.randn(output_features, input_features))
        # 也可以定义一个偏置
        self.bias = nn.Parameter(torch.zeros(output_features))

    def forward(self, x):
        # 使用自定义参数进行操作
        return F.linear(x, self.weight, self.bias)

layer = MyCustomLayer(5, 3)
# 打印所有可学习参数
for name, param in layer.named_parameters():
    print(f"Name: {name}, Shape: {param.shape}, Requires Grad: {param.requires_grad}")

第二部分：张量初始化：播下成功的种子 🌱

正确的初始化是模型成功训练的第一步。

1. torch.empty() & torch.zeros()：创建张量画布

  • torch.empty(*size): 创建一个指定形状的张量，但不进行任何初始化。它的值是分配内存时存在的任意值，因此速度很快，但使用前必须填充数据。
  • torch.zeros(*size): 创建一个指定形状且所有元素都为 0 的张量。

用法示例:

# 创建一个未初始化的 2x3 张量
uninitialized_tensor = torch.empty(2, 3)
print("Empty Tensor:\n", uninitialized_tensor)

# 创建一个全为 0 的 2x3 张量
zero_tensor = torch.zeros(2, 3)
print("Zero Tensor:\n", zero_tensor)

2. torch.nn.init.kaiming_uniform_：智能的权重初始化

这是一种专门为使用 ReLU 及其变体激活函数的网络设计的权重初始化方法（也称 He 初始化）。它可以帮助缓解梯度消失或爆炸问题。函数名末尾的下划线 _ 表示它是一个 in-place 操作，会直接修改输入的张量。

用法示例:

import torch.nn.init as init
import math

# 首先创建一个权重张量
weight = torch.empty(3, 5)

# 使用 Kaiming 均匀分布进行 in-place 初始化
init.kaiming_uniform_(weight, a=math.sqrt(5))
print("Kaiming Initialized Weight:\n", weight)

第三部分：常用神经网络函数 🧠

这些函数位于 torch.nn.functional（通常别名为 F）中，是构建网络计算图的核心部分。

1. F.linear(input, weight, bias=None)：无状态的线性变换

与 nn.Linear 模块不同，F.linear 是一个纯函数。它不持有任何内部状态（如权重或偏置），你需要手动将这些参数传入。这在权重是动态生成或需要共享时非常有用。

用法示例:

import torch.nn.functional as F

input_tensor = torch.randn(4, 5)  # (batch_size, in_features)
weight_matrix = torch.randn(3, 5) # (out_features, in_features)
bias_vector = torch.randn(3)      # (out_features)

# 应用线性变换
output = F.linear(input_tensor, weight_matrix, bias_vector)
print("Output shape:", output.shape) # torch.Size([4, 3])

2. torch.Tensor.softmax(dim)：将输出转换为概率

Softmax 函数能将一组数值转换为一个概率分布，其中所有值的和为 1。dim 参数至关重要，它指定了在哪一个维度上进行 Softmax 操作。

用法示例:

# 假设这是模型对3个类别的原始输出（logits）
logits = torch.tensor([[1.0, 2.0, 0.5], [3.0, 1.0, 1.5]])

# 沿着最后一个维度（类别维度）计算概率
probabilities = logits.softmax(dim=-1)
print("Probabilities:\n", probabilities)
print("Sum of probabilities per sample:", probabilities.sum(dim=-1))

3. F.one_hot(tensor, num_classes=-1)：分类任务的好帮手

这个函数可以将一个包含类别索引的整数张量转换为 one-hot 编码的张量。这在计算交叉熵损失或需要二进制类别表示时非常有用。

用法示例:

# 类别标签 (batch_size=4)
labels = torch.tensor([0, 2, 1, 3])

# 转换为 one-hot 编码，假设总共有 5 个类别
one_hot_labels = F.one_hot(labels, num_classes=5)
print("One-Hot Labels:\n", one_hot_labels)
# tensor([[1, 0, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [0, 1, 0, 0, 0],
#         [0, 0, 0, 1, 0]])

第四部分：张量变形与数学运算 📐

高效的数据操作是 PyTorch 的一大优势。

1. torch.Tensor.view(*shape)：高效的形状重塑

.view() 可以改变张量的形状，只要新旧形状的元素总数相同即可。它非常高效，因为它返回的是一个共享原始数据内存的“视图”，而不是一个数据副本。-1 可以用作占位符，自动推断该维度的大小。

用法示例:

x = torch.randn(4, 6) # 原始形状

# 重塑为 8x3
y = x.view(8, 3)

# 使用 -1 自动推断
z = x.view(2, 2, -1) # z 的形状将是 (2, 2, 6)

print("Original shape:", x.shape)
print("View (8, 3) shape:", y.shape)
print("View (2, 2, -1) shape:", z.shape)

2. torch.topk(input, k, dim=None)：寻找最重要的K个

此函数用于查找张量在指定维度上最大（或最小）的 k 个元素及其对应的索引。

用法示例:

x = torch.tensor([[1, 5, 2], [9, 3, 6]])

# 寻找每行中最大的 2 个元素
values, indices = torch.topk(x, k=2, dim=1)

print("Top-2 Values:\n", values)     # tensor([[5, 2], [9, 6]])
print("Top-2 Indices:\n", indices)   # tensor([[1, 2], [0, 2]])

3. sum(), mean(), div_()：基础数学聚合

这些是张量上最常用的数学运算：

  • .sum(): 求和。
  • .mean(): 求平均值。
  • .div_(): In-place 除法。带下划线的方法会直接修改张量本身。

用法示例:

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

# 沿行求和 (dim=1)，并保持维度
row_sum = x.sum(dim=1, keepdim=True)
print("Row Sum (keepdim=True):\n", row_sum)

# 计算所有元素的平均值
total_mean = x.mean()
print("Total Mean:", total_mean)

# 对 x 进行 in-place 除以 2
x.div_(2)
print("x after in-place division:\n", x)

第五部分：高级索引：scatter_add_ 的魔力 ✨

这是一个功能强大但相对复杂的 in-place 操作，用于实现高级的聚合功能。

torch.Tensor.scatter_add_(dim, index, src)

它的工作方式可以理解为“根据 index 将 src 张量中的值加到 self 张量中”。

  • self: 目标张量，它将被修改。
  • dim: 进行操作的维度。
  • index: 索引张量，指示 src 中的每个值应该被添加到 self 的哪个位置。
  • src: 源张量，其值将被“散布”出去。

一个直观的例子: 假设你想统计一个班级中，来自不同城市（用 ID 表示）的学生人数。

# 假设有 3 个城市 (ID: 0, 1, 2)，初始化计数器
city_counts = torch.zeros(3, dtype=torch.long)

# 8个学生，他们各自的城市ID
# 学生0来自城市1，学生1来自城市0，学生2来自城市1...
student_city_ids = torch.tensor([1, 0, 1, 2, 0, 0, 1, 2])

# 我们要为每个学生加 1
ones_to_add = torch.ones(8, dtype=torch.long)

# 使用 scatter_add_ 进行聚合
# 沿着维度 0，根据 student_city_ids，将 ones_to_add 的值加到 city_counts
city_counts.scatter_add_(dim=0, index=student_city_ids, src=ones_to_add)

print(f"城市 0 的学生人数: {city_counts[0]}") # 3
print(f"城市 1 的学生人数: {city_counts[1]}") # 3
print(f"城市 2 的学生人数: {city_counts[2]}") # 2

希望这份指南能帮助你更自如地运用 PyTorch。祝你编码愉快！

Back to top