注意力机制 Transformer

# -*- coding: utf-8 -*-
import torch
import logging
import os
from torch import nn
from d2l import torch as d2l

# 配置基础日志。若外部项目已经配置日志，这里不会重复覆盖。
if not logging.getLogger().handlers:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )
logger = logging.getLogger(__name__)


def setup_writable_d2l_data_dir():
    """将 d2l 默认数据目录重定向到当前脚本目录下的可写路径。"""
    data_dir = os.path.join(os.path.dirname(__file__), "data")
    os.makedirs(data_dir, exist_ok=True)

    original_download = d2l.download

    def download_to_local(url, folder=data_dir, sha1_hash=None):
        # 统一把下载目录固定到 data_dir，避免写入只读目录 ../data。
        return original_download(url, folder=folder, sha1_hash=sha1_hash)

    d2l.download = download_to_local
    logger.info("d2l 数据目录已重定向到: %s", data_dir)

#@save
class AttentionDecoder(d2l.Decoder):
    """带有注意力机制解码器的基本接口"""
    def __init__(self, **kwargs):
        super(AttentionDecoder, self).__init__(**kwargs)

    @property
    def attention_weights(self):
        raise NotImplementedError
    
class Seq2SeqAttentionDecoder(AttentionDecoder):
    """带加性注意力(Additive Attention)的Seq2Seq解码器。

    主要流程：
    1. 先将输入token id做embedding；
    2. 每个时间步使用上一时刻隐藏状态作为query，和编码器输出做注意力计算；
    3. 将注意力上下文向量与当前输入embedding拼接后送入GRU；
    4. 通过全连接层映射到词表维度，得到每一步的预测分布。
    """
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        # 兼容不同版本 d2l 的 AdditiveAttention 构造函数：
        # 新版常见签名: (num_hiddens, dropout)
        # 旧版常见签名: (key_size, query_size, num_hiddens, dropout)
        try:
            self.attention = d2l.AdditiveAttention(num_hiddens, dropout)
        except TypeError:
            self.attention = d2l.AdditiveAttention(
                num_hiddens, num_hiddens, num_hiddens, dropout)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(
            embed_size + num_hiddens, num_hiddens, num_layers,
            dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)
        # 是否打印 init_state 的详细日志（训练时建议关闭，避免每个 batch 都打印）
        self.enable_state_logging = False
        # 是否打印解码器每一步的详细 shape 日志（训练时建议关闭，避免日志过多）
        self.enable_step_logging = False
        # 是否打印“首轮首个 batch”的向量变化（默认关闭，需要时手动打开）
        self.enable_first_epoch_vector_logging = False
        self._has_logged_first_train_batch = False
        logger.info(
            "初始化解码器: vocab_size=%d, embed_size=%d, num_hiddens=%d, num_layers=%d, dropout=%s",
            vocab_size, embed_size, num_hiddens, num_layers, dropout
        )

    def init_state(self, enc_outputs, enc_valid_lens, *_args):
        # 编码器会返回两部分：
        # 1) 每个时间步的输出 outputs，原始形状是 (num_steps, batch_size, num_hiddens)
        # 2) 最后一层（及各层）的隐藏状态 hidden_state，形状是 (num_layers, batch_size, num_hiddens)
        # 保留可变参数是为了兼容上层接口，此处不使用。
        _ = _args
        outputs, hidden_state = enc_outputs
        if self.enable_state_logging:
            logger.info(
                "init_state输入: enc_outputs=%s, hidden_state=%s, enc_valid_lens=%s",
                tuple(outputs.shape), tuple(hidden_state.shape), enc_valid_lens
            )
        # 后续注意力计算希望编码器输出按 batch 放在第一维，
        # 所以把 outputs 从 (num_steps, batch_size, num_hiddens)
        # 调整成 (batch_size, num_steps, num_hiddens)。
        state = (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)
        if self.enable_state_logging:
            logger.info(
                "init_state输出: permuted_enc_outputs=%s, hidden_state=%s",
                tuple(state[0].shape), tuple(state[1].shape)
            )
        return state

    def forward(self, X, state):
        """单次前向传播。

        参数:
            X: 解码器输入token id，形状(batch_size, num_steps)
            state: (enc_outputs, hidden_state, enc_valid_lens)
        返回:
            outputs: 词表预测，形状(batch_size, num_steps, vocab_size)
            new_state: 更新后的解码器状态
        """
        # state 中包含：
        # enc_outputs: 编码器所有时间步输出，形状 (batch_size, num_steps, num_hiddens)
        # hidden_state: 当前解码器隐藏状态，形状 (num_layers, batch_size, num_hiddens)
        enc_outputs, hidden_state, enc_valid_lens = state
        if self.enable_step_logging:
            logger.info(
                "forward开始: X=%s, enc_outputs=%s, hidden_state=%s",
                tuple(X.shape), tuple(enc_outputs.shape), tuple(hidden_state.shape)
            )
        # 先把 token id 映射成词向量，再转成“按时间步遍历”更方便的形状：
        # (batch_size, num_steps) -> embedding -> (batch_size, num_steps, embed_size)
        # -> permute -> (num_steps, batch_size, embed_size)
        X = self.embedding(X).permute(1, 0, 2)
        if self.enable_step_logging:
            logger.info("embedding后X=%s", tuple(X.shape))

        # 只在训练开始阶段打印一次“向量变化”，帮助理解训练中的数值流动。
        trace_vectors = (
            self.training
            and self.enable_first_epoch_vector_logging
            and not self._has_logged_first_train_batch
        )
        if trace_vectors:
            logger.info("首轮首batch向量追踪开始（仅打印一次）")
        outputs, self._attention_weights = [], []
        for step, x in enumerate(X):
            token_embed = x
            # 用“当前时刻的解码器隐藏状态”作为 query 去做注意力检索。
            # 这里只取最后一层隐藏状态 hidden_state[-1]，并扩维成 (batch_size, 1, num_hiddens)。
            query = torch.unsqueeze(hidden_state[-1], dim=1)
            prev_h = hidden_state[-1]
            # context 是从编码器输出中“加权汇总”得到的上下文向量，
            # 形状为 (batch_size, 1, num_hiddens)。
            context = self.attention(
                query, enc_outputs, enc_outputs, enc_valid_lens)
            # 把上下文向量和当前输入词向量拼接，作为 GRU 的输入特征。
            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
            # GRU 期望输入是 (num_steps, batch_size, input_size)，
            # 这里单步输入，所以 num_steps=1，形状变为 (1, batch_size, embed_size + num_hiddens)。
            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
            outputs.append(out)
            self._attention_weights.append(self.attention.attention_weights)
            if self.enable_step_logging:
                logger.info(
                    "step=%d: query=%s, context=%s, rnn_out=%s, hidden_state=%s",
                    step, tuple(query.shape), tuple(context.shape),
                    tuple(out.shape), tuple(hidden_state.shape)
                )
            if trace_vectors:
                embed_sample = token_embed[0].detach()
                context_sample = context[0, 0].detach()
                prev_h_sample = prev_h[0].detach()
                new_h_sample = hidden_state[-1, 0].detach()
                logger.info(
                    (
                        "vector_step=%d | "
                        "embed_norm=%.4f, context_norm=%.4f, "
                        "h_prev_norm=%.4f, h_new_norm=%.4f, "
                        "delta_h_norm=%.4f | "
                        "embed_head=%s | context_head=%s | h_new_head=%s"
                    ),
                    step,
                    embed_sample.norm().item(),
                    context_sample.norm().item(),
                    prev_h_sample.norm().item(),
                    new_h_sample.norm().item(),
                    (new_h_sample - prev_h_sample).norm().item(),
                    embed_sample[:4].cpu().tolist(),
                    context_sample[:4].cpu().tolist(),
                    new_h_sample[:4].cpu().tolist(),
                )
        # 把每个时间步的 GRU 输出拼起来，再通过线性层映射到词表大小，
        # 得到每个时间步对每个词的打分 logits，形状是 (num_steps, batch_size, vocab_size)。
        outputs = self.dense(torch.cat(outputs, dim=0))
        final_outputs = outputs.permute(1, 0, 2)
        if self.enable_step_logging:
            logger.info(
                "forward结束: logits=%s, final_outputs=%s, attention_steps=%d",
                tuple(outputs.shape), tuple(final_outputs.shape),
                len(self._attention_weights)
            )
        if trace_vectors:
            self._has_logged_first_train_batch = True
            logger.info("首轮首batch向量追踪结束")
        return final_outputs, [enc_outputs, hidden_state, enc_valid_lens]

    @property
    def attention_weights(self):
        return self._attention_weights


