Skip to content

575 Fix load pretrained weights issue #592

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 2, 2022
Merged

575 Fix load pretrained weights issue #592

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3a7a35f
fix load pretrain
yiheng-wang-nv Mar 1, 2022
9b67f68
del loaded weights
yiheng-wang-nv Mar 2, 2022
d8e7374
Merge branch 'master' into 575-fix-ssl-pretraining
yiheng-wang-nv Mar 2, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions self_supervised_pretraining/ssl_finetuning_train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,15 +190,15 @@ def main():
vit_dict = torch.load(pretrained_path)
vit_weights = vit_dict['state_dict']

# Delete the following variable names conv3d_transpose.weight, conv3d_transpose.bias,
# conv3d_transpose_1.weight, conv3d_transpose_1.bias as they were used to match dimensions
# while pretraining with ViTAutoEnc and are not a part of ViT backbone (this is used in UNETR)
vit_weights.pop('conv3d_transpose_1.bias')
vit_weights.pop('conv3d_transpose_1.weight')
vit_weights.pop('conv3d_transpose.bias')
vit_weights.pop('conv3d_transpose.weight')

model.vit.load_state_dict(vit_weights)
# Remove items of vit_weights if they are not in the ViT backbone (this is used in UNETR).
# For example, some variables names like conv3d_transpose.weight, conv3d_transpose.bias,
# conv3d_transpose_1.weight and conv3d_transpose_1.bias are used to match dimensions
# while pretraining with ViTAutoEnc and are not a part of ViT backbone.
model_dict = model.vit.state_dict()
vit_weights = {k: v for k, v in vit_weights.items() if k in model_dict}
model_dict.update(vit_weights)
model.vit.load_state_dict(model_dict)
del model_dict, vit_weights, vit_dict
print('Pretrained Weights Succesfully Loaded !')

elif use_pretrained==0:
Expand Down