Reorganize Dirs

mmrazor/engine/__init__.py

@@ -1,5 +1,14 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .hook import DumpSubnetHook
-from .optim import SeparateOptimWrapperConstructor
+from .hooks import DumpSubnetHook
+from .optimizers import SeparateOptimWrapperConstructor
+from .runner import (AutoSlimValLoop, DartsEpochBasedTrainLoop,
+                     DartsIterBasedTrainLoop, EvolutionSearchLoop,
+                     GreedySamplerTrainLoop, SingleTeacherDistillValLoop,
+                     SlimmableValLoop)
 
-__all__ = ['SeparateOptimWrapperConstructor', 'DumpSubnetHook']
+__all__ = [
+    'SeparateOptimWrapperConstructor', 'DumpSubnetHook',
+    'SingleTeacherDistillValLoop', 'DartsEpochBasedTrainLoop',
+    'DartsIterBasedTrainLoop', 'SlimmableValLoop', 'EvolutionSearchLoop',
+    'GreedySamplerTrainLoop', 'AutoSlimValLoop'
+]

mmrazor/runners/evolution_search_loop.py → mmrazor/engine/runner/evolution_search_loop.py

@@ -13,15 +13,11 @@
 from mmengine.utils import is_list_of
 from torch.utils.data import DataLoader
 
-from mmrazor.models.subnet import Candidates, FlopsEstimator, export_fix_subnet
-from mmrazor.models.subnet.random_subnet import (MULTI_MUTATORS_RANDOM_SUBNET,
-                                                 SINGLE_MUTATOR_RANDOM_SUBNET)
 from mmrazor.registry import LOOPS
+from mmrazor.structures import Candidates, FlopsEstimator, export_fix_subnet
+from mmrazor.utils import SupportRandomSubnet
 from .utils import crossover
 
-random_subnet_type = Union[SINGLE_MUTATOR_RANDOM_SUBNET,
-                           MULTI_MUTATORS_RANDOM_SUBNET]
-
 
 @LOOPS.register_module()
 class EvolutionSearchLoop(EpochBasedTrainLoop):
@@ -215,14 +211,14 @@ def gen_crossover_candidates(self) -> List:
                 crossover_candidates.append(crossover_candidate)
         return crossover_candidates
 
-    def _mutation(self) -> random_subnet_type:
+    def _mutation(self) -> SupportRandomSubnet:
         """Mutate with the specified mutate_prob."""
         candidate1 = random.choice(self.top_k_candidates.subnets)
         candidate2 = self.model.sample_subnet()
         candidate = crossover(candidate1, candidate2, prob=self.mutate_prob)
         return candidate
 
-    def _crossover(self) -> random_subnet_type:
+    def _crossover(self) -> SupportRandomSubnet:
         """Crossover."""
         candidate1 = random.choice(self.top_k_candidates.subnets)
         candidate2 = random.choice(self.top_k_candidates.subnets)
@@ -292,7 +288,7 @@ def _save_searcher_ckpt(self) -> None:
                     if osp.isfile(ckpt_path):
                         os.remove(ckpt_path)
 
-    def _check_constraints(self, random_subnet: random_subnet_type) -> bool:
+    def _check_constraints(self, random_subnet: SupportRandomSubnet) -> bool:
         """Check whether is beyond constraints.
 
         Returns:

mmrazor/runners/subnet_sampler_loop.py → mmrazor/engine/runner/subnet_sampler_loop.py

@@ -12,13 +12,9 @@
 from mmengine.utils import is_list_of
 from torch.utils.data import DataLoader
 
-from mmrazor.models.subnet import (MULTI_MUTATORS_RANDOM_SUBNET,
-                                   SINGLE_MUTATOR_RANDOM_SUBNET, Candidates,
-                                   FlopsEstimator, export_fix_subnet)
 from mmrazor.registry import LOOPS
-
-random_subnet_type = Union[SINGLE_MUTATOR_RANDOM_SUBNET,
-                           MULTI_MUTATORS_RANDOM_SUBNET]
+from mmrazor.structures import Candidates, FlopsEstimator, export_fix_subnet
+from mmrazor.utils import SupportRandomSubnet
 
 
 class BaseSamplerTrainLoop(IterBasedTrainLoop):
@@ -48,7 +44,7 @@ def __init__(self,
             self.model = runner.model
 
     @abstractmethod
-    def sample_subnet(self) -> random_subnet_type:
+    def sample_subnet(self) -> SupportRandomSubnet:
         """Sample a subnet to train the supernet."""
 
     def run_iter(self, data_batch: Sequence[dict]) -> None:
@@ -197,7 +193,7 @@ def run(self) -> None:
         self.runner.call_hook('after_train_epoch')
         self.runner.call_hook('after_train')
 
-    def sample_subnet(self) -> random_subnet_type:
+    def sample_subnet(self) -> SupportRandomSubnet:
         """Sample a subnet from top_k candidates one by one, then to train the
         surpernet with the subnet.
 