class TrainSeq2SeqNetAdapter(nn.Module):
    """兼容不同 d2l 版本返回值差异的训练适配器。

    某些版本的 EncoderDecoder.forward 只返回 Y_hat，
    但 train_seq2seq 期望拿到 (Y_hat, state)。
    这里统一转换为二元组，避免训练阶段解包报错。
    """

    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, enc_X, dec_X, *args):
        out = self.model(enc_X, dec_X, *args)
        if isinstance(out, tuple):
            return out
        return out, None
    
# ===== 训练配置 =====
# embed_size: 词向量维度；num_hiddens: 隐藏层维度；num_layers: RNN层数；dropout: 随机失活比例
embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
# batch_size: 每次喂给模型多少样本；num_steps: 每个样本截断/填充后的时间步长度
batch_size, num_steps = 64, 10
# lr: 学习率；num_epochs: 训练轮数；device: 自动选择 GPU(可用时) 否则 CPU
lr, num_epochs, device = 0.005, 250, d2l.try_gpu()
logger.info(
    "训练配置: embed_size=%d, num_hiddens=%d, num_layers=%d, dropout=%.3f, "
    "batch_size=%d, num_steps=%d, lr=%.4f, num_epochs=%d, device=%s",
    embed_size, num_hiddens, num_layers, dropout,
    batch_size, num_steps, lr, num_epochs, device
)
setup_writable_d2l_data_dir()

