support loading deit weights (open-mmlab#538)

Author: 谢昕辰

mmseg/models/backbones/vit.py

@@ -325,6 +325,8 @@ def init_weights(self, pretrained=None):
             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
 

mmseg/models/necks/multilevel_neck.py

@@ -54,7 +54,6 @@ def __init__(self,
 
     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)