@@ -298,18 +294,18 @@ def _val_candidate(self) -> Dict:
         metrics = self.evaluator.evaluate(len(self.dataloader_val.dataset))
         return metrics
 
-    def _sample_from_supernet(self) -> random_subnet_type:
+    def _sample_from_supernet(self) -> SupportRandomSubnet:
         """Sample from the supernet."""
         subnet = self.model.sample_subnet()
         return subnet
 
-    def _sample_from_candidates(self) -> random_subnet_type:
+    def _sample_from_candidates(self) -> SupportRandomSubnet:
         """Sample from the candidates."""
         assert len(self.candidates) > 0
         subnet = random.choice(self.candidates)
         return subnet
 
-    def _check_constraints(self, random_subnet: random_subnet_type) -> bool:
+    def _check_constraints(self, random_subnet: SupportRandomSubnet) -> bool:
         """Check whether is beyond constraints.
 
         Returns:

mmrazor/runners/utils/genetic.py → mmrazor/engine/runner/utils/genetic.py

@@ -3,12 +3,12 @@
 
 import numpy as np
 
-from mmrazor.models.subnet import SINGLE_MUTATOR_RANDOM_SUBNET
+from mmrazor.utils import SingleMutatorRandomSubnet
 
 
-def crossover(random_subnet1: SINGLE_MUTATOR_RANDOM_SUBNET,
-              random_subnet2: SINGLE_MUTATOR_RANDOM_SUBNET,
-              prob: float = 0.5) -> SINGLE_MUTATOR_RANDOM_SUBNET:
+def crossover(random_subnet1: SingleMutatorRandomSubnet,
+              random_subnet2: SingleMutatorRandomSubnet,
+              prob: float = 0.5) -> SingleMutatorRandomSubnet:
     """Crossover in genetic algorithm.
 
     Args:

mmrazor/models/__init__.py

@@ -6,4 +6,3 @@
 from .mutables import *  # noqa: F401,F403
 from .mutators import *  # noqa: F401,F403
 from .ops import *  # noqa: F401,F403
-from .subnet import *  # noqa: F401,F403

mmrazor/models/algorithms/nas/autoslim.py

@@ -11,10 +11,10 @@
 
 from mmrazor.models.distillers import ConfigurableDistiller
 from mmrazor.models.mutators import OneShotChannelMutator
-from mmrazor.models.subnet import SINGLE_MUTATOR_RANDOM_SUBNET
 from mmrazor.models.utils import (add_prefix,
                                   reinitialize_optim_wrapper_count_status)
 from mmrazor.registry import MODEL_WRAPPERS, MODELS
+from mmrazor.utils import SingleMutatorRandomSubnet
 from ..base import BaseAlgorithm
 
 VALID_MUTATOR_TYPE = Union[OneShotChannelMutator, Dict]