# 读取机器翻译数据集，返回：
# train_iter: 训练数据迭代器（按 batch 提供 src/tgt）
# src_vocab: 源语言词表；tgt_vocab: 目标语言词表
train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
logger.info("数据集加载完成: src_vocab_size=%d, tgt_vocab_size=%d", len(src_vocab), len(tgt_vocab))
sample_batch = next(iter(train_iter))
logger.info(
    "首个batch形状: src=%s, src_valid_len=%s, tgt=%s, tgt_valid_len=%s",
    tuple(sample_batch[0].shape), tuple(sample_batch[1].shape),
    tuple(sample_batch[2].shape), tuple(sample_batch[3].shape)
)

# 编码器负责把源语言序列编码成上下文表示。
encoder = d2l.Seq2SeqEncoder(
    len(src_vocab), embed_size, num_hiddens, num_layers, dropout)
# 解码器基于注意力机制，按时间步生成目标语言序列。
decoder = Seq2SeqAttentionDecoder(
    len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
# 需要观察 state 初始化细节时可打开（训练时默认关闭，避免日志过于庞大）：
# decoder.enable_state_logging = True
# 需要观察单步注意力细节时可打开（训练时默认关闭，避免日志过于庞大）：
# decoder.enable_step_logging = True
# 打开“首轮首个 batch”向量变化打印（推荐：理解训练过程时开启）
decoder.enable_first_epoch_vector_logging = True

# 将编码器和解码器封装成完整的 Encoder-Decoder 网络。
net = d2l.EncoderDecoder(encoder, decoder)
train_net = TrainSeq2SeqNetAdapter(net)
# 某些模块（如 LazyModule）在首个 forward 前参数尚未初始化，直接 numel() 会报错。
# 这里做“安全统计”：只统计已初始化参数，并把未初始化参数个数单独打日志。
num_params = 0
uninitialized_count = 0
for p in train_net.parameters():
    if not p.requires_grad:
        continue
    try:
        num_params += p.numel()
    except ValueError:
        uninitialized_count += 1
logger.info(
    "模型构建完成: initialized_trainable_params=%d, uninitialized_params=%d",
    num_params, uninitialized_count
)
# 启动训练：内部会完成前向、损失计算、反向传播和参数更新。
logger.info("开始训练...")
d2l.train_seq2seq(train_net, train_iter, lr, num_epochs, tgt_vocab, device)
d2l.plt.show()
logger.info("训练结束。")

# 保存到 checkpoint 的训练/模型配置，后续可用于复现实验或重建模型结构。
config = {
    "embed_size": embed_size,  # 词向量维度
    "num_hiddens": num_hiddens, # 隐藏层维度
    "num_layers": num_layers,   # RNN层数
    "dropout": dropout,          # 随机失活比例
    "batch_size": batch_size,    # 每次喂给模型多少样本
    "num_steps": num_steps,      # 每个样本截断/填充后的时间步长度
    "lr": lr,                    # 学习率
    "num_epochs": num_epochs,    # 训练轮数
    "device": str(device),       # 设备类型（GPU 或 CPU）
}

# 训练完成后保存 checkpoint，便于后续继续训练或直接做推理。
torch.save(
    {
        # 仅保存参数权重（推荐做法），加载时需先构建同结构模型再 load_state_dict。
        "model": net.state_dict(),
        # 同时保存源语言和目标语言词表，保证推理时 token/id 映射一致。
        "src_vocab": src_vocab,
        "tgt_vocab": tgt_vocab,
        # 保存关键超参数配置，方便恢复训练环境与模型定义。
        "config": config,
    },
    # checkpoint 文件名；可按需改成带时间戳或轮次的名字。
    "checkpoint.pth",
)
logger.info("模型已保存到 checkpoint.pth")

# -*- coding: utf-8 -*-
import argparse
import logging
from typing import List

import torch
from torch import nn
from d2l import torch as d2l


if not logging.getLogger().handlers:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )
logger = logging.getLogger(__name__)


