[v2.0] Refactor code hierarchy (part 2) (#2987)

nni/nas_utils.py → nni/common/nas_utils.py

@@ -4,7 +4,7 @@
 import functools
 import logging
 
-from . import trial
+from .. import trial
 
 
 _logger = logging.getLogger(__name__)

nni/compression/torch/__init__.py → nni/compression/pytorch/__init__.py

@@ -2,6 +2,4 @@
 # Licensed under the MIT license.
 
 from .speedup import ModelSpeedup
-from .pruning import *
-from .quantization import *
 from .compressor import Compressor, Pruner, Quantizer

nni/compression/torch/speedup/compressor.py → nni/compression/pytorch/speedup/compressor.py

@@ -3,8 +3,8 @@
 
 import logging
 import torch
-from nni.compression.torch.utils.mask_conflict import fix_mask_conflict
-from nni.compression.torch.utils.utils import get_module_by_name
+from nni.compression.pytorch.utils.mask_conflict import fix_mask_conflict
+from nni.compression.pytorch.utils.utils import get_module_by_name
 from .compress_modules import replace_module
 from .infer_shape import ModuleMasks, infer_from_mask, infer_from_inshape, infer_from_outshape, set_conv_prune_dim
 
@@ -29,7 +29,7 @@ class ModelSpeedup:
         map_location : str
             the device on which masks are placed, same to map_location in ```torch.load```
         """
-        from nni._graph_utils import build_module_graph
+        from nni.common.graph_utils import build_module_graph
 
         self.bound_model = model
         self.masks = torch.load(masks_file, map_location)

nni/compression/torch/utils/counter.py → nni/compression/pytorch/utils/counter.py

@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn as nn
-from nni.compression.torch.compressor import PrunerModuleWrapper
+from nni.compression.pytorch.compressor import PrunerModuleWrapper
 
 try:
     from thop import profile
@@ -132,4 +132,4 @@ custom_mask_ops = {
     nn.Conv2d: count_convNd_mask,
     nn.Conv3d: count_convNd_mask,
     nn.Linear: count_linear_mask,
-}
+}
\ No newline at end of file

nni/compression/torch/utils/sensitivity_analysis.py → nni/compression/pytorch/utils/sensitivity_analysis.py

@@ -9,7 +9,7 @@ from collections import OrderedDict
 import numpy as np
 import torch.nn as nn
 
-from ..pruning.constants_pruner import PRUNER_DICT
+# FIXME: I don't know where "utils" should be
 SUPPORTED_OP_NAME = ['Conv2d', 'Conv1d']
 SUPPORTED_OP_TYPE = [getattr(nn, name) for name in SUPPORTED_OP_NAME]
 
@@ -63,6 +63,8 @@ class SensitivityAnalysis:
             This value is effective only when the early_stop_mode is set.
 
         """
+        from nni.algorithms.compression.pytorch.pruning.constants_pruner import PRUNER_DICT
+
         self.model = model
         self.val_func = val_func
         self.target_layer = OrderedDict()

nni/compression/torch/utils/shape_dependency.py → nni/compression/pytorch/utils/shape_dependency.py

@@ -17,7 +17,7 @@ class Dependency:
         """
         Build the graph for the model.
         """
-        from nni._graph_utils import TorchModuleGraph
+        from nni.common.graph_utils import TorchModuleGraph
 
         # check if the input is legal
         if traced_model is None:

nni/compression/tensorflow/__init__.py

@@ -2,4 +2,3 @@
 # Licensed under the MIT license.
 
 from .compressor import Compressor, Pruner
-from .pruning import *

nni/nnicli/nni_client.py → nni/experiment/nni_client.py

@@ -5,7 +5,7 @@
 
 Example:
 
-from nni.nnicli import Experiment
+from nni.experiment import Experiment
 
 exp = Experiment()
 exp.start_experiment('../../../../examples/trials/mnist-pytorch/config.yml')
@@ -196,16 +196,16 @@ class TrialJob:
         Trial job id.
     status: str
         Job status.
-    hyperParameters: list of `nnicli.TrialHyperParameters`
-        See `nnicli.TrialHyperParameters`.
+    hyperParameters: list of `nni.experiment.TrialHyperParameters`
+        See `nni.experiment.TrialHyperParameters`.
     logPath: str
         Log path.
     startTime: int
         Job start time (timestamp).
     endTime: int
         Job end time (timestamp).
-    finalMetricData: list of `nnicli.TrialMetricData`
-        See `nnicli.TrialMetricData`.
+    finalMetricData: list of `nni.experiment.TrialMetricData`
+        See `nni.experiment.TrialMetricData`.
     parameter_index: int
         Parameter index.
     """