@@ -70,10 +70,10 @@ def _build_distiller(
 
         return distiller
 
-    def sample_subnet(self) -> SINGLE_MUTATOR_RANDOM_SUBNET:
+    def sample_subnet(self) -> SingleMutatorRandomSubnet:
         return self.mutator.sample_choices()
 
-    def set_subnet(self, subnet: SINGLE_MUTATOR_RANDOM_SUBNET) -> None:
+    def set_subnet(self, subnet: SingleMutatorRandomSubnet) -> None:
         self.mutator.set_choices(subnet)
 
     def set_max_subnet(self) -> None:

mmrazor/models/algorithms/nas/darts.py

@@ -9,10 +9,9 @@
 from torch.nn.modules.batchnorm import _BatchNorm
 
 from mmrazor.models.mutators import DiffModuleMutator
-from mmrazor.models.subnet import (FIX_MUTABLE, export_fix_subnet,
-                                   load_fix_subnet)
 from mmrazor.models.utils import add_prefix
 from mmrazor.registry import MODEL_WRAPPERS, MODELS
+from mmrazor.utils import FixMutable
 from ..base import BaseAlgorithm
 
 
@@ -47,7 +46,7 @@ class Darts(BaseAlgorithm):
     def __init__(self,
                  architecture: Union[BaseModel, Dict],
                  mutator: Optional[Union[DiffModuleMutator, Dict]] = None,
-                 fix_subnet: Optional[FIX_MUTABLE] = None,
+                 fix_subnet: Optional[FixMutable] = None,
                  unroll: bool = False,
                  norm_training: bool = False,
                  data_preprocessor: Optional[Union[dict, nn.Module]] = None,
@@ -57,6 +56,9 @@ def __init__(self,
         # Darts has two training mode: supernet training and subnet retraining.
         # fix_subnet is not None, means subnet retraining.
         if fix_subnet:
+            # Avoid circular import
+            from mmrazor.structures import load_fix_subnet
+
             # According to fix_subnet, delete the unchosen part of supernet
             load_fix_subnet(self.architecture, fix_subnet)
             self.is_supernet = False
@@ -86,6 +88,9 @@ def __init__(self,
     def search_subnet(self):
         """Search subnet by mutator."""
 
+        # Avoid circular import
+        from mmrazor.structures import export_fix_subnet
+
         subnet = self.mutator.sample_choices()
         self.mutator.set_choices(subnet)
         return export_fix_subnet(self)

mmrazor/models/algorithms/nas/spos.py

@@ -8,9 +8,8 @@
 from torch.nn.modules.batchnorm import _BatchNorm
 
 from mmrazor.models.mutators import OneShotModuleMutator
-from mmrazor.models.subnet import (SINGLE_MUTATOR_RANDOM_SUBNET,
-                                   VALID_FIX_MUTABLE_TYPE, load_fix_subnet)
 from mmrazor.registry import MODELS
+from mmrazor.utils import SingleMutatorRandomSubnet, ValidFixMutable
 from ..base import BaseAlgorithm, LossResults
 
 
@@ -68,7 +67,7 @@ class SPOS(BaseAlgorithm):
     def __init__(self,
                  architecture: Union[BaseModel, Dict],
                  mutator: Optional[Union[OneShotModuleMutator, Dict]] = None,
-                 fix_subnet: Optional[VALID_FIX_MUTABLE_TYPE] = None,
+                 fix_subnet: Optional[ValidFixMutable] = None,
                  norm_training: bool = False,
                  data_preprocessor: Optional[Union[dict, nn.Module]] = None,
                  init_cfg: Optional[dict] = None):
@@ -77,6 +76,9 @@ def __init__(self,
         # SPOS has two training mode: supernet training and subnet retraining.
         # fix_subnet is not None, means subnet retraining.
         if fix_subnet:
+            # Avoid circular import
+            from mmrazor.structures import load_fix_subnet
+
             # According to fix_subnet, delete the unchosen part of supernet
             load_fix_subnet(self.architecture, fix_subnet)
             self.is_supernet = False
@@ -101,11 +103,11 @@ def __init__(self,
 
         self.norm_training = norm_training
 
-    def sample_subnet(self) -> SINGLE_MUTATOR_RANDOM_SUBNET:
+    def sample_subnet(self) -> SingleMutatorRandomSubnet:
         """Random sample subnet by mutator."""
         return self.mutator.sample_choices()
 
-    def set_subnet(self, subnet: SINGLE_MUTATOR_RANDOM_SUBNET):
+    def set_subnet(self, subnet: SingleMutatorRandomSubnet):
         """Set the subnet sampled by :meth:sample_subnet."""
         self.mutator.set_choices(subnet)
 

mmrazor/models/algorithms/pruning/slimmable_network.py

@@ -12,7 +12,6 @@
 from torch import nn
 
 from mmrazor.models.mutators import SlimmableChannelMutator
-from mmrazor.models.subnet import load_fix_subnet
 from mmrazor.models.utils import (add_prefix,
                                   reinitialize_optim_wrapper_count_status)
 from mmrazor.registry import MODEL_WRAPPERS, MODELS
@@ -69,6 +68,8 @@ def __init__(self,
 
         # must after `prepare_from_supernet`
         if len(channel_cfg_paths) == 1:
+            # Avoid circular import
+            from mmrazor.structures import load_fix_subnet
             load_fix_subnet(self.architecture, channel_cfg_paths[0])
             self.is_deployed = True
         else:

mmrazor/models/distillers/configurable_distiller.py

@@ -6,8 +6,8 @@
 from mmengine.model import BaseModel
 from torch import nn
 
-from mmrazor.core import DistillDeliveryManager, RecorderManager
 from mmrazor.registry import MODELS
+from mmrazor.structures import DistillDeliveryManager, RecorderManager
 from ..algorithms.base import LossResults
 from .base_distiller import BaseDistiller
 

mmrazor/models/mutators/channel_mutator/channel_mutator.py

@@ -5,8 +5,8 @@
 
 from torch.nn import Module
 
-from mmrazor.core.tracer import ConcatNode, DepthWiseConvNode, PathList
 from mmrazor.registry import MODELS, TASK_UTILS
+from mmrazor.structures import ConcatNode, DepthWiseConvNode, PathList
 from ...mutables import MutableChannel
 from ..base_mutator import BaseMutator
 from ..utils import DEFAULT_MODULE_CONVERTERS

mmrazor/models/mutators/channel_mutator/slimmable_channel_mutator.py

@@ -7,10 +7,10 @@
 from torch.nn import Module
 from torch.nn.modules.batchnorm import _BatchNorm
 
-from mmrazor.core.tracer import PathList
 from mmrazor.models.architectures.dynamic_op import DynamicBatchNorm
 from mmrazor.models.mutables import SlimmableMutableChannel
 from mmrazor.registry import MODELS
+from mmrazor.structures import PathList
 from ..utils import switchable_bn_converter
 from .channel_mutator import ChannelMutator