class AttentionDecoder(d2l.Decoder):
    """带有注意力机制解码器的基础接口。"""

    def __init__(self, **kwargs):
        super(AttentionDecoder, self).__init__(**kwargs)

    @property
    def attention_weights(self):
        raise NotImplementedError


class Seq2SeqAttentionDecoder(AttentionDecoder):
    """与训练脚本一致的注意力解码器定义，用于加载 checkpoint。"""

    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers, dropout=0, **kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        # 兼容不同版本 d2l 的 AdditiveAttention 构造参数。
        try:
            self.attention = d2l.AdditiveAttention(num_hiddens, dropout)
        except TypeError:
            self.attention = d2l.AdditiveAttention(
                num_hiddens, num_hiddens, num_hiddens, dropout
            )
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(
            embed_size + num_hiddens, num_hiddens, num_layers, dropout=dropout
        )
        self.dense = nn.Linear(num_hiddens, vocab_size)
        self._attention_weights = []

    def init_state(self, enc_outputs, enc_valid_lens, *_args):
        _ = _args
        outputs, hidden_state = enc_outputs
        return outputs.permute(1, 0, 2), hidden_state, enc_valid_lens

    def forward(self, X, state):
        enc_outputs, hidden_state, enc_valid_lens = state
        X = self.embedding(X).permute(1, 0, 2)
        outputs, self._attention_weights = [], []
        for x in X:
            query = torch.unsqueeze(hidden_state[-1], dim=1)
            context = self.attention(query, enc_outputs, enc_outputs, enc_valid_lens)
            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
            outputs.append(out)
            self._attention_weights.append(self.attention.attention_weights)
        outputs = self.dense(torch.cat(outputs, dim=0))
        return outputs.permute(1, 0, 2), [enc_outputs, hidden_state, enc_valid_lens]

    @property
    def attention_weights(self):
        return self._attention_weights


def load_model_from_checkpoint(ckpt_path: str, device: torch.device):
    """从 checkpoint 加载词表、配置和模型参数。"""
    # 兼容 PyTorch 2.6+ 默认 weights_only=True 的行为变化。
    # 该 checkpoint 来自本地训练流程，属于可信来源，因此允许完整反序列化。
    try:
        checkpoint = torch.load(ckpt_path, map_location=device)
    except Exception:
        checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
    config = checkpoint["config"]
    src_vocab = checkpoint["src_vocab"]
    tgt_vocab = checkpoint["tgt_vocab"]

    encoder = d2l.Seq2SeqEncoder(
        len(src_vocab),
        config["embed_size"],
        config["num_hiddens"],
        config["num_layers"],
        config["dropout"],
    )
    decoder = Seq2SeqAttentionDecoder(
        len(tgt_vocab),
        config["embed_size"],
        config["num_hiddens"],
        config["num_layers"],
        config["dropout"],
    )
    net = d2l.EncoderDecoder(encoder, decoder)
    net.load_state_dict(checkpoint["model"])
    net.to(device)
    net.eval()
    logger.info("已加载模型: %s", ckpt_path)
    return net, src_vocab, tgt_vocab, config


def run_inference(
    ckpt_path: str,
    sentences: List[str],
    device: torch.device,
    num_steps: int = None,
):
    net, src_vocab, tgt_vocab, model_config = load_model_from_checkpoint(ckpt_path, device)
    infer_num_steps = num_steps if num_steps is not None else model_config["num_steps"]
    logger.info("checkpoint配置: num_steps=%s", model_config.get("num_steps"))
    logger.info("开始推理: device=%s, num_steps=%d, 样本数=%d", device, infer_num_steps, len(sentences))

    for text in sentences:
        translation, _ = d2l.predict_seq2seq(
            net, text, src_vocab, tgt_vocab, infer_num_steps, device, True
        )
        print(f"{text} => {translation}")


