mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

support loading deit weights (#538)

Browse Source
This commit is contained in:
谢昕辰 2021-05-09 11:34:18 +08:00 committed by GitHub
parent f253451b54
commit 1052f8d5d3
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mmseg/models/backbones/vit.py
View File

@ -325,6 +325,8 @@ class VisionTransformer(nn.Module):
checkpoint = _load_checkpoint(pretrained, logger=logger)
if 'state_dict' in checkpoint:
state_dict = checkpoint['state_dict']
elif 'model' in checkpoint:
state_dict = checkpoint['model']
else:
state_dict = checkpoint

1
mmseg/models/necks/multilevel_neck.py
View File

@ -54,7 +54,6 @@ class MultiLevelNeck(nn.Module):
def forward(self, inputs):
assert len(inputs) == len(self.in_channels)
print(inputs[0].shape)
inputs = [
lateral_conv(inputs[i])
for i, lateral_conv in enumerate(self.lateral_convs)