【Scala PyTorch深度学习】PyTorch On Scala系列课程 第十四章 29 PyTorch模型扩展自定义Module【AI Infra3】[PyTorch Scala硕士研一课程】

PyTorch Scala高校计算机硕士研一课程

使用自定义模块扩展 torch.nn

虽然 PyTorch 提供了像 torch.nn.Linear 或 torch.nn.Conv2d 这样的基本构建块，以及像 torch.nn.Sequential 这样的容器，但应用程序通常会通过将更复杂或专门的逻辑封装到可重用组件中来获得好处。扩展 torch.nn.Module 是 PyTorch 用于创建这些自定义层或网络部分的标准机制。这种方法提升了模块化、代码组织性和可重用性，使得管理复杂的模型架构变得更加容易。它不仅允许您定义层，还可以定义具体的正向计算逻辑，包括控制流、子组件之间的配合，以及与自定义操作的整合。

自定义模块的构成

其核心是，自定义模块是一个继承自 torch.nn.Module 的 Python 类。您通常会重写的两个最重要的方法是：

这是一个基本结构：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MyCustomModule extends nn.Module:
def __init__(input_features : Int, output_features : Int, hidden_units : Int):
super().__init__() # 必不可少的第一步

// 定义子模块（层）
val layer1 = nn.Linear(input_features, hidden_units)
val activation = nn.ReLU()
val layer2 = nn.Linear(hidden_units, output_features)

// 直接定义可学习参数（如果需要）
// 示例：一个可学习的缩放因子
val scale = nn.Parameter(torch.randn(1))

// 定义不可学习的状态（缓冲区）
// 示例：一个用于正向传播的计数器（仅作演示）
register_buffer('forward_count', torch.zeros(1, dtype=torch.long))

def forward(x : Tensor):
// 使用已初始化的组件定义计算流程
val x = layer1(x)
val x = activation(x)
val x = layer2(x)

// 使用自定义参数
val x = x * scale

// 更新缓冲区（如果需要，请确保设备兼容性）
// 注意：这样的直接修改在标准训练循环中可能不常见，
// 但它展示了缓冲区的用法。
forward_count += 1

return x

// 使用示例：
val input_dim = 64
val output_dim = 10
val hidden_dim = 128

val model = MyCustomModule(input_dim, output_dim, hidden_dim)
println(model)

// 测试正向传播
val dummy_input = torch.randn(4, input_dim) // 批量大小为 4
val output = model(dummy_input)
println("输出形状:", output.shape)
println("正向计数:", model.forward_count)

// 参数和缓冲区均被跟踪
for (name, param) <- model.named_parameters():
println(f"参数: {name}, 形状: {param.shape}")
for (name, buf) <- model.named_buffers():
println(f"缓冲区: {name}, 值: {buf}")

实践建议

创建有效的自定义模块不仅仅是继承自 nn.Module。请考虑以下做法：

• 初始化用于定义： __init__ 主要用于定义组件（子模块、参数、缓冲区）。避免在此处执行大量计算。所有定义为 nn.Module 实例或 nn.Parameter 实例的属性都会自动注册。这意味着它们会出现在 model.parameters() 中，其状态会被保存到 model.state_dict() 中，并且 model.to(device) 等方法可以正确移动它们。
• 使用 register_buffer 管理状态： 对于模块状态的一部分但不应由优化器更新的张量（如运行统计数据或固定常数），请使用 self.register_buffer('缓冲区名称', 张量)。与持有张量的普通 Python 属性不同，缓冲区能被 state_dict 和设备放置方法（.to()、.cuda()、.cpu()）正确处理。
• forward 定义计算： forward 方法封装了模块的运行时逻辑。它可以包含任何有效的 Python 代码，包括条件语句（if/else）和循环（for），从而实现动态计算行为。确保梯度计算所需的张量是在 forward 中创建或操作的，或作为参数传递的。
• 模块化和专注： 设计模块以执行特定、明确定义的功能。更小、更专注的模块更易于测试、调试和重用。复杂的网络可以通过组合这些经过良好测试的单元来构建。
• 清晰的输入/输出： 在模块的文档字符串中记录输入和输出张量的预期形状和类型。这提高了可用性和集成性。
• 可组合性： 考虑您的模块如何与其他模块结合。nn.Sequential、nn.ModuleList 和 nn.ModuleDict 等标准 PyTorch 容器都可以与自定义模块配合使用。

MyCustomBlock 内部输入 (B, Seq, Din)我的自定义模块(nn.Module)nn.Linear(Dhidden, Dout)输出 (B, Seq, Dout)nn.Linear(Din, Dhidden)自定义功能（例如，注意力机制，自定义逻辑）nn.LayerNorm(Dhidden)

一个使用自定义 nn.Module (MyCustomBlock) 与标准 PyTorch 层组合网络的视图。自定义块封装了内部层和逻辑。

示例：一个简单的缩放点积注意力模块

让我们实现一个基本的缩放点积注意力机制，它是 Transformer 中的一个基本组成部分，作为一个自定义模块。这演示了如何定义参数（隐式地在 nn.Linear 内部）以及在 forward 中实现特定的数学运算。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SimpleScaledDotProductAttention extends nn.Module:
""" 计算简单的缩放点积注意力。 """
def __init__( d_model: Int, d_k: Int, dropout_p: Float = 0.1):
"""
参数:
d_model (int): 输入嵌入的维度。
d_k (int): 键和查询的维度（通常为 d_model // num_heads）。
dropout_p (float): Dropout 概率。
"""
super().__init__()
self.d_k = d_k
// 用于将输入投影到 Q、K、V 空间的线性层
val query_proj = nn.Linear(d_model, d_k)
val key_proj = nn.Linear(d_model, d_k)
val value_proj = nn.Linear(d_model, d_k) // 通常 d_v = d_k
val dropout = nn.Dropout(dropout_p)

def forward( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
"""
参数:
query (torch.Tensor): 查询张量，形状为 (Batch, Seq_len_q, d_model)。
(torch.Tensor): 张量，形状为 (Batch, Seq_len_k, d_model)。
value (torch.Tensor): 值张量，形状为 (Batch, Seq_len_v, d_model)。
通常 Seq_len_k == Seq_len_v。
mask (torch.Tensor, optional): 用于阻止注意力机制关注
某些位置（例如，填充）的掩码张量。
形状为 (Batch, Seq_len_q, Seq_len_k)。
对于被关注的位置，值应为 0；对于被掩码的位置，值应为 -inf。
返回:
torch.Tensor: 注意力机制后的输出张量，形状为 (Batch, Seq_len_q, d_k)。
torch.Tensor: 注意力权重，形状为 (Batch, Seq_len_q, Seq_len_k)。
"""
// 1. 投影输入
val Q = query_proj(query) // (B, Seq_q, d_k)
val K = key_proj(key) // (B, Seq_k, d_k)
val V = value_proj(value) // (B, Seq_v, d_k)

// 2. 计算注意力分数 (QK^T / sqrt(d_k))
// K.transpose(-2, -1) 结果形状为 (B, d_k, Seq_k)
val scores = torch.matmul(Q, K.transpose(–2, –1)) / math.sqrt(self.d_k)
// 分数形状：(B, Seq_q, Seq_k)

// 3. 应用掩码（如果提供）
if mask is not None then
// 确保掩码具有兼容的维度，可能需要unsqueeze
// 示例：如果掩码是 (B, Seq_k)，为 Seq_q 广播添加维度
// mask = mask.unsqueeze(1) // -> (B, 1, Seq_k)
scores = scores.masked_fill(mask == 0, float('-inf')) # 常见约定：0 表示掩码

// 4. 应用 softmax 以获得注意力权重
val attn_weights = F.softmax(scores, dim=–1) // (B, Seq_q, Seq_k)

// 5. 对注意力权重应用 dropout
val attn_weights_dropped = dropout(attn_weights) // (B, Seq_q, Seq_k)

// 6. 计算值的加权和
val output = torch.matmul(attn_weights_dropped, V) // (B, Seq_q, Seq_k) @ (B, Seq_v, d_k) -> (B, Seq_q, d_k)
// 假设 Seq_k == Seq_v

return output, attn_weights

// 使用示例：
val batch_size = 4
val seq_len = 10
val embed_dim = 128
val key_dim = 64

val attention_module = SimpleScaledDotProductAttention(d_model=embed_dim, d_k=key_dim)

// 创建虚拟输入（自注意力机制通常使用相同的张量）
val q_input = torch.randn(batch_size, seq_len, embed_dim)
val k_input = torch.randn(batch_size, seq_len, embed_dim)
val v_input = torch.randn(batch_size, seq_len, embed_dim)

val (output, weights) = attention_module(q_input, k_input, v_input)

println("注意力输出形状:", output.shape) // 预期：(4, 10, 64)
println("注意力权重形状:", weights.shape) // 预期：(4, 10, 10)

此示例将注意力逻辑封装在一个单一模块中，使其易于集成到像 Transformer 编码器或解码器层这样更大的模型中。

进阶考量

• 可变输入大小： 设计模块以处理不同长度的序列，通常使用填充和掩码，如注意力示例中所示。像 torch.nn.utils.rnn.pack_padded_sequence 或自适应池化层这样的技术也可能根据应用而相关。
• 钩子： nn.Module 提供了一个钩子机制（register_forward_hook、register_backward_hook、register_forward_pre_hook），允许您在 forward 传递之前或之后，或在 backward 传递期间执行自定义代码，而无需修改模块的核心 forward 代码。钩子对于调试、可视化或实现某些归一化技术很有用。
• 与自定义操作的接口： 自定义 nn.Module 的 forward 方法是调用专门的 C++ 或 CUDA 扩展（本章其他部分介绍）或自定义 autograd.Function 实例的自然位置，当性能或特定梯度计算需要它们时。模块结构整齐地封装了标准 PyTorch 组件与这些自定义后端之间的配合。

通过掌握 torch.nn.Module 的扩展，您获得了实现几乎任何网络架构或组件的灵活性，能够以清晰、可重用和可维护的方式组织您的代码，这对于应对高级深度学习项目不可或缺。

扩展 torch.optim，使用自定义优化器

尽管 PyTorch 在 torch.optim 中提供了从 SGD 到 Adam 的多种成熟优化算法，但研究和实际用途常能从自定义优化策略中获益。您可能需要实现一篇近期论文中的新算法，调整现有优化器以适应特定限制（例如，仅凭参数组无法处理的分层自适应学习率），或结合不同优化器的步骤。提供了关于如何通过继承 torch.optim.Optimizer 来创建自定义优化器的详细指导，这将使您能够全面掌控参数更新过程。

torch.optim.Optimizer 基类

核心来说，任何 PyTorch 优化器都继承自 torch.optim.Optimizer 基类。理解其结构对于构建您自己的优化器很必要。主要组成部分包括：

实现自定义优化器：带有动量的 SGD 示例

让我们从头开始实现带有动量的随机梯度下降（SGD），以说明整个过程。带有动量的 SGD 更新规则如下：

vt+1=μvt+gt+1v**t+1=μvt+g**t+1pt+1=pt−αvt+1p**t+1=p**t−αvt+1

说明：

• ptp**t 是时间步 tt 的参数。
• gt+1g**t+1 是时间步 t+1t+1 损失对 pp 的梯度。
• vtv**t 是时间步 tt 的动量缓冲区（速度）。
• μμ 是动量系数。
• αα 是学习率。

您可以这样实现它：

import torch
import torch.optim.Optimizer
import defaultdict

class CustomSGD extends Optimizer:
"""实现带有动量的随机梯度下降。"""

def __init__(params: Iterable[Tensor], lr:Float=0.01, momentum=0.0, weight_decay=0.0):
if lr < 0.0:
raise ValueError(f"无效学习率: {lr}")
if momentum < 0.0:
raise ValueError(f"无效动量值: {momentum}")
if weight_decay < 0.0:
raise ValueError(f"无效 weight_decay 值: {weight_decay}")

defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay)
super().__init__(params, defaults)

// 初始化状态（尽管通常在 step 中惰性初始化）
// 如果我们在 step 中初始化，这里并非严格需要
// for group in self.param_groups:
// for p in group['params']:
// self.state[p] = dict(momentum_buffer=None)

@torch.no_grad() // 重要：在优化器步骤内禁用梯度跟踪
def step(self, closure=None):
"""执行一次优化步骤。

参数:
closure (callable, optional): 一个可调用函数，用于重新评估模型
并返回损失。
"""
val loss = None
if closure is not None:
with torch.enable_grad(): # 确保为闭包启用梯度
loss = closure()

for group <- self.param_groups:
val lr = group['lr']
val momentum = group['momentum']
val weight_decay = group['weight_decay']

for p <- group['params']:
if p.grad is None:
continue // 跳过没有梯度的参数

var grad = p.grad // 获取梯度张量`

// 如果指定，应用权重衰减（L2 惩罚）
// 注意：这是标准方法，修改梯度
if weight_decay != 0 then
grad = grad.add(p, alpha=weight_decay)

// 访问并更新参数状态（动量缓冲区）
val param_state = self.state(p)

if 'momentum_buffer' not in param_state then
// 在第一步中惰性初始化动量缓冲区
param_state('momentum_buffer') = torch.clone(grad).detach()
else:
param_state('momentum_buffer').mul_(momentum).add_(grad) // v = mu*v + grad

// 获取更新后的动量缓冲区
val momentum_buffer = param_state('momentum_buffer')

// 执行参数更新步骤
// p = p – lr * momentum_buffer
p.add_(momentum_buffer, alpha=–lr)

return loss

实现要点：

• @torch.no_grad(): 使用 @torch.no_grad() 装饰 step 方法很重要。优化步骤不应成为 autograd 跟踪的计算图的一部分。
• 迭代： 代码会遍历 param_groups，然后遍历每个组内的 params。
• 梯度检查： 它会检查 if p.grad is None:，因为模型中的某些参数可能不会接收到梯度（例如，如果它们未在正向传播中使用或与计算图分离）。
• 状态管理： momentum_buffer 存储在 self.state[p] 中。它被惰性初始化（参数第一次调用 step 时），以避免在参数从未获得梯度时预先分配内存。
• 原地更新： 参数更新使用 add_ 和 mul_ 等原地操作。这会直接修改参数张量，而不会创建新张量，这对于优化器实际更新模型权重非常必要。
• 超参数： 超参数（lr、momentum、weight_decay）从 group 字典中获取，允许不同组拥有不同的设置。
• 权重衰减： 权重衰减（L2 正则化）通常通过在主要更新步骤之前将缩放后的参数值添加到梯度来实现。

使用自定义优化器

使用您的自定义优化器就像使用内置优化器一样：

// 假设 'model' 是您的 torch.nn.Module
// 实例化自定义优化器
val optimizer = new CustomSGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e–4)

