[Fix] Compatible with mme in model loading pretrained. (#1329)

* [Fix] Compatible with mme in model loading pretrained.

* add unittest of pretrained

* fix unittest error

Author: 274869388

mmpose/models/detectors/associative_embedding.py

@@ -65,7 +65,8 @@ def __init__(self,
         self.test_cfg = test_cfg
         self.use_udp = test_cfg.get('use_udp', False)
         self.parser = HeatmapParser(self.test_cfg)
-        self.init_weights(pretrained=pretrained)
+        self.pretrained = pretrained
+        self.init_weights()
 
     @property
     def with_keypoint(self):
@@ -74,7 +75,9 @@ def with_keypoint(self):
 
     def init_weights(self, pretrained=None):
         """Weight initialization for model."""
-        self.backbone.init_weights(pretrained)
+        if pretrained is not None:
+            self.pretrained = pretrained
+        self.backbone.init_weights(self.pretrained)
         if self.with_keypoint:
             self.keypoint_head.init_weights()
 

mmpose/models/detectors/mesh.py

@@ -72,11 +72,14 @@ def __init__(self,
         self.test_cfg = test_cfg
 
         self.loss_mesh = builder.build_loss(loss_mesh)
-        self.init_weights(pretrained=pretrained)
+        self.pretrained = pretrained
+        self.init_weights()
 
     def init_weights(self, pretrained=None):
         """Weight initialization for model."""
-        self.backbone.init_weights(pretrained)
+        if pretrained is not None:
+            self.pretrained = pretrained
+        self.backbone.init_weights(self.pretrained)
         self.mesh_head.init_weights()
         if self.with_gan:
             self.discriminator.init_weights()

mmpose/models/detectors/multi_task.py

@@ -46,8 +46,8 @@ def __init__(self,
         for head in heads:
             assert head is not None
             self.heads.append(builder.build_head(head))
-
-        self.init_weights(pretrained=pretrained)
+        self.pretrained = pretrained
+        self.init_weights()
 
     @property
     def with_necks(self):
@@ -56,7 +56,9 @@ def with_necks(self):
 
     def init_weights(self, pretrained=None):
         """Weight initialization for model."""
-        self.backbone.init_weights(pretrained)
+        if pretrained is not None:
+            self.pretrained = pretrained
+        self.backbone.init_weights(self.pretrained)
         if self.with_necks:
             for neck in self.necks:
                 if hasattr(neck, 'init_weights'):

mmpose/models/detectors/pose_lifter.py

@@ -86,8 +86,8 @@ def __init__(self,
         if self.semi:
             assert keypoint_head is not None and traj_head is not None
             self.loss_semi = builder.build_loss(loss_semi)
-
-        self.init_weights(pretrained=pretrained)
+        self.pretrained = pretrained
+        self.init_weights()
 
     @property
     def with_neck(self):
@@ -125,13 +125,15 @@ def causal(self):
 
     def init_weights(self, pretrained=None):
         """Weight initialization for model."""
-        self.backbone.init_weights(pretrained)
+        if pretrained is not None:
+            self.pretrained = pretrained
+        self.backbone.init_weights(self.pretrained)
         if self.with_neck:
             self.neck.init_weights()
         if self.with_keypoint:
             self.keypoint_head.init_weights()
         if self.with_traj_backbone:
-            self.traj_backbone.init_weights(pretrained)
+            self.traj_backbone.init_weights(self.pretrained)
         if self.with_traj_neck:
             self.traj_neck.init_weights()
         if self.with_traj:

mmpose/models/detectors/top_down.py

@@ -66,8 +66,8 @@ def __init__(self,
                 keypoint_head['loss_keypoint'] = loss_pose
 
             self.keypoint_head = builder.build_head(keypoint_head)
-
-        self.init_weights(pretrained=pretrained)
+        self.pretrained = pretrained
+        self.init_weights()
 
     @property
     def with_neck(self):
@@ -81,7 +81,9 @@ def with_keypoint(self):
 
     def init_weights(self, pretrained=None):
         """Weight initialization for model."""
-        self.backbone.init_weights(pretrained)
+        if pretrained is not None:
+            self.pretrained = pretrained
+        self.backbone.init_weights(self.pretrained)
         if self.with_neck:
             self.neck.init_weights()
         if self.with_keypoint:

tests/test_models/test_bottom_up_forward.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import numpy as np
+import pytest
 import torch
 
 from mmpose.models.detectors import AssociativeEmbedding
@@ -60,6 +61,9 @@ def test_ae_forward():
                                     model_cfg['test_cfg'],
                                     model_cfg['pretrained'])
 
+    with pytest.raises(TypeError):
+        detector.init_weights(pretrained=dict())
+    detector.pretrained = model_cfg['pretrained']
     detector.init_weights()
 
     input_shape = (1, 3, 256, 256)

tests/test_models/test_mesh_forward.py

@@ -3,6 +3,7 @@
 import tempfile
 
 import numpy as np
+import pytest
 import torch
 
 from mmpose.core.optimizer import build_optimizers
@@ -49,6 +50,10 @@ def test_parametric_mesh_forward():
         loss_gan=None)
 
     detector = ParametricMesh(**model_cfg)
+
+    with pytest.raises(TypeError):
+        detector.init_weights(pretrained=dict())
+    detector.pretrained = model_cfg['pretrained']
     detector.init_weights()
 
     optimizers_config = dict(generator=dict(type='Adam', lr=0.0001))

tests/test_models/test_multitask_forward.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import numpy as np
+import pytest
 import torch
 
 from mmpose.models.detectors import MultiTask
@@ -24,7 +25,9 @@ def test_multitask_forward():
         pretrained=None,
     )
     model = MultiTask(**model_cfg)
-
+    with pytest.raises(TypeError):
+        model.init_weights(pretrained=dict())
+    model.pretrained = model_cfg['pretrained']
     # build inputs and target
     mm_inputs = _demo_mm_inputs()
     inputs = mm_inputs['img']

tests/test_models/test_pose_lifter_forward.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import mmcv
 import numpy as np
+import pytest
 import torch
 
 from mmpose.models import build_posenet
@@ -71,6 +72,9 @@ def test_pose_lifter_forward():
     cfg = mmcv.Config({'model': model_cfg})
     detector = build_posenet(cfg.model)
 
+    with pytest.raises(TypeError):
+        detector.init_weights(pretrained=dict())
+    detector.pretrained = model_cfg['pretrained']
     detector.init_weights()
 
     inputs = _create_inputs(

tests/test_models/test_top_down_forward.py

@@ -2,6 +2,7 @@
 import copy
 
 import numpy as np
+import pytest
 import torch
 
 from mmpose.models.detectors import PoseWarper, TopDown
@@ -79,6 +80,9 @@ def test_topdown_forward():
                        model_cfg['train_cfg'], model_cfg['test_cfg'],
                        model_cfg['pretrained'])
 
+    with pytest.raises(TypeError):
+        detector.init_weights(pretrained=dict())
+    detector.pretrained = model_cfg['pretrained']
     detector.init_weights()
 
     input_shape = (1, 3, 256, 256)