互联网大厂经典面试题：手撕Transformer

最新推荐文章于 2025-05-30 13:43:20 发布

原创最新推荐文章于 2025-05-30 13:43:20 发布 · 575 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #机器学习 #深度学习 #transformer #python #面试经验

互联网大厂面试中，往往涉及手撕代码环节。Transformer作为现在大模型的基本架构，在学术界以及工业界都有很广泛的应用，因此成为了一个重要考点，本文着重介绍如何快速理解transformer以及通过python“手撕”实现（以演示为主，不能直接运行）。

1. 自注意力机制（Self-Attention）

自注意力机制的核心思想是：每个位置的输出由所有位置的信息加权得到。自注意力机制使得模型能够关注输入序列中不同位置之间的关系。

计算步骤：

查询（Query）、键（Key）和值（Value）：每个输入元素通过 线性变换 生成查询、键和值。
计算注意力分数：通过查询和键的点积来计算注意力分数，然后通过 Softmax 转化为权重。
加权求和：使用注意力权重加权值（Value），得到每个位置的输出。

公式：

$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

以下是完整的代码实现，面试中通常只会涉及其中的一部分：

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

# 位置编码模块，用于为输入的词嵌入加上位置信息
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()

        # 创建一个位置编码矩阵，形状为 (max_len, d_model)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()  # 位置从 0 到 max_len-1
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(torch.log(torch.tensor(10000.0)) / d_model))

        # 偶数维度使用正弦函数，奇数维度使用余弦函数
        pe[:, 0::2] = torch.sin(position * div_term)  # 位置编码的偶数维度
        pe[:, 1::2] = torch.cos(position * div_term)  # 位置编码的奇数维度
        self.register_buffer('pe', pe)  # 将位置编码保存在 buffer 中，这样它不会被训练

    def forward(self, x):
        # 将位置编码加到输入 x 上，x 是输入的词嵌入
        return x + self.pe[:x.size(1)].detach()  # 不计算梯度（detach）

# 自注意力机制（Self-Attention）模块
class SelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(SelfAttention, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.depth = d_model // n_heads  # 每个头的深度

        # 定义查询、键、值的线性变换层
        self.query = nn.Linear(d_model, d_model)
        self.key = nn.Linear(d_model, d_model)
        self.value = nn.Linear(d_model, d_model)

        # 输出的线性变换层
        self.out = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        # 将输入张量 x 按照头数（n_heads）进行拆分，方便多头注意力计算
        batch_size = x.size(0)  # 获取批次大小
        x = x.view(batch_size, -1, self.n_heads, self.depth)  # (batch_size, seq_len, n_heads, depth)
        return x.permute(0, 2, 1, 3)  # 变换维度，使其变成 (batch_size, n_heads, seq_len, depth)

    def forward(self, x):
        # 将输入 x（形状为 batch_size x seq_len x d_model）分为多个头
        query = self.split_heads(self.query(x))  # (batch_size, n_heads, seq_len, depth)
        key = self.split_heads(self.key(x))      # (batch_size, n_heads, seq_len, depth)
        value = self.split_heads(self.value(x))  # (batch_size, n_heads, seq_len, depth)

        # 计算缩放点积注意力
        score = torch.matmul(query, key.transpose(-2, -1)) / self.depth**0.5  # (batch_size, n_heads, seq_len, seq_len)
        attention = F.softmax(score, dim=-1)  # 对 score 进行 softmax 归一化

        # 根据注意力权重加权值
        output = torch.matmul(attention, value)  # (batch_size, n_heads, seq_len, depth)
        output = output.permute(0, 2, 1, 3).contiguous().view(x.size(0), -1, self.d_model)  # 将多个头的输出拼接
        return self.out(output)  # 通过线性层得到最终输出

# 前馈神经网络模块，通常包含两个全连接层和 ReLU 激活函数
class FeedForward(nn.Module):
    def __init__(self, d_model):
        super(FeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_model * 4)  # 第一层，通常会扩大维度
        self.fc2 = nn.Linear(d_model * 4, d_model)  # 第二层，恢复到 d_model

    def forward(self, x):
        x = F.relu(self.fc1(x))  # 使用 ReLU 激活函数
        return self.fc2(x)  # 输出

# Transformer 编码器层，由自注意力机制和前馈神经网络组成
class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads):
        super(EncoderLayer, self).__init__()
        self.self_attention = SelfAttention(d_model, n_heads)  # 自注意力机制
        self.ffn = FeedForward(d_model)  # 前馈神经网络
        self.norm1 = nn.LayerNorm(d_model)  # 层归一化
        self.norm2 = nn.LayerNorm(d_model)  # 层归一化

    def forward(self, x):
        # 自注意力计算
        attn_out = self.self_attention(x)
        x = self.norm1(x + attn_out)  # 残差连接 + 层归一化

        # 前馈神经网络计算
        ffn_out = self.ffn(x)
        return self.norm2(x + ffn_out)  # 残差连接 + 层归一化