// 在您的训练循环中：
for inputs, targets <- dataloader:
optimizer.zero_grad()
val outputs = model(inputs)
val loss = criterion(outputs, targets)
loss.backward()
optimizer.step() # 使用 CustomSGD 逻辑执行更新

进阶考量

• 参数组： 您可以在初始化优化器时定义参数组，以应用不同的设置：

  val optimizer = new CustomSGD([
  {'params': model.base.parameters()},
  {'params': model.classifier.parameters(), 'lr': 1e–3} # 分类器使用不同的学习率
  ], lr=1e–4, momentum=0.9) # 其他参数（例如，基础部分）使用默认学习率

• 闭包： 如果您的算法需要多次损失评估（如 L-BFGS），您将定义一个 closure 函数并将其传递给 optimizer.step(closure)。您的 step 实现必须适当地调用 closure()，可能多次，通常在 with torch.enable_grad(): 块内。

闭包使用示例（具体优化器逻辑不同）

def closure():
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
return loss

在 step 方法中（对于类似 L-BFGS 的优化器）

# loss = closure() # 内部可能被多次调用
# 使用损失和梯度来更新参数…
```

• 与学习率调度器的互动： 正确管理 self.param_groups 的自定义优化器可以与 PyTorch 的学习率调度器 (torch.optim.lr_scheduler) 配合使用。调度器会修改每个 param_group 中的 'lr' 值，然后您的自定义 step 函数会读取该值。
• 复杂状态： 像 Adam 或 AdamW 这样的优化器需要每个参数更多的状态（例如，一阶和二阶矩估计，可能还有步数计数）。您可以通过向 self.state[p] 字典添加更多条目来管理这一点。
• 性能： 尽管 Python 允许快速原型化新的优化器思路，但计算量大的更新规则可能会成为瓶颈。如果分析显示优化器步骤很慢，您可以考虑将核心逻辑实现为自定义 C++ 或 CUDA 扩展以获得最佳性能，正如本章其他部分所讨论的。然而，对于大多数算法，Python 迭代的开销与梯度计算相比可以忽略不计，纯 Python 实现完全足够并且更容易维护。

通过继承 torch.optim.Optimizer，您能够实现几乎任何参数更新规则，将新颖的优化研究直接整合到您的 PyTorch 训练流程中，并根据您的特定需求微调学习过程。

外部函数接口 (FFI)

收藏

虽然构建自定义 C++ 或 CUDA 扩展能提供与 PyTorch 最紧密的结合，尤其是对于需要自动求导支持的操作，但在某些情况下，您需要与现有 C 库对接，而无需将它们重写为完整的 PyTorch 扩展。这时，外部函数接口 (FFI) 就派上用场了。FFI 允许 Python 代码调用用其他语言（最常见的是 C 或 C++）编写并编译成共享库的函数。

Python 的标准库包含 ctypes 模块，它是为此目的设计的一个强大工具。它能够加载共享库（Linux/macOS 上的 .so 文件，Windows 上的 .dll 文件），并直接从 Python 调用其中的函数。这种方法在以下情况中特别有用：

使用 ctypes 进行 C 集成

ctypes 的核心流程包含以下步骤：

与 PyTorch 张量对接

将 C 库与 PyTorch 集成时最常见的需求是传递张量数据。由于 C 函数操作原始内存缓冲区，您需要提供指向张量数据的指针。PyTorch 张量为此提供了 data_ptr() 方法。

import torch
import ctypes

// 假设 'mylib.so' 或 'mylib.dll' 包含一个函数：
// void process_data(float* data_ptr, int size);

// 加载共享库
try:
// 根据需要调整路径/名称
lib = ctypes.CDLL('./mylib.so')
catch OSError as e:
println(f"加载共享库错误: {e}")
// 适当处理错误，例如在 Windows 上尝试 .dll 等。
exit()

// 定义函数签名
try:
val process_data_func = lib.process_data
val process_data_func.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_int]
val process_data_func.restype = None // void 返回类型
catch AttributeError as e:
println(f"查找函数或设置签名错误: {e}")
// 处理错误：库中可能不存在该函数
exit()

// 创建一个 PyTorch 张量
val tensor = torch.randn(100, dtype=torch.float32)

// — 重要部分：确保张量数据布局兼容 —
// 许多 C 函数期望连续的 C 风格数组。
if not tensor.is_contiguous() then
tensor = tensor.contiguous()

// 获取数据指针（作为 void 指针，然后进行类型转换）
val data_ptr_void = tensor.data_ptr()
// 将 void 指针转换为 C 函数期望的特定类型
val data_ptr_c = ctypes.cast(data_ptr_void, ctypes.POINTER(ctypes.c_float))

// 调用 C 函数
val size = tensor.numel()
try:
process_data_func(data_ptr_c, ctypes.c_int(size))
println("成功调用 C 函数。")
// 'tensor' 的数据可能被 C 函数原地修改
// print(tensor)
catch Exception as e:
println(f"C 函数执行错误: {e}")

重要注意事项：

• 内存管理： 当您将 tensor.data_ptr() 传递给 C 时，您是在共享内存。C 代码直接读取或写入由 PyTorch 管理的内存。您必须确保 PyTorch 张量在 C 函数使用其指针的整个过程中保持分配和有效。在传递指针后在 Python 中修改张量的大小或存储可能导致崩溃或数据损坏。
• 数据连续性： C 函数通常期望数据数组在内存中是连续的（就像标准的 C 数组一样）。某些操作（例如切片、转置）产生的 PyTorch 张量可能不连续。在为需要连续数据的 C 函数获取数据指针之前，始终使用 tensor.is_contiguous() 进行检查，并在必要时调用 tensor.contiguous()。
• 类型匹配： 精确匹配 C 函数的参数类型与 ctypes 定义（c_float、c_int、POINTER(…) 等），并确保 PyTorch 张量的 dtype 与使用的指针类型一致（例如，torch.float32 对应 ctypes.POINTER(ctypes.c_float)）。不匹配会导致未定义行为。
• 错误处理： C 函数可能不会引发 Python 异常。您可能需要设计您的 C 函数返回错误代码或状态指示器，并在您的 Python 包装代码中显式检查这些信息。
• 全局解释器锁 (GIL)： 通过 ctypes 调用 C 函数可以释放 Python 的 GIL，这意味着 C 代码可以与其他 Python 线程（如果有的话）并发执行。如果 C 代码是 CPU 密集型且耗时较长，这可以提供并行性。然而，如果 C 代码频繁回调 Python C API，它可能会重新获取 GIL，从而限制并发优势。
• 复杂性： 尽管 ctypes 很强大，但为具有数据结构、回调或错误处理的复杂 C API 创建绑定可能变得有难度。

图示：FFI 数据流

PyTorch 张量(dtype=float32,is_contiguous=True)tensor.data_ptr()(获取内存地址)ctypes 包装器(定义签名,将指针转换为POINTER(c_float))传递地址C 函数(void process_data(float* data, int size))用转换后的指针和大小调用函数共享库(mylib.so / mylib.dll)在其中执行原地修改数据(共享内存)

流程图说明了如何通过 ctypes 获取 PyTorch 张量的内存地址并传递给编译好的 C 共享库中的函数。C 函数直接操作张量的内存。

替代方案：CFFI

Python 中另一个常用的 FFI 库是 CFFI（C 外部函数接口）。CFFI 通常要求您提供 C 函数声明（例如，在 Python 字符串中使用 C 语法），并处理大部分类型转换和接口生成。对于复杂 API 而言，它有时可能更易于使用，并且在某些情况下可能比 ctypes 提供更好的性能。然而，ctypes 是标准库的一部分，无需额外安装。

何时使用 FFI 与自定义扩展

使用 ctypes 或 CFFI 的 FFI 通常最适合集成现有、自包含的 C/C++ 库，且这些 C 函数不需要自动求导支持。如果您需要与 PyTorch 的自动求导引擎紧密结合，需要为您的 C/C++ 代码实现自定义梯度计算，或者专门为 PyTorch 构建性能关键组件，那么编写原生 C++ 或 CUDA 扩展（如章节“构建自定义 C++ 扩展”和“构建自定义 CUDA 扩展”中所述）是更合适、功能更强的方法。当运用外部预编译代码是主要目标时，FFI 可作为实用的桥梁。

实践：构建一个简单的 CUDA 扩展

构建一个简单的 CUDA 扩展涉及创建用于基本操作（标量向量加法）的自定义 CUDA 核函数，编写必要的 C++ 绑定，编译它，并最终使用 PyTorch 张量从 Python 调用它。这个过程展示了在 GPU 上加速特定计算的基本步骤。

目标

我们的目标是实现一个函数 scaled_add(alpha, x, y)，计算 z=α∗x+yz=α∗x+y，其中 αα 是一个标量，x,y,zx,y,z 是向量（一维张量）。我们将把核心计算写成 CUDA 核函数，并将其作为 PyTorch C++ 扩展进行集成。

前提条件

请确保您已安装并配置以下各项：

文件结构

让我们组织代码。创建一个这样的目录结构：

simple_cuda_extension/
├── setup.py
└── src/
├── scaled_add.cpp
└── scaled_add_kernel.cu

步骤 1：编写 CUDA 核函数 (scaled_add_kernel.cu)

此文件包含实际的 GPU 代码。我们定义一个 CUDA 核函数，逐元素执行标量加法。

// src/scaled_add_kernel.cu
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h> // 如果需要，用于 CUDA 数学函数，但本例中并非严格需要

// CUDA 核函数定义
// 为每个元素计算 z = alpha * x + y
__global__ void scaled_add_kernel(const float* x, const float* y, float* z, float alpha, int N) {
// 计算全局线程索引
int index = blockIdx.x * blockDim.x + threadIdx.x;
int stride = gridDim.x * blockDim.x; // 网格中的线程总数

// 使用网格步进循环确保所有元素都被处理
// 即使 N 大于启动的线程数。
for (int i = index; i < N; i += stride) {
z[i] = alpha * x[i] + y[i];
}
}

// C++ 包装函数（可选但推荐）
// 这可以从主要的 C++ 绑定代码中调用。
// 它设置核函数的启动配置。
void scaled_add_cuda_launcher(const float* x, const float* y, float* z, float alpha, int N) {
// 定义块和网格维度
// 通常，选择块大小为 32（warp 大小）的倍数
// 常见选择有 128, 256, 512, 1024
int blockSize = 256;
// 计算覆盖所有 N 个元素所需的网格大小
// 等同于 ceil(N / blockSize)
int gridSize = (N + blockSize – 1) / blockSize;

// 启动核函数
scaled_add_kernel<<<gridSize, blockSize>>>(x, y, z, alpha, N);

// 可选：检查核函数启动错误（对调试有用）
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
fprintf(stderr, "CUDA kernel launch failed: %s\\n", cudaGetErrorString(err));
// 在实际应用中，可以考虑在此处抛出异常
}
// 可选：如果立即需要，同步设备（等待核函数完成）
// cudaDeviceSynchronize(); // 如果后续操作使用相同的流，通常不需要
}

说明：

• __global__ void scaled_add_kernel(…)：定义一个在 GPU 上运行的函数。
• blockIdx.x、blockDim.x、threadIdx.x、gridDim.x：内置的 CUDA 变量，为每个线程在其启动的线程网格中提供唯一的 ID 和上下文。
• index = blockIdx.x * blockDim.x + threadIdx.x：计算每个线程的唯一全局索引。
• stride = gridDim.x * blockDim.x：网格中的线程总数。
• 网格步进循环： for 循环 (for (int i = index; i < N; i += stride)) 非常重要。它允许固定数量的线程（可能少于 N）通过让每个线程处理多个间隔 stride 的元素来处理所有 N 个元素。这比假设 N 可以被块大小完美整除，或者网格大小与 N / blockSize 精确匹配要准确。
• scaled_add_cuda_launcher：一个辅助 C++ 函数，用于配置并启动核函数。它根据输入大小 N 和选择的块大小 (blockSize) 计算所需的块数量 (gridSize)。<<<gridSize, blockSize>>> 是启动核函数的 CUDA 语法。

步骤 2：编写 C++ 绑定 (scaled_add.cpp)

此文件将 CUDA 代码与 PyTorch 连接起来。它定义了一个可从 Python 调用的函数，处理张量数据访问，并调用 CUDA 启动器。

// src/scaled_add.cpp
#include <torch/extension.h>
#include <vector>

// 从 scaled_add_kernel.cu 转发声明 CUDA 启动函数
void scaled_add_cuda_launcher(const float* x, const float* y, float* z, float alpha, int N);

// 将绑定到 Python 的 C++ 接口函数
// 它接受 PyTorch 张量作为输入
torch::Tensor scaled_add(torch::Tensor x, torch::Tensor y, float alpha) {
// 输入验证：确保张量在 GPU 上且具有相同的形状/数据类型
TORCH_CHECK(x.device().is_cuda(), "输入张量 x 必须是 CUDA 张量");
TORCH_CHECK(y.device().is_cuda(), "输入张量 y 必须是 CUDA 张量");
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "输入张量 x 必须是 float32");
TORCH_CHECK(y.scalar_type() == torch::kFloat32, "输入张量 y 必须是 float32");
TORCH_CHECK(x.is_contiguous(), "输入张量 x 必须是连续的");
TORCH_CHECK(y.is_contiguous(), "输入张量 y 必须是连续的");
TORCH_CHECK(x.sizes() == y.sizes(), "输入张量 x 和 y 必须具有相同的形状");
TORCH_CHECK(x.dim() == 1, "输入张量 x 必须是 1D"); // 本例的简单检查
TORCH_CHECK(y.dim() == 1, "输入张量 y 必须是 1D"); // 本例的简单检查

// 获取元素数量
int N = x.numel();

// 创建输出张量（与输入在同一设备上）
auto z = torch::empty_like(x); // 创建具有相同形状、数据类型、设备的张量

// 获取原始数据指针
// .data_ptr<float>() 允许访问底层 C++ float* 数据
const float* x_ptr = x.data_ptr<float>();
const float* y_ptr = y.data_ptr<float>();
float* z_ptr = z.data_ptr<float>();

// 调用 .cu 文件中定义的 CUDA 核函数启动器函数
scaled_add_cuda_launcher(x_ptr, y_ptr, z_ptr, alpha, N);

return z;
}

// 使用 PYBIND11_MODULE 宏的绑定代码
// 这将创建名为 'simple_cuda_extension_cpp' 的 Python 模块
// 第二个参数 'm' 是模块对象
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// 将 C++ 'scaled_add' 函数公开为 Python 中的 'scaled_add'
m.def("scaled_add", &scaled_add, "在 CUDA 上计算的标量向量加法 (alpha * x + y)");
}

说明：

• #include <torch/extension.h>：PyTorch C++ 扩展的主要头文件。
• 前向声明：我们声明 scaled_add_cuda_launcher，以便编译器在使用它之前知道它。实际实现位于 .cu 文件中，并将在稍后链接。
• scaled_add(torch::Tensor x, torch::Tensor y, float alpha)：向 Python 公开的函数。它接受 PyTorch 张量和一个浮点数。
• TORCH_CHECK(…)：PyTorch 的断言宏。它检查条件并在失败时抛出描述性的 C++ 异常（这些异常会被转换为 Python 异常）。我们验证设备、数据类型、连续性、形状和维度。连续性很重要，因为 CUDA 核函数通常假设数据在内存中是连续排列的。
• torch::empty_like(x)：创建与 x 具有相同属性（大小、数据类型、设备）的输出张量 z，但不初始化内存内容。
• .data_ptr<float>()：获取指向张量底层数据缓冲区的原始 C 风格指针。这是传递给 CUDA 核函数所必需的。
• scaled_add_cuda_launcher(…)：调用在我们的 .cu 文件中定义的函数。
• PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)：这个宏（由 torch/extension.h 提供，它包含 pybind11）为 Python 模块创建入口点。TORCH_EXTENSION_NAME 是一个占位符，它将被 setup.py 中指定的模块名称替换。
• m.def("scaled_add", …)：将 C++ 函数 scaled_add 绑定到模块 m 中 Python 名称 scaled_add。该字符串是 Python 函数的文档字符串。

步骤 3：创建构建脚本 (setup.py)

此脚本使用 Python 的 setuptools 和 PyTorch 的工具来将 C++ 和 CUDA 代码编译成一个 Python 扩展模块。

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
name='simple_cuda_extension_cpp', # 包名，可以是任意名称
ext_modules=[
CUDAExtension(
name='simple_cuda_extension_cpp', # 用户将导入的 Python 模块名称
sources=[
'src/scaled_add.cpp',
'src/scaled_add_kernel.cu',
]
)
],
cmdclass={
'build_ext': BuildExtension
}
)

说明：

• from torch.utils.cpp_extension import BuildExtension, CUDAExtension：从 PyTorch 导入必要的构建工具。
• CUDAExtension(…)：指定我们正在构建一个包含 CUDA 代码的扩展。
  • name：生成的 Python 模块的名称（例如，import simple_cuda_extension_cpp）。这必须与 PYBIND11_MODULE 宏内部使用的 TORCH_EXTENSION_NAME 占位符匹配。
  • sources：扩展所需的所有源文件（.cpp 和 .cu）列表。
• cmdclass={'build_ext': BuildExtension}：告知 setuptools 使用 PyTorch 的自定义构建命令，该命令知道如何处理 CUDA 编译（nvcc）和与 PyTorch 库链接。

步骤 4：编译扩展

在您的终端中导航到 simple_cuda_extension 目录（包含 setup.py 的目录），然后运行构建命令：

# 选项 1：构建并安装到您的 Python 环境中
python setup.py install

# 选项 2：就地构建（在当前目录创建 .so 或 .pyd 文件）
# 对开发有用
python setup.py build_ext –inplace

如果成功，此命令将调用 C++ 编译器和 nvcc 来编译您的代码，并将其链接到 PyTorch 库，生成一个共享对象文件（例如，在 Linux 上为 simple_cuda_extension_cpp.cpython-39-x86_64-linux-gnu.so，在 Windows 上为 simple_cuda_extension_cpp.pyd），可供 Python 导入。

步骤 5：在 Python 中使用扩展

现在您可以像使用任何其他 Python 模块一样导入和使用您的自定义 CUDA 函数了。

# test_extension.py（放置在 simple_cuda_extension 目录之外，或安装后放置）
import torch
import time

// 尝试导入已编译的扩展
try:
import simple_cuda_extension_cpp
println("成功导入 CUDA 扩展。")
catch ImportError:
println("导入 CUDA 扩展时出错。您是否成功编译了它？")
println("运行：python setup.py build_ext –inplace（在扩展目录中）")
exit()

// 首先在 CPU 上定义输入张量
val N = 1024 * 1024 // 向量大小
val alpha = 2.5
val x_cpu = torch.randn(N, dtype=torch.float32)
val y_cpu = torch.randn(N, dtype=torch.float32)

// 将张量移动到 GPU
if torch.cuda.is_available() then
val device = torch.device('cuda')
val x_gpu = x_cpu.to(device)
val y_gpu = y_cpu.to(device)
println(f"使用设备: {device}")
else
println("CUDA 不可用。退出。")
exit()

// 确保输入是连续的（对 .data_ptr() 很重要）
x_gpu = x_gpu.contiguous()
y_gpu = y_gpu.contiguous()

// — 使用自定义 CUDA 扩展 —
println("\\n测试自定义 CUDA 扩展：")
// GPU 预热
val _ = simple_cuda_extension_cpp.scaled_add(alpha, x_gpu, y_gpu)
torch.cuda.synchronize() // 等待预热完成

val start_time = time.time()
val z_gpu_custom = simple_cuda_extension_cpp.scaled_add(alpha, x_gpu, y_gpu)
torch.cuda.synchronize() // 等待核函数完成再停止计时器
val end_time = time.time()
println(f"自定义 CUDA 扩展时间：{(end_time – start_time)*1000:.4f} 毫秒")

// — 使用标准 PyTorch 操作进行验证 —
println("\\n测试标准 PyTorch 操作：")
// GPU 预热
val _ = alpha * x_gpu + y_gpu
torch.cuda.synchronize()

val start_time = time.time()
val z_gpu_pytorch = alpha * x_gpu + y_gpu
torch.cuda.synchronize()
val end_time = time.time()
println(f"标准 PyTorch 时间：{(end_time – start_time)*1000:.4f} 毫秒")

// — 验证 —
// 检查结果是否接近（允许浮点差异）
val difference = torch.abs(z_gpu_custom – z_gpu_pytorch).mean()
println(f"\\n自定义和 PyTorch 结果之间的平均绝对差异：{difference.item()}")

if torch.allclose(z_gpu_custom, z_gpu_pytorch, atol=1e–6):
print("结果匹配！")
else:
print("结果不匹配！")

// 示例：如果需要，打印前几个元素
// println("自定义输出（前 10 个）:", z_gpu_custom[:10])
// println("PyTorch 输出（前 10 个）:", z_gpu_pytorch[:10])

运行测试： 保存上述 Python 代码（例如，保存为 test_extension.py），然后运行它：python test_extension.py。

您应该会看到输出，显示导入是否成功、自定义核函数和标准 PyTorch 操作的执行时间，以及确认结果在数值上非常接近的检查。对于现代 GPU 上的这种简单操作，标准 PyTorch 操作是高度优化的，因此如果 PyTorch 版本更快或相似，请不要惊讶。自定义扩展的好处在复杂、非标准的操作或可以融合到单个核函数中的操作序列中会变得更明显。

总结

此实践练习展示了创建 PyTorch CUDA 扩展的端到端过程：

“虽然这个例子很简单，但它建立了基本的工作流程。扩展通常涉及更复杂的核函数，可能处理不同的数据类型、多个维度，如果需要自动求导支持，则需要定义自定义反向传播（请参考第 1 章关于自定义自动求导函数的内容）。构建扩展需要仔细注意内存管理、数据类型、设备放置和同步，但它提供了一种有效的方法来优化 PyTorch 模型中性能关键的部分。”