def parse_args():
    parser = argparse.ArgumentParser(description="加载 checkpoint.pth 并执行翻译推理")
    parser.add_argument(
        "--ckpt",
        default="/home/lxg/code/AI/code/20260423/checkpoint.pth",
        help="checkpoint 文件路径",
    )
    parser.add_argument(
        "--num_steps",
        type=int,
        default=None,
        help="推理时最大时间步；不传则使用 checkpoint 中保存的配置",
    )
    parser.add_argument(
        "--text",
        nargs="*",
        default=None,
        help="待翻译英文句子，可传多个；不传则使用内置样例",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    device = d2l.try_gpu()
    texts = args.text if args.text else ["go .", "i lost .", "he's calm .", "i'm home ."]
    run_inference(args.ckpt, texts, device, args.num_steps)

import math
import torch
from torch import nn
from d2l import torch as d2l

#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        # 头数h。多头注意力的核心思想：把同一个表示空间拆成h个子空间并行计算，
        # 让不同头关注不同位置/关系，再把结果拼回去。
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        # 下面四个线性层的作用：
        # 1) W_q/W_k/W_v：把输入映射到同一隐藏维num_hiddens，方便做点积注意力
        # 2) W_o：把多头拼接后的结果再做一次融合映射
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状:
        # (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状:
        # (batch_size，)或(batch_size，查询的个数)
        #
        # 下面三行先做线性投影，再把“多头维”拆出来：
        # 原始: (B, N, H)
        # 拆头后: (B*h, N, H/h)
        # 这样可以把每个头当作一个独立样本，直接并行送进同一个点积注意力模块。
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数，
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            # 原因：我们把batch维从B扩成了B*h，每个样本对应h个头，
            # 所以mask长度也要复制h份，保证每个头用到同样的有效长度约束。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        # 将各头结果从(B*h, N, H/h)还原并拼接到最后一维，得到(B, N, H)。
        output_concat = transpose_output(output, self.num_heads)
        # 最后再过W_o，让不同头的信息进行线性融合。
        return self.W_o(output_concat)

#@save
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 记号说明：
    # B = batch_size, N = 序列长度(查询数或键值对数), H = num_hiddens, h = num_heads
    # 目标：把(B, N, H)变成(B*h, N, H/h)，便于并行计算每个头。

    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    # 这一步把最后一维H拆成(h, H/h)。
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    # 交换维度，把“头维”提前，方便下一步与batch合并。
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    # 合并(B, h) -> (B*h)，让每个头都像一个独立样本进入attention。
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    # 输入X: (B*h, N, H/h)
    # 先还原成(B, h, N, H/h)...
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    # ...再交换成(B, N, h, H/h)...
    X = X.permute(0, 2, 1, 3)
    # ...最后拼接头维，得到(B, N, H)。
    return X.reshape(X.shape[0], X.shape[1], -1)


#@save
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 与上方同名函数一致：将(B, N, H)拆成(B*h, N, H/h)以并行计算每个头。
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    # 与上方同名函数一致：把(B*h, N, H/h)还原为(B, N, H)。
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
# 这里是“自注意力”形式：Q来自X，K/V都来自Y（若X==Y则是标准自注意力）。
# 输出形状应为(B, 查询数, H) => (2, 4, 100)
attention(X, Y, Y, valid_lens).shape

import math
import torch
from torch import nn
from d2l import torch as d2l

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                                   num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
attention(X, X, X, valid_lens).shape

import math
import torch
from torch import nn
from d2l import torch as d2l


def log_tensor_info(name, tensor, max_items=8):
    """打印张量的形状、dtype和前几个元素，便于跟踪中间结果。"""
    flat = tensor.reshape(-1)
    preview = flat[:max_items].detach().cpu().tolist()
    print(f"[LOG] {name}: shape={tuple(tensor.shape)}, dtype={tensor.dtype}")
    print(f"[LOG] {name} 前{len(preview)}个元素: {preview}")


#@save
class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        # dropout作用：训练时对“token表示 + 位置编码”的结果做随机失活，
        # 防止过拟合；推理模式下（eval）会自动关闭。
        self.dropout = nn.Dropout(dropout)

        # 创建一个足够长的位置编码表 P，形状为 (1, max_len, num_hiddens)。
        # 第0维保留batch维，方便后续与输入X做广播相加。
        self.P = torch.zeros((1, max_len, num_hiddens))

        # X中每个元素对应公式里的 pos / (10000^(2i/num_hiddens))。
        # 其中：
        # - 行索引是位置pos（0 ~ max_len-1）
        # - 列索引是偶数维对应的频率索引i
        # 得到X形状：(max_len, num_hiddens/2)
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)

        # 偶数维用sin，奇数维用cos。
        # 这样每个位置都对应一组周期不同的三角函数值，模型可据此感知相对/绝对位置信息。
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

        # 初始化日志：帮助你确认位置编码表已正确构建。
        print("[LOG] PositionalEncoding 初始化完成")
        print(f"[LOG] num_hiddens={num_hiddens}, dropout={dropout}, max_len={max_len}")
        log_tensor_info("self.P", self.P, max_items=10)

    def forward(self, X):
        # X形状通常为 (batch_size, num_steps, num_hiddens)。
        print("\n[LOG] ===== 进入 PositionalEncoding.forward =====")
        log_tensor_info("输入X", X)

        # 截取与当前序列长度匹配的位置编码：(1, num_steps, num_hiddens)
        pos_slice = self.P[:, :X.shape[1], :].to(X.device)
        log_tensor_info("截取的位置编码pos_slice", pos_slice)

        # 把位置编码加到输入上（按batch维广播）。
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        log_tensor_info("相加后的X", X)

        # 经过dropout后返回。eval()模式下这里不会随机置零。
        out = self.dropout(X)
        log_tensor_info("dropout后的输出", out)
        print("[LOG] ===== 退出 PositionalEncoding.forward =====\n")
        return out
    
