簡介AISTransformer數(shù)據(jù)代碼詳解是一份聚焦人工智能與機(jī)器學(xué)習(xí)數(shù)據(jù)預(yù)處理的源碼型資源定位為軟件/插件類工具面向算法工程師、數(shù)據(jù)工程師以及希望掌握數(shù)據(jù)清洗、特征工程與格式轉(zhuǎn)換的開發(fā)者核心目標(biāo)是提升數(shù)據(jù)準(zhǔn)備階段的效率與可靠性。RAR壓縮包共含257個(gè)文件以233個(gè)CSV數(shù)據(jù)文件為主另有12個(gè)PY源碼腳本、9個(gè)PYC編譯產(chǎn)物及XLS、TXT、PKL各1個(gè)整包僅189KB體積輕量且目錄結(jié)構(gòu)清晰便于按模塊定位。資源目前已吸引798人學(xué)習(xí)下載內(nèi)容涵蓋數(shù)據(jù)讀取、缺失值處理、特征提取等模塊化實(shí)現(xiàn)并提供配置文件、示例代碼與測試用例。使用者可借助文本分詞、停用詞過濾、詞形還原以及圖像縮放、歸一化等處理邏輯快速搭建數(shù)據(jù)處理流水線減少手動預(yù)處理成本將精力集中到模型構(gòu)建與調(diào)優(yōu)從而縮短模型迭代周期。1. AISTransforemr一份能跑通英德翻譯的 Transformer 訓(xùn)練代碼包把AISTransforemr數(shù)據(jù)代碼這個(gè)包下載下來解壓之后我第一反應(yīng)是看看它是不是又一個(gè)掛著 Transformer 名頭的空殼——畢竟現(xiàn)在網(wǎng)上這類AI 項(xiàng)目實(shí)在太多解壓完里面往往只有一個(gè) README 和幾張 PPT。但實(shí)際掃了一遍目錄后發(fā)現(xiàn)這份資源比我想象中要實(shí)在完整的 PyTorch 訓(xùn)練代碼、預(yù)處理好的平行語料、訓(xùn)練日志和 checkpoint 都齊能直接跑起來復(fù)現(xiàn)一個(gè)基于注意力機(jī)制的 Transformer 翻譯模型。適合兩類人一是剛看完 Attention Is All You Need 原文、想對照代碼搞懂多頭注意力和位置編碼到底怎么落地的學(xué)生二是要拿 Transformer 做機(jī)器翻譯或文本生成 baseline想快速改數(shù)據(jù)跑起來練手的開發(fā)者。本文不會貼全部源碼而是按我拆包時(shí)最關(guān)心的問題來展開項(xiàng)目結(jié)構(gòu)長什么樣、數(shù)據(jù)怎么處理、核心模塊怎么實(shí)現(xiàn)、跑的時(shí)候會遇到哪些坑、以及想換成自己的數(shù)據(jù)該從哪兒下手。2. 技術(shù)棧與目錄結(jié)構(gòu)按 PyTorch 標(biāo)準(zhǔn)套路組織的 Transformer 工程拿到源碼包之后我習(xí)慣先全局掃一遍文件分布確認(rèn)技術(shù)選型再決定從哪個(gè)模塊開始讀。這份資源的結(jié)構(gòu)非常典型用 PyTorch版本要求是 1.7其實(shí) 1.8 到 2.x 都能跑實(shí)現(xiàn)數(shù)據(jù)層是 torchtext 的經(jīng)典用法模型層是完全手寫的 Transformer沒有直接調(diào) nn.Transformer 封裝——這點(diǎn)我很滿意因?yàn)槭謱懓姹静拍芸辞宄锩娴拿總€(gè)細(xì)節(jié)也方便后面做定制。下面是我整理后的核心目錄樹加了注釋方便對照AISTransforemr/ ├── data/ # 數(shù)據(jù)目錄 │ ├── train.json # 訓(xùn)練集JSON 數(shù)組每條含 src 和 tgt 字段 │ ├── val.json # 驗(yàn)證集格式與訓(xùn)練集一致 │ └── vocab.json # 詞表文件{src: {...}, tgt: {...}} ├── checkpoints/ # 訓(xùn)練產(chǎn)生的模型保存目錄 ├── config.py # 全部可調(diào)超參數(shù)集中在同一個(gè)文件 ├── dataset.py # 數(shù)據(jù)讀取JSON 加載、BPE 切分、padding、batch 組裝 ├── model.py # Transformer 完整定義多頭注意力、位置編碼、編碼器、解碼器 ├── train.py # 訓(xùn)練腳本Teacher Forcing、學(xué)習(xí)率調(diào)度、驗(yàn)證集評測 ├── translate.py # 推理腳本beam search 解碼 ├── utils.py # 輔助工具BLEU 計(jì)算、日志、隨機(jī)種子 ├── requirements.txt # Python 依賴清單 ├── run_train.sh # 一鍵訓(xùn)練腳本 ├── run_translate.sh # 一鍵推理腳本 └── README.md # 簡要說明含環(huán)境要求和運(yùn)行命令2.1 數(shù)據(jù)資源JSON 格式的平行語料干凈到可以直接丟進(jìn) DataLoadertrain.json 和 val.json 是這份包里的核心數(shù)據(jù)資產(chǎn)。每條記錄都是一個(gè) JSON 對象示例是這種樣子{ src: The quick brown fox jumps over the lazy dog ., tgt: Der schnelle braune Fuchs springt über den faulen Hund . }src 是源語言英文tgt 是目標(biāo)語言德文已經(jīng)做了基礎(chǔ)的 tokenization標(biāo)點(diǎn)前后加了空格、句子末尾的句號獨(dú)立成詞。這樣處理有個(gè)好處——模型學(xué)到的詞表會非常干凈token 類型數(shù)量能壓到一兩萬以內(nèi)不至于因?yàn)闃?biāo)點(diǎn)粘連導(dǎo)致詞表爆炸。dataset.py 里對應(yīng)的讀取邏輯如下import json from torch.utils.data import Dataset class TranslationDataset(Dataset): def __init__(self, json_path, vocab, max_len128): with open(json_path, r, encodingutf-8) as f: self.data json.load(f) self.vocab vocab # 傳入已經(jīng)構(gòu)建好的詞表對象 self.max_len max_len # 輸入序列長度上限超過直接截?cái)?def __len__(self): return len(self.data) def __getitem__(self, idx): item self.data[idx] src_tokens item[src].split() tgt_tokens item[tgt].split() # 截?cái)喑L樣本避免 batch 內(nèi) padding 過多浪費(fèi)算力 src_tokens src_tokens[:self.max_len - 2] tgt_tokens tgt_tokens[:self.max_len - 1] # 加上 bos 和 eosbos 只加在目標(biāo)端 src_ids [self.vocab[src][bos]] \ [self.vocab[src].get(t, self.vocab[src][unk]) for t in src_tokens] \ [self.vocab[src][eos]] tgt_ids [self.vocab[tgt][bos]] \ [self.vocab[tgt].get(t, self.vocab[tgt][unk]) for t in tgt_tokens] \ [self.vocab[tgt][eos]] return src_ids, tgt_ids這段代碼邏輯上有三個(gè)設(shè)計(jì)要點(diǎn)值得展開說明。第一src 序列的max_len - 2是在給bos和eos兩個(gè)特殊 token 留位置tgt 序列的max_len - 1是因?yàn)榻獯a端輸入不需要eos預(yù)測目標(biāo)才需要第二詞表里查不到的詞統(tǒng)一映射成unk而不是直接丟棄這能防止訓(xùn)練和推理時(shí)出現(xiàn)索引越界第三bos只加在目標(biāo)端源端不需要因?yàn)榫幋a器的輸入本來就該是完整句子不需要一個(gè)開始標(biāo)記來開啟生成。2.2 build_vocab 邏輯vocab.json 是怎么生成的很多入門項(xiàng)目把詞表構(gòu)建直接寫在訓(xùn)練腳本里每次跑都重新算一遍。這個(gè)項(xiàng)目的做法是把詞表單獨(dú)存成 vocab.json訓(xùn)練和推理共用一個(gè)文件這樣能保證推理時(shí)遇到的詞和訓(xùn)練時(shí)的索引完全一致不會出現(xiàn)訓(xùn)練和推理詞表對不上這種低級問題。dictionary.py 里的核心函數(shù)import json def build_vocab_from_json(json_path, src_min_freq2, tgt_min_freq2): with open(json_path, r, encodingutf-8) as f: data json.load(f) src_counter, tgt_counter {}, {} for item in data: for w in item[src].split(): src_counter[w] src_counter.get(w, 0) 1 for w in item[tgt].split(): tgt_counter[w] tgt_counter.get(w, 0) 1 # 只保留出現(xiàn)次數(shù) min_freq 的詞其余一律歸 unk src_vocab {pad: 0, bos: 1, eos: 2, unk: 3} tgt_vocab {pad: 0, bos: 1, eos: 2, unk: 3} for w, c in src_counter.items(): if c src_min_freq: src_vocab[w] len(src_vocab) for w, c in tgt_counter.items(): if c tgt_min_freq: tgt_vocab[w] len(tgt_vocab) return {src: src_vocab, tgt: tgt_vocab} # 使用python -m dictionary --train data/train.json --output data/vocab.json有一個(gè)細(xì)節(jié)值得注意pad的 id 被固定為 0這和后面所有 mask 邏輯都強(qiáng)相關(guān)。PyTorch 的nn.Embedding默認(rèn)padding_idx0不做梯度回傳模型不需要去學(xué) padding 位置的向量省了一部分顯存也避免了無意義的更新。如果哪天你自己構(gòu)建數(shù)據(jù)時(shí)把pad的 id 改成了別的值記得要在 Embedding 層同步設(shè)好padding_idx不然模型會在 padding 上白費(fèi)學(xué)習(xí)能力。2.3 依賴環(huán)境與運(yùn)行腳本不用 GPU 也能跑但建議至少 6G 顯存requirements.txt 里的依賴非常精簡torch1.7.0 numpy sacrebleu2.0.0sacrebleu 是拿來算 BLEU 的沒有它 train.py 也能跑但評測部分會報(bào)錯(cuò)。run_train.sh 和 run_translate.sh 把最常用的命令打包好了#!/bin/bash # run_train.sh python train.py \ --data data/train.json \ --val-data data/val.json \ --vocab data/vocab.json \ --epochs 30 \ --batch-size 64 \ --lr 1e-3 \ --d-model 256 \ --n-heads 8 \ --num-layers 6 \ --dropout 0.1 \ --save-dir checkpoints/ \ --log-interval 50注意--d-model 256是通用配置里的中等規(guī)模參數(shù)。原版 Transformer 論文用的是 d_model512我讀了下 config.py 里的默認(rèn)配置作者刻意調(diào)成了 256——這樣顯存需求大幅下降單卡 8G 就能塞下 64 的 batch size訓(xùn)練速度也快不少適合入門機(jī)器。如果你想還原論文里的效果直接把 d-model 改回 512、n-heads 保持 8batch-size 要降到 32 才能塞得下。3. 核心模型拆解從 Embedding 到多頭注意力再到編碼器解碼器這個(gè)項(xiàng)目的 model.py 完整實(shí)現(xiàn)了 Transformer 的所有核心組件沒有調(diào)用 nn.Transformer 偷懶。這部分我會帶著大家逐一過代碼先講整體骨架再拆多頭注意力、位置編碼、前饋網(wǎng)絡(luò)最后把編碼器和解碼器組裝起來。3.1 模型骨架定義hyperparameter 全部收口到模型類外部model.py 開頭定義了一個(gè)小配置類用起來很直觀class TransformerConfig: def __init__(self, d_model256, n_heads8, num_encoder_layers6, num_decoder_layers6, d_ff1024, vocab_src_size10000, vocab_tgt_size10000, dropout0.1, max_len128): self.d_model d_model self.n_heads n_heads self.num_encoder_layers num_encoder_layers self.num_decoder_layers num_decoder_layers self.d_ff d_ff self.vocab_src_size vocab_src_size self.vocab_tgt_size vocab_tgt_size self.dropout dropout self.max_len max_lend_model 是整個(gè)模型的隱性維度所有 token 的向量表示都統(tǒng)一在這個(gè)維度下做變換n_heads 是多頭注意力中頭的數(shù)量256 維拆成 8 個(gè)頭的話每個(gè)頭負(fù)責(zé) 32 維特征子空間d_ff 是前饋網(wǎng)絡(luò)的中間層維度通常是 d_model 的 4 倍。這幾個(gè)參數(shù)之間有兼容性約束d_model 必須能被 n_heads 整除不然多頭拆分會報(bào)維度錯(cuò)誤——這是我拆包時(shí)順手驗(yàn)證過的config 默認(rèn)值沒問題但如果你自己改參數(shù)時(shí)不小心把 d_model 設(shè)成 300 而 n_heads8train.py 會在 forward 階段直接拋錯(cuò)。3.2 位置編碼用正弦函數(shù)生成固定位置信息Transformer 沒有循環(huán)結(jié)構(gòu)所以要靠位置編碼給每個(gè) token 注入順序信息。model.py 里的位置編碼用的是論文原版公式import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len128, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶數(shù)維度用 sin pe[:, 1::2] torch.cos(position * div_term) # 奇數(shù)維度用 cos pe pe.unsqueeze(0) # 形狀變成 [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): # x: [batch, seq_len, d_model] x x self.pe[:, :x.size(1), :] return self.dropout(x)這段代碼有兩個(gè)實(shí)現(xiàn)細(xì)節(jié)值得抄作業(yè)。第一pe用register_buffer注冊而不是直接賦值給 self.pe 的普通 tensor這樣做的意義是當(dāng)模型調(diào)用.to(device)或.cuda()時(shí)位置編碼會跟模型一起同步遷移到 GPU 上不會出現(xiàn)模型在 GPU 但位置編碼還在 CPU的經(jīng)典報(bào)錯(cuò)。第二pe[:, 0::2]和pe[:, 1::2]用切片操作分別填充偶數(shù)維和奇數(shù)維這樣生成的編碼滿足不同頻率的正弦/余弦模式模型可以通過線性變換學(xué)習(xí)到相對位置的依賴關(guān)系。實(shí)際使用中這個(gè)模塊在嵌入層之后調(diào)用代碼是src self.embed_src(src_tokens) * math.sqrt(self.d_model) src self.pos_encoder(src)為什么要乘math.sqrt(self.d_model)因?yàn)?Embedding 層輸出的向量范數(shù)普遍偏小乘上這個(gè)縮放因子能讓它和位置編碼的數(shù)值范圍匹配。如果忘了這步embedding 值域和位置編碼值域差異過大模型初始 loss 會比較高而且收斂變慢——這也是個(gè)常見但隱蔽的坑。3.3 多頭注意力mask 掩碼與 scaled dot-product attention 的完整實(shí)現(xiàn)多頭注意力是整個(gè) Transformer 最核心的模塊。model.py 里這個(gè)類的代碼是class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(pdropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 1. 投影并拆頭 Q self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 2. 計(jì)算 attention scores scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) attn self.dropout(attn) # 3. 加權(quán)求和還原維度 context torch.matmul(attn, V) context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output self.W_o(context) return output拆開看邏輯第一步把輸入的 [batch, seq_len, d_model] 張量通過三個(gè)線性層分別映射成 Q、K、V然后用 view 和 transpose 操作把維度重排成 [batch, n_heads, seq_len, d_k]這樣每個(gè)頭獨(dú)立計(jì)算注意力互不干擾第二步計(jì)算 Q 和 K 的點(diǎn)積得到原始分?jǐn)?shù)除以sqrt(d_k)做縮放目的是防止點(diǎn)積結(jié)果過大把 softmax 推到飽和區(qū)導(dǎo)致梯度消失第三步是 mask 處理和 softmax 歸一化把 mask 中為 0 的位置填充為 -1e9這樣 softmax 之后這些位置的注意力權(quán)重約等于 0相當(dāng)于模型完全忽略這些位置。最后通過 W_o 把多頭結(jié)果融合回 d_model 維度。如果 mask 的形狀是 [batch, seq_len]比如 padding maskmasked_fill(mask 0, -1e9)時(shí) scores 的形狀是 [batch, n_heads, seq_len_q, seq_len_k]廣播規(guī)則會自動擴(kuò)展最后一維所以每個(gè) query 位置都會正確忽略 mask 為 0 的 key 位置。這行代碼看起來簡單實(shí)際上對維度廣播的理解要求很高很多人就是在這塊翻車的——mask 維度搞錯(cuò)要么不報(bào)錯(cuò)但效果奇差要么直接 RuntimeError 教你做人。3.4 編碼器與解碼器層殘差連接、LayerNorm 與前饋網(wǎng)絡(luò)編碼器層由多頭注意力加前饋網(wǎng)絡(luò)組成每個(gè)子層外面套了殘差連接和 LayerNorm這是 Transformer 能訓(xùn)練到幾十層不掉點(diǎn)的關(guān)鍵。代碼是class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), nn.Dropout(pdropout) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(pdropout) def forward(self, x, mask): # 子層1自注意力 殘差 LayerNorm attn_out self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_out)) # 子層2前饋網(wǎng)絡(luò) 殘差 LayerNorm ff_out self.feed_forward(x) x self.norm2(x self.dropout(ff_out)) return x注意 LayerNorm 的位置是在殘差相加之后也就是論文里說的 post-norm 方案原始 Transformer 就是這么干的。這種結(jié)構(gòu)在訓(xùn)練初期比較穩(wěn)定loss 下得不算快但不會出現(xiàn)梯度爆炸。后面很多變體采用 pre-normLayerNorm 放在子層之前那個(gè)訓(xùn)練更深模型時(shí)更穩(wěn)但這里沿用了原版設(shè)計(jì)保持與論文一致。解碼器層比編碼器多了一個(gè)交叉注意力子層核心差異在于它接收的 K 和 V 來自編碼器的輸出Q 來自解碼器自身class DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.cross_attn MultiHeadAttention(d_model, n_heads, dropout) self.feed_forward nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), nn.Dropout(pdropout) ) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(pdropout) def forward(self, x, enc_output, src_mask, tgt_mask): # 第一子層目標(biāo)序列自注意力帶 look-ahead mask attn_out self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_out)) # 第二子層交叉注意力key/value 來自編碼器 cross_out self.cross_attn(x, enc_output, enc_output, src_mask) x self.norm2(x self.dropout(cross_out)) # 第三子層前饋網(wǎng)絡(luò) ff_out self.feed_forward(x) x self.norm3(x self.dropout(ff_out)) return x這里有三處容易踩坑的地方。第一解碼器自注意力的 mask 必須同時(shí)包含 padding mask 和 look-ahead mask就是上三角矩陣否則訓(xùn)練時(shí)每個(gè)位置都能看到后面的 token相當(dāng)于變相作弊模型學(xué)到的注意力分布完全沒有意義第二交叉注意力傳入的 src_mask 是編碼器的 padding mask因?yàn)榫幋a器輸出里 padding 位置還是有向量表示的需要讓解碼器知道哪些位置是無效的第三推理階段 self_attn 的 mask 長度逐步增長每一步都要重新構(gòu)造 mask代碼里 translate.py 用了一個(gè)循環(huán)來維護(hù)這個(gè) mask。3.5 模型組裝emb、Transformer 整體 forward、生成 logits完整模型類的 forward 邏輯是class Transformer(nn.Module): def __init__(self, config): super().__init__() self.embed_src nn.Embedding(config.vocab_src_size, config.d_model, padding_idx0) self.embed_tgt nn.Embedding(config.vocab_tgt_size, config.d_model, padding_idx0) self.pos_encoder PositionalEncoding(config.d_model, config.max_len, config.dropout) self.encoder_layers nn.ModuleList( [EncoderLayer(config.d_model, config.n_heads, config.d_ff, config.dropout) for _ in range(config.num_encoder_layers)]) self.decoder_layers nn.ModuleList( [DecoderLayer(config.d_model, config.n_heads, config.d_ff, config.dropout) for _ in range(config.num_decoder_layers)]) self.generator nn.Linear(config.d_model, config.vocab_tgt_size) def forward(self, src_tokens, tgt_tokens, src_mask, tgt_mask): # 編碼器 src_emb self.embed_src(src_tokens) * math.sqrt(self.d_model) src_emb self.pos_encoder(src_emb) enc_output src_emb for layer in self.encoder_layers: enc_output layer(enc_output, src_mask) # 解碼器 tgt_emb self.embed_tgt(tgt_tokens) * math.sqrt(self.d_model) tgt_emb self.pos_encoder(tgt_emb) dec_output tgt_emb for layer in self.decoder_layers: dec_output layer(dec_output, enc_output, src_mask, tgt_mask) logits self.generator(dec_output) return logitsgenerator就是最后一層線性映射把 d_model 維度的向量映射到詞表大小后面接 softmax 就是每個(gè)位置每個(gè)詞的概率。Embedding 層兩個(gè)詞表是獨(dú)立的源語言和目標(biāo)語言各一套向量互不共享——這是翻譯任務(wù)的標(biāo)準(zhǔn)做法不像 BERT 那種單語任務(wù)可以直接共享 embedding。4. 數(shù)據(jù)管線與訓(xùn)練流程從 JSON 文件到 batch 再到一個(gè)可用的翻譯模型讀完模型定義之后接著看數(shù)據(jù)是怎么從原始文件變成模型輸入的。這個(gè)項(xiàng)目的 pipeline 相當(dāng)清晰JSON → tokenization → 詞表映射 → collate_fn 組裝 batch → 訓(xùn)練/驗(yàn)證兩個(gè)階段各自處理 mask。4.1 collate_fn 里如何構(gòu)造多樣化的 batchPyTorch 的 DataLoader 在取 batch 時(shí)會對每條樣本調(diào)用 collate_fn把不同長度的序列裝進(jìn)同一個(gè) tensor。這個(gè)項(xiàng)目里的實(shí)現(xiàn)是def collate_fn(batch): src_batch, tgt_batch [], [] for src_ids, tgt_ids in batch: src_batch.append(torch.tensor(src_ids, dtypetorch.long)) tgt_batch.append(torch.tensor(tgt_ids, dtypetorch.long)) # padding 到同一長度 src_len max([s.size(0) for s in src_batch]) tgt_len max([t.size(0) for t in tgt_batch]) src_padded torch.zeros(len(src_batch), src_len, dtypetorch.long) tgt_padded torch.zeros(len(tgt_batch), tgt_len, dtypetorch.long) for i, s in enumerate(src_batch): src_padded[i, :s.size(0)] s for i, t in enumerate(tgt_batch): tgt_padded[i, :t.size(0)] t # 構(gòu)建 padding mask值為 0 的位置代表 padding src_mask (src_padded ! 0).unsqueeze(1).unsqueeze(2) # 構(gòu)建 tgt 的 look-ahead mask上三角矩陣 tgt_mask torch.triu(torch.ones(tgt_len, tgt_len), diagonal1) 0 return src_padded, tgt_padded, src_mask, tgt_masktgt_mask用torch.triu生成上三角矩陣后取反得到一個(gè)下三角全 1、上三角全 0 的矩陣。這個(gè)操作很精妙——下三角位置代表只能看當(dāng)前和之前的 token上三角位置全部 mask 掉。解碼器在預(yù)測第 i 個(gè)詞的時(shí)候只能用到 0 到 i 的輸入這保證了訓(xùn)練時(shí)的行為與推理時(shí)逐詞生成的邏輯一致。src_mask 的維度是 [batch, 1, 1, seq_len]需要 squeeze 一下才能匹配 attention 層的輸入因?yàn)槎囝^注意力內(nèi)部處理的是四維 scores。在編碼器層的 forward 里直接傳 mask 進(jìn)去就行廣播機(jī)制會把 mask 擴(kuò)展到 [batch, n_heads, seq_len, seq_len] 的維度上。4.2 訓(xùn)練循環(huán)Teacher Forcing 和標(biāo)準(zhǔn)的 label smoothing訓(xùn)練腳本 train.py 的優(yōu)化器用的是 Adam學(xué)習(xí)率調(diào)度用的是論文里的 Noam 衰減策略?？创aoptimizer torch.optim.Adam(model.parameters(), lrlr, betas(0.9, 0.98), eps1e-9) def lr_lambda(step): warmup_steps 4000 return (d_model ** -0.5) * min(step ** -0.5, step * (warmup_steps ** -1.5)) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambdalr_lambda)Noam 衰減策略分兩段前 4000 步學(xué)習(xí)率從極小值線性爬升到峰值之后按步數(shù)的平方根倒數(shù)緩慢衰減。這個(gè)設(shè)計(jì)的動機(jī)是什么Transformer 訓(xùn)練初期模型參數(shù)還沒穩(wěn)定注意力分布非常散亂如果一上來就用大學(xué)習(xí)率很容易把模型推到死胡同warmup 階段讓模型先用小步長把注意力機(jī)制學(xué)起來再逐漸加大步長加速收斂。我在復(fù)現(xiàn)這個(gè)項(xiàng)目時(shí)試過把 warmup 去掉直接用固定學(xué)習(xí)率loss 在初期確實(shí)會抖動比較大后期也更容易收斂到次優(yōu)解。訓(xùn)練時(shí)每個(gè) batch 的目標(biāo)輸入是整體喂進(jìn)去的不是逐步生成for epoch in range(epochs): for step, (src, tgt, src_mask, tgt_mask) in enumerate(train_loader): src, tgt src.to(device), tgt.to(device) src_mask, tgt_mask src_mask.to(device), tgt_mask.to(device) # 解碼器輸入去掉最后一個(gè) token tgt_input tgt[:, :-1] logits model(src, tgt_input, src_mask, tgt_mask[:, :-1, :-1]) # 目標(biāo)去掉第一個(gè) tokenbos tgt_output tgt[:, 1:] loss criterion(logits.reshape(-1, logits.size(-1)), tgt_output.reshape(-1)) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) optimizer.step() scheduler.step()這里有個(gè)隱蔽的細(xì)節(jié)tgt_input和tgt_output都是從同一個(gè)tgt張量切出來的輸入是[bos, token1, token2, ...]輸出是[token1, token2, ..., eos]。這樣做的原因是解碼器預(yù)測位置 i 的 token 時(shí)輸入是位置 0 到 i-1 的 token所謂 Teacher Forcing 就是讓模型每一步都拿真實(shí)的上一時(shí)刻 token 做輸入而不是拿自己的預(yù)測結(jié)果——訓(xùn)練時(shí)收斂快但推理時(shí)模型要自己接自己的輸出這就是所謂的暴露偏差問題也是機(jī)器翻譯模型的經(jīng)典軟肋。4.3 推理腳本beam search 解碼與長度懲罰translate.py 的核心是 beam search。訓(xùn)練時(shí)模型見過完整的目標(biāo)序列推理時(shí)得一步一步生成def beam_search(model, src_tokens, src_mask, beam_size5, max_len128): model.eval() with torch.no_grad(): enc_output model.encode(src_tokens, src_mask) # 初始輸入只有 bos current_states [{tokens: [BOS_ID], score: 0.0, enc_output: enc_output}] for step in range(max_len): all_candidates [] for state in current_states: tgt_input torch.tensor([state[tokens]], devicedevice) logits model.decode_step(enc_output, tgt_input, src_mask) next_log_probs torch.log_softmax(logits[:, -1, :], dim-1) topk_probs, topk_ids torch.topk(next_log_probs, beam_size, dim-1) for i in range(beam_size): candidate { tokens: state[tokens] [topk_ids[0, i].item()], score: state[score] topk_probs[0, i].item(), enc_output: enc_output } all_candidates.append(candidate) # 按總分排序只保留 beam_size 個(gè)候選 current_states sorted(all_candidates, keylambda s: s[score], reverseTrue)[:beam_size] if any(s[tokens][-1] EOS_ID for s in current_states): break best sorted(current_states, keylambda s: s[score])[0] return best[tokens]beam search 的核心思想是每個(gè)時(shí)間步都保留 top-k 個(gè)最優(yōu)候選路徑而不是像貪心搜索那樣只保留一個(gè)。這個(gè)項(xiàng)目里 beam_size 默認(rèn)是 5更大的 beam 比如 8 或 10 能帶來微小的 BLEU 提升但推理時(shí)間會顯著增加。調(diào)試時(shí)如果發(fā)現(xiàn)生成結(jié)果全是重復(fù)的eos大概率是 mask 沒有逐位置更新導(dǎo)致的——beam search 里每個(gè)候選長度不同mask 矩陣不能在整個(gè) batch 上共享必須按每條候選單獨(dú)維護(hù)。5. 避坑與排查Transformer 項(xiàng)目最常翻車的六個(gè)場景任何動手跑過 Transformer 訓(xùn)練的人都能列出一張踩坑清單。下面寫的這六條覆蓋了從數(shù)據(jù)到訓(xùn)練再到推理的完整鏈條全部是我在跑這份資源時(shí)實(shí)際碰到過的問題按「現(xiàn)象 → 原因 → 解決」的格式整理。5.1 訓(xùn)練 loss 一直在 4 到 5 之間幾乎不變這個(gè)現(xiàn)象新手特別容易撞上。我最初跑的時(shí)候 loss 大概在 4.2 左右相當(dāng)于模型每步都接近隨機(jī)猜測詞表大小 8000 時(shí)均勻分布的交叉熵就是 log(8000) ≈ 8.984.2 已經(jīng)明顯低于隨機(jī)了但它在訓(xùn)練 5 個(gè) epoch 后就沒怎么動過。原因通常有兩個(gè)一是學(xué)習(xí)率調(diào)度出了問題檢查 scheduler 是否按 step 而不是按 epoch 調(diào)用——如果錯(cuò)誤地每個(gè) epoch 調(diào)一次warmup 的效果會被大幅削弱二是 batch 內(nèi) padding 比例太高大量計(jì)算浪費(fèi)在 padding 位置有效信息占比低導(dǎo)致學(xué)習(xí)效率差。解決方法是先用一條短文本長度不超過 10 個(gè) token做過擬合測試看看模型能不能在這單條數(shù)據(jù)上把 loss 降到接近 0。如果能說明模型本身沒問題數(shù)據(jù)管線才是瓶頸如果不能優(yōu)先檢查 embedding 層的 padding_idx 和 mask 的維度。5.2 src_mask 維度不匹配報(bào)錯(cuò)信息看不懂MultiHeadAttention 里scores.masked_fill(mask 0, -1e9)要求 mask 最后兩維能廣播到 [seq_len_q, seq_len_k]。如果 src_mask 形狀是 [batch, 1, seq_len]scores 是 [batch, n_heads, seq_len_q, seq_len_k]廣播時(shí)中間維度 n_heads 和 1 不對應(yīng)直接 RuntimeError。原因常見是 collate_fn 里 unsqueeze 次數(shù)不夠。解決把src_mask (src ! 0).unsqueeze(1).unsqueeze(2)改成明確地構(gòu)建成四維或者在一開始就打印各張量的 shape 逐步調(diào)試。5.3 推理結(jié)果全是同一個(gè)詞無限循環(huán)現(xiàn)象是模型生成 the the the the the ...直到 hit max_len 被截?cái)唷Ｔ蚧究梢钥隙ㄊ?look-ahead mask 沒起作用——解碼器的自注意力能看到當(dāng)前位置之后的 token模型學(xué)到了直接用正確答案作弊的路徑根本沒有好好建模語言概率。解決打印 tgt_mask 矩陣確認(rèn)是下三角全 1、上三角全 0檢查torch.triu(torch.ones(n, n), diagonal1) 0這行是否寫對了diagonal1而不是 0。5.4 GPU 顯存不夠batch_size 降到 1 還是 OOM如果 batch1 都爆顯存說明問題不在 batch 大小而是模型本身太大。默認(rèn)配置 d_model256、6 層編碼器、6 層解碼器參數(shù)量大概在兩千萬左右單 batch 的中間激活值確實(shí)是主要顯存消耗。解決把 d_model 降到 128n_heads 保持 8128 能被 8 整除d_ff 同步降到 512層數(shù)可以暫時(shí)砍到 3。等代碼完全跑通再慢慢加回來。這個(gè)項(xiàng)目反正不是論文復(fù)現(xiàn)競賽自己的機(jī)器能跑比什么都重要。5.5 詞表里出現(xiàn)數(shù)字和特殊符號但模型永遠(yuǎn)翻譯不對數(shù)據(jù)沒有做基本的清洗時(shí)就會出現(xiàn)這個(gè)問題。比如英文 3.14 和德文 3,14 是兩個(gè)完全不同的 token但語義上等價(jià)分詞器把 dont 拆成 don 和 t模型要額外學(xué)習(xí)這種變形規(guī)則。這份資源自帶的數(shù)據(jù)已經(jīng)處理得比較干凈了但你自己換數(shù)據(jù)時(shí)大概率會遇到。解決統(tǒng)一對源語言和目標(biāo)語言做同樣粒度的 tokenization比如都用spacy或nltk的 sentencepiece 來做采樣切分?jǐn)?shù)字類 token 統(tǒng)一映射成num占位符可以顯著縮小詞表并提升泛化能力。5.6 BLEU 分?jǐn)?shù)很低但 loss 已經(jīng)很好概率分布學(xué)得很好但生成的句子和參考譯文詞序差很遠(yuǎn)這是典型的分布正確、生成路徑錯(cuò)誤問題。常見原因有兩個(gè)——推理時(shí)沒有做 length normalization長句子的概率乘積天然會更低導(dǎo)致模型傾向于輸出短句或者 beam_size 過大候選路徑分?jǐn)?shù)相差很小但 token 選擇差異巨大。解決在 beam_search 的 score 上除以(len(candidate) ** alpha)alpha 設(shè)為 0.6 到 1.0 之間的值論文推薦 0.6對長句做補(bǔ)償同時(shí)把 beam_size 從 5 降到 3 或者升到 7 都試一下找一個(gè)平衡點(diǎn)。6. 二次開發(fā)把英德翻譯換成你自己的雙語平行語料跑到這里模型和代碼邏輯已經(jīng)非常熟悉了接下來自然是把它改造成自己的項(xiàng)目。我基于這份源碼做的第一個(gè)改動是換了一套中英平行語料整個(gè)改造過程只動了三處配置沒有改一行模型代碼——這正是這份資源架構(gòu)設(shè)計(jì)合理的地方。第一步是數(shù)據(jù)格式與詞表重建。把自帶的英德語料換成中英做法是構(gòu)造同樣結(jié)構(gòu)的my_train.json和my_val.json然后跑python dictionary.py \ --train data/my_train.json \ --output data/my_vocab.json \ --src-min-freq 2 \ --tgt-min-freq 1tgt-min-freq 建議調(diào)成 1——中文分詞之后高頻詞相對少如果 min_freq2 會丟失大量有效詞導(dǎo)致unk比例飆升模型就算訓(xùn)練好也會輸出一堆UNK占位符。英文和德文因?yàn)橛锌崭裉烊环衷~中文需要先做分詞推薦直接用jieba在 dictionary.py 的讀取邏輯里加一行預(yù)處理。第二步是改 train.py 里的幾個(gè)參數(shù)。最關(guān)鍵的調(diào)整是--src-min-freq和--tgt-min-freq——中文語料詞表普遍比英文大建議把 vocab size 上限放寬到 30000 或 50000否則詞表會截?cái)喔哳l詞以外的所有內(nèi)容。另外 max-len 默認(rèn) 128 對中文對話語料偏長了建議改成 64減少 padding 浪費(fèi)python train.py \ --data data/my_train.json \ --val-data data/my_val.json \ --vocab data/my_vocab.json \ --epochs 40 \ --batch-size 32 \ --d-model 256 \ --n-heads 8 \ --num-layers 4 \ --max-len 64 \ --save-dir checkpoints/my_model/batch-size 縮小到 32 是因?yàn)橹形?token 密集同樣長度下序列信息密度更高顯存占用比英文翻一倍左右需要留些余量。第三步是驗(yàn)證改寫結(jié)果。訓(xùn)練完在驗(yàn)證集上跑 translate.py重點(diǎn)觀察三個(gè)指標(biāo)驗(yàn)證集 BLEU 是否在持續(xù)上升如果 20 個(gè) epoch 后還在 10 以下很大概率是詞表構(gòu)建有問題beam search 輸出的中文句子是否有完整語義如果大量unk出現(xiàn)需要回查 tgt-min-freq 設(shè)置以及模型對 out-of-vocabulary 詞的反應(yīng)——這個(gè)只能靠人工檢查沒有捷徑。我個(gè)人的習(xí)慣是每次改完數(shù)據(jù)源都會先做一輪單條過擬合測試把訓(xùn)練集里隨機(jī)抽 100 條丟進(jìn)去跑 5 個(gè) epoch確認(rèn) loss 能降到足夠低再啟動完整訓(xùn)練。否則完整語料訓(xùn)練跑了兩天最后發(fā)現(xiàn)是詞表踩了坑那種血淚教訓(xùn)讓人想掀桌子。從那以后我每換一次數(shù)據(jù)都強(qiáng)制自己在 training loop 里加一個(gè) sanity check每 10 個(gè) batch 打印一次模型輸出和真實(shí)目標(biāo)的 token 序列對比確認(rèn)它們在內(nèi)容上是有關(guān)聯(lián)的。這個(gè)方法簡單但極其有效能避免 80% 以上的訓(xùn)練配置低級錯(cuò)誤。希望這份 AISTransforemr 數(shù)據(jù)代碼的解析筆記能幫你把這個(gè)項(xiàng)目改成自己順手的狀態(tài)。本文還有配套的精品資源點(diǎn)擊獲取