Ошибка несоответствия размера при загрузке State Dict для TransformerDecoderModel ⇐ Python

[Anonymous]

Я столкнулся с проблемой при попытке загрузить сохраненный словарь состояний модели в мою 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]).
This happens because there's a discrepancy between the vocab_size used during training and the one defined at the time of loading the model.

Вот соответствующие части моего кода:
Обучение: модель была обучена с vocab_size, равным 10000.
Загрузка: после при загрузке модели создается экземпляр с vocab_size, который кажется другим, а именно 6313.
Разница в размере словаря, вероятно, связана с изменениями в токенизаторе после обучения. Как я могу решить эту проблему? Есть ли способ обеспечить согласованность 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
Что я пробовал
Я стремился устранить ошибку несоответствия размеров, гарантируя, что модель сможет загружаться, несмотря на различия в размере словаря. Чтобы добиться этого, я изменил функцию load_model, чтобы сделать процесс загрузки более гибким.
Загрузка нестрогого состояния: изначально при попытке загрузки я столкнулся с ошибкой несоответствия размеров. словарь состояния с помощью model.load_state_dict(контрольная точка). Чтобы решить эту проблему, я изменил строку на model.load_state_dict(checkpoint, strict=False) в надежде, что это позволит модели игнорировать несовпадающие параметры (embed.weight, fc.weight и fc.bias) и загружать остальные. весов успешно. Однако это не решило проблему, поскольку эти уровни имеют решающее значение для работы модели и зависят от правильного размера словаря.
Оптимизация кода: Кроме того, я проверил свой код, чтобы убедиться в отсутствии других потенциальных проблем, вызывающих несоответствие. Я убедился, что токенизатор, используемый для обучения и загрузки, один и тот же, и что параметр vocab_size постоянно устанавливается на 10000 на обоих этапах.
Чего я ожидал
Я ожидал, что установив strict=False, Модель сможет правильно загрузить большую часть своих весов, что позволит мне запустить модель, не столкнувшись с ошибкой несоответствия размеров. Несмотря на то, что я знал, что этот подход не устранит основную причину разницы в размере словаря, я надеялся, что он, по крайней мере, предоставит обходной путь, который позволит модели работать нормально.
Однако даже после этого Несмотря на эти изменения, модель по-прежнему не загружается должным образом из-за критической роли слоев внедрения и вывода по отношению к размеру словаря. Поэтому мне нужно руководство о том, как обеспечить согласованность vocab_size между обучением и загрузкой или как настроить загруженные веса в соответствии с новым vocab_size.

Подробнее здесь: https://stackoverflow.com/questions/793 ... codermodel