# 示例参数
encoding_dim, num_steps = 32, 60

# 构造位置编码模块并切到推理模式（便于稳定观察数值）
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()

# 用全零输入，只保留“纯位置编码”效果，便于可视化观察
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]

# 打印关键张量信息，帮助理解形状流转
log_tensor_info("示例输出X", X)
log_tensor_info("用于绘图的P", P)

# 画部分维度（第6~9维）随位置变化的曲线
d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
         figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])

d2l.plt.show()

# 打印前8个位置编号的二进制表示：
# 常用于类比“不同频率基函数组合后可编码位置信息”这一直觉。
for i in range(8):
    print(f'{i}的二进制是：{i:>03b}')

# d2l.show_heatmaps 需要四维输入：(batch, heads, query_len, key_len)
# 这里把二维位置编码扩成(1, 1, num_steps, encoding_dim)做热力图展示。
P = P[0, :, :].unsqueeze(0).unsqueeze(0)
log_tensor_info("heatmap输入P", P)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
                  ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')
d2l.plt.show()

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import matplotlib.gridspec as gridspec
import matplotlib.font_manager as fm


def configure_matplotlib_chinese_font():
    """自动选择可用中文字体，避免图中中文显示异常。"""
    candidates = [
        "Noto Sans CJK SC",
        "Noto Sans CJK",
        "WenQuanYi Micro Hei",
        "WenQuanYi Zen Hei",
        "SimHei",
        "Microsoft YaHei",
        "PingFang SC",
        "Source Han Sans CN",
        "Arial Unicode MS",
    ]
    installed = {f.name for f in fm.fontManager.ttflist}
    for name in candidates:
        if name in installed:
            plt.rcParams["font.family"] = "sans-serif"
            plt.rcParams["font.sans-serif"] = [name, "DejaVu Sans"]
            break
    # 避免负号显示为方块
    plt.rcParams["axes.unicode_minus"] = False


configure_matplotlib_chinese_font()

# ===============================
# 1. 位置编码生成函数
# ===============================
def get_pe(seq_len, dim):
    pe = np.zeros((seq_len, dim))
    position = np.arange(seq_len)[:, np.newaxis]
    # 频率指数：低维度对应高频率，高维度对应低频率
    div_term = np.exp(np.arange(0, dim, 2) * -(np.log(10000.0) / dim))
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    return pe

num_steps = 200    # 足够长的步数
encoding_dim = 32  # 编码总维度
pe_matrix = get_pe(num_steps, encoding_dim)

# ===============================
# 2. 画布初始化（使用 GridSpec 布局）
# ===============================
fig = plt.figure(figsize=(10, 12))
# 创建 3x2 的布局，最后一行用来展示向量
gs = gridspec.GridSpec(3, 2, height_ratios=[1, 1, 0.8])

# A. 四个圆周表盘
axes_circles = []
dims_to_show = [0, 8, 16, 30] # 选择四个观察的维度起始点
colors = ['red', 'blue', 'green', 'orange']
pointers = []

for i in range(4):
    row, col = i // 2, i % 2
    ax = fig.add_subplot(gs[row, col])
    d = dims_to_show[i]
    ax.set_xlim(-1.2, 1.2)
    ax.set_ylim(-1.2, 1.2)
    ax.set_aspect('equal')
    ax.grid(True, linestyle='--', alpha=0.5)
    
    # 绘制背景圆圈
    circle = plt.Circle((0, 0), 1, fill=False, color='gray', alpha=0.3)
    ax.add_artist(circle)
    
    # 设置标题和速度描述
    speed_desc = ["Fast (区分邻居)", "Medium", "Medium-Slow", "Slow (感知大局)"][i]
    ax.set_title(f"Compass {i+1}: Dim {d}-{d+1}\nSpeed: {speed_desc}")
    
    # 初始化指针
    p, = ax.plot([], [], 'o-', lw=3, color=colors[i], markersize=8)
    pointers.append(p)
    axes_circles.append(ax)

