diff --git a/src/sdk/pynni/nni/compression/speedup/__init__.py b/src/sdk/pynni/nni/compression/speedup/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/sdk/pynni/nni/compression/speedup/torch/compress_modules.py b/src/sdk/pynni/nni/compression/speedup/torch/compress_modules.py index 540fe115cf..5bfcc16804 100644 --- a/src/sdk/pynni/nni/compression/speedup/torch/compress_modules.py +++ b/src/sdk/pynni/nni/compression/speedup/torch/compress_modules.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. import torch -from .infer_shape import CoarseMask, ModuleMasks +from .infer_shape import ModuleMasks replace_module = { 'BatchNorm2d': lambda module, mask: replace_batchnorm2d(module, mask), diff --git a/src/sdk/pynni/nni/compression/speedup/torch/compressor.py b/src/sdk/pynni/nni/compression/speedup/torch/compressor.py index 1686a5c209..ae6b7ce015 100644 --- a/src/sdk/pynni/nni/compression/speedup/torch/compressor.py +++ b/src/sdk/pynni/nni/compression/speedup/torch/compressor.py @@ -379,7 +379,7 @@ def _find_successors(self, module_name): def infer_module_mask(self, module_name, mask=None, in_shape=None, out_shape=None): """ Infer input shape / output shape based on the module's weight mask / input shape / output shape. - + For a module: Infer its input and output shape from its weight mask Infer its output shape from its input shape diff --git a/src/sdk/pynni/nni/compression/speedup/torch/infer_shape.py b/src/sdk/pynni/nni/compression/speedup/torch/infer_shape.py index 995dcf997f..701d1f58e6 100644 --- a/src/sdk/pynni/nni/compression/speedup/torch/infer_shape.py +++ b/src/sdk/pynni/nni/compression/speedup/torch/infer_shape.py @@ -56,7 +56,7 @@ def merge_index(index_a, index_b): s.add(num) for num in index_b: s.add(num) - return torch.tensor(sorted(s)) + return torch.tensor(sorted(s)) # pylint: disable=not-callable def merge(self, cmask): """ @@ -98,7 +98,7 @@ def __init__(self, module_name): self.param_masks = dict() self.input_mask = None self.output_mask = None - + def set_param_masks(self, name, mask): """ Parameters @@ -217,7 +217,7 @@ def view_inshape(module_masks, mask, shape): TODO: consider replace tensor.view with nn.Flatten, because tensor.view is not included in module, thus, cannot be replaced by our framework. - + Parameters ---------- module_masks : ModuleMasks @@ -250,7 +250,7 @@ def view_inshape(module_masks, mask, shape): step_size = shape['in_shape'][2] * shape['in_shape'][3] for loc in mask.mask_index[1]: index.extend([loc * step_size + i for i in range(step_size)]) - output_cmask.add_index_mask(dim=1, index=torch.tensor(index)) + output_cmask.add_index_mask(dim=1, index=torch.tensor(index)) # pylint: disable=not-callable module_masks.set_output_mask(output_cmask) return output_cmask @@ -373,7 +373,6 @@ def convert_to_coarse_mask(mask): """ assert 'weight' in mask assert isinstance(mask['weight'], torch.Tensor) - cmask = None weight_mask = mask['weight'] shape = weight_mask.size() ones = torch.ones(shape[1:]).to(weight_mask.device) @@ -451,7 +450,7 @@ def conv2d_outshape(module_masks, mask): The ModuleMasks instance of the conv2d mask : CoarseMask The mask of its output tensor - + Returns ------- CoarseMask