TransformerDecoderModel加载状态字典时出现尺寸不匹配错误的解决方案求助

大家好，我最近在加载自己训练的TransformerDecoderModel状态字典时遇到了一个棘手的问题，折腾了半天没解决，来求助各位大佬！报错信息如下：

RuntimeError: Error(s) in loading state_dict for TransformerDecoderModel:
size mismatch for embed.weight: copying a param with shape torch.Size([10000, 128]) from checkpoint, the shape in current model is torch.Size([6313, 128]).
size mismatch for fc.weight: copying a param with shape torch.Size([10000, 128]) from checkpoint, the shape in current model is torch.Size([6313, 128]).
size mismatch for fc.bias: copying a param with shape torch.Size([10000]) from checkpoint, the shape in current model is torch.Size([6313]).

查了一下原因，是训练模型时用的vocab_size是10000，但加载时模型的vocab_size变成了6313——应该是训练完之后分词器（tokenizer）有改动，导致词汇表大小不一致了。现在我想解决这个问题，要么能保证训练和加载时vocab_size完全一致，要么能调整预训练的权重适配新的vocab_size，求各位给点思路！

相关代码与环境 #

完整代码片段

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import json
from bpe_tokenizer import BpeTokenizer
import os

class TransformerDecoderModel(nn.Module):
    def __init__(self, vocab_size, embed_size, num_heads, hidden_dim, num_layers):
        super(TransformerDecoderModel, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.positional_encoding = PositionalEncoding(embed_size)
        decoder_layer = nn.TransformerDecoderLayer(d_model=embed_size, nhead=num_heads, dim_feedforward=hidden_dim)
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(embed_size, vocab_size)

    def forward(self, src, tgt):
        # Embedding and positional encoding for both src and tgt
        src = self.embed(src) * torch.sqrt(torch.tensor(src.size(-1)).float())
        tgt = self.embed(tgt) * torch.sqrt(torch.tensor(tgt.size(-1)).float())
        src = self.positional_encoding(src)
        tgt = self.positional_encoding(tgt)
        out = self.transformer_decoder(tgt, src)
        out = self.fc(out)
        return out

def train_model(data_path, tokenizer_path, model_path, vocab_size, min_freq, epochs=1, batch_size=2, grad_accum_steps=16):
    if os.path.exists(tokenizer_path):
        tokenizer = BpeTokenizer(tokenizer_path)
        print("Loaded existing tokenizer.")
    else:
        tokenizer = BpeTokenizer()
        tokenizer.train([data_path], vocab_size, min_freq)
        tokenizer.save(tokenizer_path)
        print("Trained and saved tokenizer.")

    with open(data_path, 'r', encoding='utf-8') as f:
        data = json.load(f)['data']

    dataset = ChatDataset(data, tokenizer)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    device = torch.device('cpu')
    model = TransformerDecoderModel(vocab_size, embed_size=128, num_heads=2, hidden_dim=256, num_layers=2).to(device)
    criterion = nn.CrossEntropyLoss(ignore_index=-100)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for i, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, input_ids)  # Assuming src and tgt are the same here
            loss = criterion(outputs.view(-1, vocab_size), labels.view(-1))
            loss = loss / grad_accum_steps
            loss.backward()

            if (i + 1) % grad_accum_steps == 0:
                optimizer.step()
                optimizer.zero_grad()

            total_loss += loss.item() * grad_accum_steps

        print(f'Epoch {epoch + 1}, Loss: {total_loss / len(dataloader)}')

    torch.save(model.state_dict(), model_path)

def load_model(model_path, tokenizer_path):
    tokenizer = BpeTokenizer(tokenizer_path)

    vocab_size = len(tokenizer)
    embed_size = 128
    num_heads = 2
    hidden_dim = 256
    num_layers = 2

    model = TransformerDecoderModel(vocab_size=vocab_size, embed_size=embed_size, num_heads=num_heads, hidden_dim=hidden_dim, num_layers=num_layers)

    checkpoint = torch.load(model_path, map_location=torch.device('cpu'), weights_only=True)
    model.load_state_dict(checkpoint, strict=False)

    return model, tokenizer

if __name__ == "__main__":
    data_path = 'train_data.json'
    tokenizer_path = 'tokenizer.json'
    model_path = 'chat_model.pth'
    vocab_size = 10000
    min_freq = 2
    
    train_model(data_path, tokenizer_path, model_path, vocab_size, min_freq)

    # Attempt to load the trained model
    try:
        model, tokenizer = load_model(model_path, tokenizer_path)
    except RuntimeError as e:
        print(e)

运行环境

• Python版本：3.10.11
• PyTorch版本：2.5.1

我已经尝试过的解决思路 #

1. 非严格加载状态字典：最开始直接调用model.load_state_dict(checkpoint)会直接抛出尺寸不匹配的错误，后来改成model.load_state_dict(checkpoint, strict=False)，本来想着能忽略不匹配的参数，先把其他核心层的权重加载进来凑合用，但后来发现embed和fc层是模型的核心，没有正确的权重根本没法正常生成结果，这个方法只能算是自欺欺人。
2. 检查代码一致性：我反复确认了训练和加载时用的是同一个tokenizer.json文件，训练时也确实把vocab_size设成了10000，但加载时len(tokenizer)返回的却是6313——难道是我保存分词器的时候没把vocab_size的信息存进去？还是分词器加载时的逻辑有问题？

我的需求 #

现在我有两个方向的诉求：

1. 最好能彻底解决根源问题，保证训练和加载时的vocab_size完全一致，以后再也不出现这种情况；
2. 如果实在没法回到原来的vocab_size，有没有办法调整预训练的embed和fc层的权重，适配现在的6313大小的词汇表？

真心希望各位大佬能给我一些建议，谢谢大家了！

备注：内容来源于stack exchange，提问作者Mlshs