From 5ade1e405593d4acc857a704b3fdbb346e1506c0 Mon Sep 17 00:00:00 2001 From: duzx16 Date: Thu, 13 Apr 2023 16:10:10 +0800 Subject: [PATCH] Update loading instructions --- ptuning/README.md | 5 +++-- ptuning/main.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ptuning/README.md b/ptuning/README.md index ffa45ba..acd8ca7 100644 --- a/ptuning/README.md +++ b/ptuning/README.md @@ -152,7 +152,8 @@ model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remot prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin")) new_prefix_state_dict = {} for k, v in prefix_state_dict.items(): - new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v + if k.startswith("transformer.prefix_encoder."): + new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) ``` 注意你可能需要将 `pre_seq_len` 改成你训练时的实际值。如果你是[从本地加载模型的话](https://github.com/THUDM/ChatGLM-6B#%E4%BB%8E%E6%9C%AC%E5%9C%B0%E5%8A%A0%E8%BD%BD%E6%A8%A1%E5%9E%8B)，需要将 `THUDM/chatglm-6b` 改成本地的模型路径（注意不是checkpoint路径）。 @@ -160,7 +161,7 @@ model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) (2) 如果需要加载的是旧 Checkpoint（包含 ChatGLM-6B 以及 PrefixEncoder 参数），或者进行的全参数微调，则直接加载整个 Checkpoint： ```python -model = AutoModel.from_pretrained(CHECKPOINT_PATH, config=config, trust_remote_code=True) +model = AutoModel.from_pretrained(CHECKPOINT_PATH, trust_remote_code=True) ``` 之后根据需求可以进行量化，也可以直接使用： diff --git a/ptuning/main.py b/ptuning/main.py index 193a60d..43ecdf8 100644 --- a/ptuning/main.py +++ b/ptuning/main.py @@ -118,7 +118,8 @@ def main(): prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) new_prefix_state_dict = {} for k, v in prefix_state_dict.items(): - new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v + if k.startswith("transformer.prefix_encoder."): + new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) else: model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)