# B. 位置向量可视化区域 (占据最后一整行)
ax_vec = fig.add_subplot(gs[2, :])
ax_vec.set_xlim(-0.5, encoding_dim - 0.5)
ax_vec.set_ylim(-1.1, 1.1)
ax_vec.set_title(f"Current Position Vector P (Dim 0-{encoding_dim-1})")
ax_vec.set_xticks(np.arange(encoding_dim))
ax_vec.set_xticklabels([f"D{i}" for i in range(encoding_dim)], rotation=45, fontsize=8)
ax_vec.set_ylabel("PE Value")
ax_vec.grid(axis='y', linestyle='--', alpha=0.5)

# 初始化向量柱状图
bars = ax_vec.bar(np.arange(encoding_dim), np.zeros(encoding_dim), color='purple', alpha=0.7)

# 全局位置和文本显示
info_text = fig.text(0.5, 0.015, '', ha='center', fontsize=12, fontweight='bold', color='purple')

# ===============================
# 3. 动画更新
# ===============================
def update(frame):
    artists = []
    current_pe_vector = pe_matrix[frame, :]
    
    # 3.1 更新四个圆周表盘的指针
    for i, p in enumerate(pointers):
        d = dims_to_show[i]
        # x 为 sin, y 为 cos
        x = current_pe_vector[d]
        y = current_pe_vector[d+1]
        p.set_data([0, x], [0, y])
        artists.append(p)
    
    # 3.2 更新向量柱状图
    for i, bar in enumerate(bars):
        val = current_pe_vector[i]
        bar.set_height(val)
        # 根据数值给柱子点色（正数绿色，负数红色）
        bar.set_color('green' if val > 0 else 'red')
        artists.append(bar)
    
    # 3.3 更新文本信息
    info_text.set_text(f"--- Position Index: {frame} ---")
    artists.append(info_text)
    
    return artists

ani = FuncAnimation(fig, update, frames=num_steps, interval=60, blit=True)

# 保存为GIF（保存在当前脚本同目录）
gif_path = os.path.join(os.path.dirname(__file__), "demo_07_animation.gif")
ani.save(gif_path, writer=PillowWriter(fps=15))
print(f"GIF已保存: {gif_path}")

plt.tight_layout(rect=[0, 0.04, 1, 0.96]) # 留出文本空间
plt.show()

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import matplotlib.font_manager as fm
import os


def configure_matplotlib_chinese_font():
    """自动选择可用中文字体，避免中文显示异常。"""
    candidates = [
        "Noto Sans CJK SC",
        "Noto Sans CJK",
        "WenQuanYi Micro Hei",
        "WenQuanYi Zen Hei",
        "SimHei",
        "Microsoft YaHei",
        "PingFang SC",
        "Source Han Sans CN",
        "Arial Unicode MS",
    ]
    installed = {f.name for f in fm.fontManager.ttflist}
    for name in candidates:
        if name in installed:
            plt.rcParams["font.family"] = "sans-serif"
            plt.rcParams["font.sans-serif"] = [name, "DejaVu Sans"]
            break
    plt.rcParams["axes.unicode_minus"] = False


configure_matplotlib_chinese_font()

# ===============================
# 1. 模拟参数
# ===============================
num_positions = 100   # 序列长度（相当于时间轴）
dim = 32              # 向量维度（位置矩阵每行32维）
x_axis = np.linspace(0, 10, dim)
# 位置纹理由多个正弦分量叠加而成（可视化重点）
# 第3个分量频率不宜过高，否则在32个采样点下会出现“看起来不像正弦”的锯齿感。
frequencies = [1.0, 2.5, 4.0]   # 不同频率（平滑且层次分明）
amplitude = 0.3                 # 每个分量的振幅

# 模拟一个“词向量”：比如它是一个平滑的低频波
def get_word_vector():
    # 模拟“猫”这个词的语义信号（相对稳定）
    return 0.5 * np.sin(x_axis * 0.8) + 0.2 * np.cos(x_axis * 0.3)

