正确地将转换后的pytorch模型加载到拥抱面变压器中

时间:2020-08-29 14:57:10

标签: python tensorflow machine-learning pytorch huggingface-transformers

我使用以下函数将预先训练的tf模型转换为pytorch。

def convert_tf_checkpoint_to_pytorch(*, tf_checkpoint_path, albert_config_file, pytorch_dump_path):
    # Initialise PyTorch model
    config = AlbertConfig.from_json_file(albert_config_file)
    print("Building PyTorch model from configuration: {}".format(str(config)))
    model = AlbertForPreTraining(config)

    # Load weights from tf checkpoint
    load_tf_weights_in_albert(model, config, tf_checkpoint_path)

    # Save pytorch-model
    print("Save PyTorch model to {}".format(pytorch_dump_path))
    torch.save(model.state_dict(), pytorch_dump_path)

我通过以下方式加载转换后的模型并编码句子:

def vectorize_sentence(text):
    albert_tokenizer = AlbertTokenizer.from_pretrained("albert-base-v2")
    config = AlbertConfig.from_pretrained(config_path, output_hidden_states=True)
    model = TFAlbertModel.from_pretrained(pytorch_dir, config=config, from_pt=True)
    e = albert_tokenizer.encode(text, max_length=512)
    model_input = tf.constant(e)[None, :]  # Batch size 1
    output = model(model_input)

    v = [0] * 768
    # generate sentence vectors by averaging the word vectors
    for i in range(1, len(model_input[0]) - 1):
        v = v + output[0][0][i].numpy()

    vector = v/len(model_input[0])
    return vector

但是在加载模型时,会出现警告:

PyTorch模型TFAlbertModel的某些权重或缓冲区不是 从TF 2.0模型初始化并重新初始化: ['predictions.LayerNorm.bias','predictions.dense.weight', 'predictions.LayerNorm.weight','sop_classifier.classifier.bias', 'predictions.dense.bias','sop_classifier.classifier.weight', 'predictions.decoder.bias','predictions.bias', 'predictions.decoder.weight']您可能应该在 下游任务,以便能够将其用于预测和推断。

有人可以告诉我我做错了什么吗?警告是什么意思?我在《变形金刚》的github仓库上看到了问题#5588。不知道我的问题是否与此相同。

1 个答案:

答案 0 :(得分:1)

我认为您可以尝试使用

model = AlbertModel.from_pretrained

代替

model = TFAlbertModel.from_pretrained

VectorizeSentence定义中的

AlbertModel是pytorch格式模型的类的名称,而TFAlbertModel是tensorflow格式模型的类的名称。

我不确定load_tf_weights_in_albert()的确切功能,但是我认为一旦完成,您的模型就是pytorch格式。