终面倒计时10分钟：候选人用`PyTorch`实现Transformer，P8考官追问其计算图优化-优快云博客

场景设定

在一间安静的终面室里，候选人小华正紧张地准备展示他如何使用 PyTorch 实现一个简单的 Transformer 模型。P8 考官坐在对面，面带微笑，但目光锐利，显然对候选人的技术能力充满期待。小华深吸一口气，开始了他的演示。

第一轮：实现 Transformer

小华：（打开笔记本，展示代码）

import torch
import torch.nn as nn

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        attn_output, _ = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
        src = src + self.dropout1(attn_output)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.linear1(src)))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

小华：这是 Transformer 的编码器层实现，它主要包括多头注意力机制和前馈神经网络。通过 nn.MultiheadAttention 和 nn.Linear，我们可以轻松搭建这样的结构。

P8 考官：很好，那你能否解释一下 PyTorch 在构建这个 Transformer 时，如何管理和优化计算图？比如动态图和静态图的区别？

第二轮：计算图优化

小华：（稍微思考了一下）
嗯…… PyTorch 的动态图机制很灵活！就像我在写代码的时候，每一步操作都会自动记录到一个“大脑”里，这个“大脑”就是计算图。当我们调用 backward() 时，PyTorch 会根据这个计算图自动计算梯度。比如在 Transformer 中，每次计算注意力的时候，PyTorch 都会把相关的张量操作记录下来。

P8 考官：（追问）听起来很灵活，但动态图的缺点也很明显。比如每次前向传播时都要重新构建计算图，这会带来性能开销。你提到的优化策略是什么？

小华：（挠挠头）对……这个……其实我们可以用 torch.jit 来优化。就像把代码“冷冻”成一个静态图，这样就不需要每次都重新构建了。比如我们可以用 torch.jit.script 或 torch.jit.trace 来把 Transformer 的前向传播部分编译成静态图。

P8 考官：那你能否详细解释一下 torch.jit.trace 和 torch.jit.script 的区别？以及它们在 Transformer 中的应用？

小华：（有点紧张）好的……torch.jit.trace 是通过记录一个示例输入的计算图来生成静态图，而 torch.jit.script 是通过直接编译 Python 代码来生成静态图。在 Transformer 中，我们可以用 torch.jit.trace 来优化前向传播，因为它会记录所有的操作，包括注意力计算和前馈网络。而 torch.jit.script 可以用于更复杂的情况，比如需要条件分支的地方。

P8 考官：（微微点头）那你觉得在实际部署中，动态图和静态图哪个更高效？为什么？

小华：（稍微犹豫）嗯……静态图肯定更高效，因为它避免了每次构建计算图的开销。而且静态图可以进行一些高级优化，比如算子融合和内存布局优化。但是动态图更灵活，适合开发和调试，尤其是在需要动态控制流的场景中。

P8 考官：（严肃地）很好，那你能否举一个具体的例子，说明如何在 Transformer 中使用 torch.jit 来优化计算图？

第三轮：具体优化方案

小华：（深吸一口气）
好的，我可以用 torch.jit.trace 来优化 Transformer 的前向传播。比如：

import torch
import torch.nn as nn
import torch.jit

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        # ... （省略初始化代码）

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # ... （省略前向传播代码）

# 创建模型并示例输入
model = TransformerEncoderLayer(d_model=512, nhead=8)
x = torch.randn(10, 32, 512)  # (seq_len, batch_size, d_model)
mask = torch.ones(10, 10).tril() == 1

# 使用 torch.jit.trace 优化
traced_model = torch.jit.trace(model, (x, mask))

# 检查优化后的模型
print(traced_model.graph)

P8 考官：（露出微笑）你提到的 torch.jit.trace 做得很好。那你觉得除了 torch.jit，还有哪些方法可以进一步优化 Transformer 的内存管理和计算效率？

小华：（信心满满）还有几个方法可以优化：