# 模拟位置编码：随位置变化的三角函数
def get_pe_components(pos):
    """返回当前位置下的各正弦分量列表。"""
    components = []
    for f in frequencies:
        # 相位随位置变化：位置越靠后，波形整体相位越向前推进
        comp = amplitude * np.sin(x_axis * f + pos * (f / 5.0))
        components.append(comp)
    return components


def get_pe_signal(pos, dim):
    # 位置纹理信号 = 多个正弦分量逐维相加
    components = get_pe_components(pos)
    return np.sum(components, axis=0)

word_vec = get_word_vector()

# ===============================
# 2. 画布初始化
# ===============================
fig, (ax1, ax2, ax3, ax4) = plt.subplots(
    4, 1, figsize=(10, 14), gridspec_kw={"height_ratios": [1, 1, 1, 0.9]}
)
# 增加子图间距与底部留白，避免第2图图例与第3图重叠。
plt.subplots_adjust(hspace=0.72, bottom=0.10, top=0.95)

# 词向量图
line1, = ax1.plot(x_axis, word_vec, lw=2, color='blue')
ax1.set_title("1. 原始语义信号 (Word Embedding: 'CAT')", fontsize=12)
ax1.set_ylim(-1.5, 1.5)
ax1.grid(True, alpha=0.3)

# 位置编码拆解图：多个正弦分量 + 总叠加曲线
component_colors = ['tab:blue', 'tab:green', 'tab:red']
component_lines = []
for idx, f in enumerate(frequencies):
    comp_line, = ax2.plot(
        x_axis,
        np.zeros(dim),
        lw=1.5,
        color=component_colors[idx % len(component_colors)],
        alpha=0.9,
        linestyle='--',
        label=f"分量{idx+1}: sin(freq={f})",
    )
    component_lines.append(comp_line)

line2, = ax2.plot(x_axis, np.zeros(dim), lw=2.6, color='orange', label="分量叠加结果")
ax2.set_title("2. 位置纹理信号拆解：多个正弦波叠加", fontsize=12)
ax2.set_ylim(-1.5, 1.5)
ax2.grid(True, alpha=0.3)
ax2.legend(loc='upper center', bbox_to_anchor=(0.5, -0.16), ncol=4, frameon=False, fontsize=8)

# 叠加结果图
line3, = ax3.plot(x_axis, np.zeros(dim), lw=3, color='purple')
ax3.set_title("3. 最终合成信号 (Embedding + PE)", fontsize=12)
ax3.set_ylim(-1.5, 1.5)
ax3.grid(True, alpha=0.3)

# 位置矩阵行向量图（参考 demo_07）
# 含义：在当前位置frame，取位置矩阵P的第frame行（即当前PE向量）并显示各维度数值。
ax4.set_title("4. 当前位置 PE 向量（位置矩阵的一行）", fontsize=12)
ax4.set_xlim(-0.5, dim - 0.5)
ax4.set_ylim(-1.5, 1.5)
ax4.set_ylabel("PE值")
ax4.set_xticks(np.arange(dim))
ax4.set_xticklabels([f"D{i}" for i in range(dim)], rotation=45, fontsize=8)
ax4.grid(axis='y', linestyle='--', alpha=0.3)
bars = ax4.bar(np.arange(dim), np.zeros(dim), color='purple', alpha=0.7)

# 文字提示（使用整张图坐标，而不是某个子图坐标）
pos_text = fig.text(0.5, 0.03, '', ha='center', fontsize=16, fontweight='bold', color='darkred')

# ===============================
# 3. 动画更新
# ===============================
def update(frame):
    # 获取当前位置的PE信号
    components = get_pe_components(frame)
    current_pe = np.sum(components, axis=0)
    
    # 更新PE拆解图（分量 + 总和）
    for comp_line, comp in zip(component_lines, components):
        comp_line.set_ydata(comp)
    line2.set_ydata(current_pe)
    
    # 更新叠加图（相加操作）
    combined = word_vec + current_pe
    line3.set_ydata(combined)

    # 更新“当前位置PE向量”柱状图
    for i, bar in enumerate(bars):
        val = current_pe[i]
        bar.set_height(val)
        bar.set_color('green' if val > 0 else 'red')
    
    pos_text.set_text(f"当前句子位置 (Position Index): {frame}")

    return [*component_lines, line2, line3, pos_text, *bars]

ani = FuncAnimation(fig, update, frames=num_positions, interval=100, blit=True)

# 保存为GIF（保存在当前脚本同目录）
gif_path = os.path.join(os.path.dirname(__file__), "demo_08_animation.gif")
ani.save(gif_path, writer=PillowWriter(fps=10))
print(f"GIF已保存: {gif_path}")

plt.show()