# Transformer 编码器，由多个编码器层堆叠组成
class Encoder(nn.Module):
    def __init__(self, d_model, n_heads, num_layers, max_len=5000):
        super(Encoder, self).__init__()
        self.d_model = d_model
        self.positional_encoding = PositionalEncoding(d_model, max_len)  # 位置编码
        self.layers = nn.ModuleList([EncoderLayer(d_model, n_heads) for _ in range(num_layers)])  # 多个编码器层

    def forward(self, x):
        x = self.positional_encoding(x)  # 添加位置编码
        for layer in self.layers:  # 遍历每一层编码器
            x = layer(x)
        return x

# 整个 Transformer 模型（仅包含编码器部分）
class Transformer(nn.Module):
    def __init__(self, d_model, n_heads, num_layers, vocab_size, max_len=5000):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)  # 嵌入层
        self.positional_encoding = PositionalEncoding(d_model, max_len)  # 位置编码
        self.encoder = Encoder(d_model, n_heads, num_layers, max_len)  # 编码器
        self.fc_out = nn.Linear(d_model, vocab_size)  # 输出层

    def forward(self, x):
        x = self.embedding(x)  # 获取词嵌入
        x = self.positional_encoding(x)  # 添加位置编码
        x = self.encoder(x)  # 通过编码器
        return self.fc_out(x)  # 最后输出

工程上的实现方式

在工程中使用 Transformer 模型，肯定不会真的有人去从底层对其进行“手撕”，通常情况下我们会通过现有的深度学习框架（如 PyTorch 或 TensorFlow）来调用已实现的 Transformer 模型。现代的深度学习框架提供了很多预训练的 Transformer 模型，并且通常支持多种预训练模型的加载和微调，可以直接利用这些现成的工具。

Hugging Face 提供了非常方便的接口来加载、微调和使用 Transformer 模型。这个库提供了许多预训练的 Transformer 模型（如 BERT、GPT 等），并且可以很方便地进行模型加载和推理。

安装 Hugging Face Transformers 库

首先，安装 Hugging Face 的 Transformers 和 torch 库：

pip install transformers torch

加载预训练的 Transformer 模型（例如 BERT）

假设你需要加载预训练的 BERT 模型，并进行文本分类任务。下面是如何在工程中调用 Transformer 的例子。

from transformers import BertTokenizer, BertForSequenceClassification
import torch

# 1. 加载预训练模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  # 选择预训练模型
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)  # 用于二分类任务

# 2. 编写输入文本并进行编码
text = "This is a sample text for classification"
inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)

# 3. 推理：将编码后的输入数据传入模型进行预测
with torch.no_grad():  # 推理时不计算梯度
    outputs = model(**inputs)

# 4. 输出预测结果
logits = outputs.logits
prediction = torch.argmax(logits, dim=-1)  # 获取最大概率的类别
print(f"Prediction: {prediction.item()}")