使用预训练的嵌入向量

先保存模型训练的嵌入向量

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)