[Fix] Add MAE LayerDecayOptimizerConstructor

open-mmlab · Apr 30, 2022 · 459b155 · 459b155
1 parent c4c0c1f
commit 459b155
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmseg/core/optimizers/layer_decay_optimizer_constructor.py b/mmseg/core/optimizers/layer_decay_optimizer_constructor.py
@@ -143,7 +143,8 @@ def add_params(self, params, module, **kwargs):
                     layer_id = get_layer_id_for_convnext(
                         name, self.paramwise_cfg.get('num_layers'))
                     logger.info(f'set param {name} as id {layer_id}')
-                elif 'BEiT' in module.backbone.__class__.__name__:
+                elif 'BEiT' in module.backbone.__class__.__name__ or \
+                     'MAE' in module.backbone.__class__.__name__:
                     layer_id = get_layer_id_for_vit(name, num_layers)
                     logger.info(f'set param {name} as id {layer_id}')
                 else: