AlibabaResearch / AlibabaResearch/AdvancedLiterateMachinery

VGT: Prefix Handling for Checkpoint State Dictionary

オープン
#175 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
C++
スター
1.8k
フォーク
195
PR マージ指標
30日以内にマージされた PR はありません

説明

When attempting to load model weights from a checkpoint in VGT, the model's state dictionary keys do not match the checkpoint keys due to potential prefixing issues. This mismatch results in parameters not being loaded correctly, which lead to a significant increase in "total_loss".

In my opinion, no prefixes are necessary. However, if prefixes are needed, we can implement a simple check to address this issue.

In MyDetectionCheckpointer.py, check if prefixes are required for the keys:
```python
if needs_prefix(checkpoint_state_dict, model_state_dict):
new_checkpoint_state_dict = {}
for k in checkpoint_state_dict.keys():
new_checkpoint_state_dict[append_prefix(k)] = checkpoint_state_dict[k]

for k in DiT_checkpoint_state_dict.keys():
new_checkpoint_state_dict[DiT_append_prefix(k)] = DiT_checkpoint_state_dict[k]

checkpoint_state_dict = new_checkpoint_state_dict
```

Here’s the function to determine if prefixes are needed:
```python
def needs_prefix(checkpoint_state_dict, model_state_dict):
for k in checkpoint_state_dict.keys():
if k not in model_state_dict:
prefixed_key = append_prefix(k)
if prefixed_key in model_state_dict:
return True
return False
```

After implementing these changes, re-training from a checkpoint works flawlessly.

コントリビューションガイド

このリポジトリのコントリビューションガイドは索引されていません

評価

この issue はまだ評価されていません。

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。