先保存模型训练的嵌入向量
a_embedding_weights = model.a_embeddings.weight.data.cpu().numpy()
b_embedding_weights = model.b_embeddings.weight.data.cpu().numpy()
np.savez('model_weights.npz', a_embedding=a_embedding_weights, b_embedding=b_embedding_weights)
首先加载训练好的嵌入向量
loaded_weights = np.load('model_weights.npz')
aembes = loaded_weights['a_embedding']
bembes = loaded_weights['b_embedding']
然后在模型中定义嵌入层
class CombinedModel(nn.Module):
def __init__(self, d_model, num_a, num_b,aembes,bembes):
super(CombinedModel, self).__init__()
# 嵌入层
self.aembes_layer = nn.Embedding(num_a, d_model)
self.bembes_layer = nn.Embedding(num_b, d_model)
self.aembes_layer.weight.data.copy_(torch.from_numpy(aembes))
self.bembes_layer.weight.data.copy_(torch.from_numpy(bembes))
def forward(self):
return None
最后定义模型
mymodel = CombinedModel(d_model, num_a, num_b,aembes,bembes)