Can nni support to prune model tranformed from onnx format

[Sugar929]

Describe the issue:
after I use some open-source tool to transform an onnx model into a pytorch model, can I still use nni to prune it?

If I can, how to set the config_list? Because the transformed pytorch model have different named_modules

[J-shang]

I think it can, could you give us an example script to get this transformed model, then we can check this?

[Sugar929]

I chose the DeepLabv3 model as an example, and the onnx2torch tool(https://github.com/ENOT-AutoDL/onnx2torch) to convert it from onnx to pytorch format.
and part of the converting result was like:

GraphModule(
  (Shape_0): OnnxShape()
  (Constant_0): OnnxConstant()
  (Gather_0): OnnxGather()
  (Shape_1): OnnxShape()
  (Constant_1): OnnxConstant()
  (Gather_1): OnnxGather()
  (Conv_0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (BatchNormalization_0): BatchNorm2d(32, eps=9.999999747378752e-06, momentum=0.10000002384185791, affine=True, track_running_stats=True)
  (LeakyRelu_0): LeakyReLU(negative_slope=0.009999999776482582)
  (Conv_1): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
  (BatchNormalization_1): BatchNorm2d(64, eps=9.999999747378752e-06, momentum=0.10000002384185791, affine=True, track_running_stats=True)
  (LeakyRelu_1): LeakyReLU(negative_slope=0.009999999776482582)
  (MaxPool_0): MaxPool2d(kernel_size=[3, 3], stride=[2, 2], padding=[1, 1], dilation=1, ceil_mode=False)
  (Conv_2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (BatchNormalization_2): BatchNorm2d(64, eps=9.999999747378752e-06, momentum=0.10000002384185791, affine=True, track_running_stats=True)
  (LeakyRelu_2): LeakyReLU(negative_slope=0.009999999776482582)
  (Conv_3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
...

def forward(self, input_1):
    shape_0 = self.Shape_0(input_1)
    constant_0 = self.Constant_0()
    gather_0 = self.Gather_0(shape_0, constant_0);  shape_0 = constant_0 = None
    shape_1 = self.Shape_1(input_1)
    constant_1 = self.Constant_1()
    gather_1 = self.Gather_1(shape_1, constant_1);  shape_1 = constant_1 = None
    conv_0 = self.Conv_0(input_1);  input_1 = None
    batch_normalization_0 = self.BatchNormalization_0(conv_0);  conv_0 = None
    leaky_relu_0 = self.LeakyRelu_0(batch_normalization_0);  batch_normalization_0 = None
    conv_1 = self.Conv_1(leaky_relu_0);  leaky_relu_0 = None
    batch_normalization_1 = self.BatchNormalization_1(conv_1);  conv_1 = None
    leaky_relu_1 = self.LeakyRelu_1(batch_normalization_1);  batch_normalization_1 = None
    max_pool_0 = self.MaxPool_0(leaky_relu_1);  leaky_relu_1 = None
    conv_2 = self.Conv_2(max_pool_0)
...

config list:

config_list = [{
        'sparsity': 0.3,
        'op_types': ['Conv2d'],
    }]

error:

File "/home/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() missing 1 required positional argument: 'input_tensor'

Environment:

NNI version:2.10
Training service (local|remote|pai|aml|etc):local
Client OS:linux
Server OS (for remote mode only):
Python version:3.8
PyTorch/TensorFlow version:pytorch
Is conda/virtualenv/venv used?:conda
Is running in Docker?:no

[Sugar929]

And the pruning code:

dummy_input = torch.randn(1, 3, 1024, 1024).to(device)
pruner = PRUNER_DICT[algorithm](model, config_list)
_, masks = pruner.compress()
pruner._unwrap_model()
ModelSpeedup(model, dummy_input, masks_file=masks).speedup_model()

[Sugar929]

Maybe I should give you the complete error:

/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:149: UserWarning: For linear and cubic interpolation in "asymmetric" and "align_corners" coordinate_transformation_moderesults might differ significantly!
  warnings.warn(
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/shape.py:44: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
  return torch.tensor(
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py:33: TracerWarning: Converting a tensor to a NumPy array might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  axes = axes.detach().cpu().numpy()
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py:36: TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
  steps = [1] * len(starts)
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py:42: TracerWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).
  for start, end, axis, step in zip(starts, ends, axes, steps):
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:72: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if sizes.nelement() != 0:
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:73: TracerWarning: Converting a tensor to a Python list might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  sizes = sizes.tolist()
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:75: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if not self.ignore_bs_ch_size and input_shape[:2] != sizes[:2]:
/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:82: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if scales.nelement() != 0:
[2023-04-03 22:13:37] start to speedup the model
[2023-04-03 22:13:40] infer module masks...
[2023-04-03 22:13:40] Update mask for Shape_0
Traceback (most recent call last):
  File "main.py", line 148, in <module>
    compress(model, config_list)
  File "/home/model_pruning/compress.py", line 37, in compress
    raise e      
  File "/home/model_pruning/compress.py", line 34, in compress
    ModelSpeedup(model, torch.randn(10, 3, 1024, 1024).to(device), masks_file=masks).speedup_model()
  File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 546, in speedup_model
    self.infer_modules_masks()
  File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 383, in infer_modules_masks
    self.update_direct_sparsity(curnode)
  File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 244, in update_direct_sparsity
    _auto_infer = AutoMaskInference(
  File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/infer_mask.py", line 80, in __init__
    self.output = self.module(*dummy_input)
  File "/home/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() missing 1 required positional argument: 'input_tensor'

When I just converted it and didn't prune it, the onnx2torch warnning was:

/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/resize.py:149: UserWarning: For linear and cubic interpolation in "asymmetric" and "align_corners" coordinate_transformation_moderesults might differ significantly!
  warnings.warn(

[J-shang]

python -m pip install --extra-index-url https://test.pypi.org/simple/ --no-cache-dir nni==3.0a1
Could you have a try with this nni version and from nni.compression.pytorch.speedup.v2 import ModelSpeedup?

[Sugar929]

pruning code:

from nni.compression.pytorch.speedup.v2 import ModelSpeedup
pruner = L1NormPruner(model, config_list)
_, masks = pruner.compress()
pruner._unwrap_model()
ModelSpeedup(model, torch.randn(1, 3, 1024, 1024).to(device), masks_or_file=masks).speedup_model()

error:

File "/home/model_pruning/compress.py", line 37, in compress
  ModelSpeedup(model, torch.randn(1, 3, 1024, 1024).to(device), masks_or_file=masks).speedup_model()
File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/v2/model_speedup.py", line 100, in __init__
  self.graph_module = graph_module if isinstance(graph_module, GraphModule) else concrete_trace(model, self.dummy_input)
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 1472, in concrete_trace
  graph = tracer.trace(root,
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 644, in trace
  fn, args, more_args, kwargs = self.create_args_for_root(fn, isinstance(root, torch.nn.Module), concrete_args)
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 506, in create_args_for_root
  raise RuntimeError(f"Tracing expected {len(arg_names)} arguments but got {len(concrete_args)} concrete arguments")
RuntimeError: Tracing expected 0 arguments but got 1 concrete arguments

[J-shang]

could you show how to build your model, so that we can find where goes wrong, seems that the forward function is not a class member function.

[Sugar929]

sorry, I directly used an onnx format from others, and I don't know how the model was built.
Then I use the deeplabv3 from pytorch to do the experiment, code:

    model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet50', pretrained=True)
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, 'net/outmodel.onnx')
    onnx_model = onnx.load('net/outmodel.onnx')
    model = convert(onnx_model)
    model.to(device)

    config_list = [{
        'sparsity': 0.3,
        'op_types': ['Conv2d'],
    }
    ]
    pruner = L1NormPruner(model, config_list)
    _, masks = pruner.compress()
    pruner._unwrap_model()
    ModelSpeedup(model, torch.randn(1, 3, 224, 224).to(device), masks_or_file=masks).speedup_model()   

error:

  File "/home/model_pruning/compress.py", line 37, in compress
    ModelSpeedup(model, torch.randn(1, 3, 224, 224).to(device), masks_or_file=masks).speedup_model()
  File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/v2/model_speedup.py", line 100, in __init__
    self.graph_module = graph_module if isinstance(graph_module, GraphModule) else concrete_trace(model, self.dummy_input)
  File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 1472, in concrete_trace
    graph = tracer.trace(root,
  File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 983, in trace
    (self.create_arg(OperatorPatcherContext.patch_run(fn, *args, *more_args, **kwargs)),),
  File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/operator_patcher.py", line 289, in patch_run
    return new_func(*args, **kwargs)
  File "<eval_with_key>.0", line 141, in forward
  File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 714, in module_call_wrapper
    return _orig_module_call(mod, *args, **kwargs)
  File "/home/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py", line 102, in forward
    return _forward()
  File "/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py", line 90, in _forward
    flip_dims, pos_axes_slices, neg_axes_slices = _get_slices(starts, ends, axes, steps)
  File "/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py", line 49, in _get_slices
    pos_axes_slices = list(slices.get(a, slice(None, None)) for a in range(max(axes) + 1))
  File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 1267, in __call__
    args = (clz(args[0]), *args[1:])
  File "/home/miniconda3/lib/python3.8/site-packages/onnx2torch/node_converters/slice.py", line 49, in <genexpr>
    pos_axes_slices = list(slices.get(a, slice(None, None)) for a in range(max(axes) + 1))
TypeError: __bool__ should return bool, returned ConcreteProxy

[liufinback]

pruning code:

from nni.compression.pytorch.speedup.v2 import ModelSpeedup
pruner = L1NormPruner(model, config_list)
_, masks = pruner.compress()
pruner._unwrap_model()
ModelSpeedup(model, torch.randn(1, 3, 1024, 1024).to(device), masks_or_file=masks).speedup_model()

error:

File "/home/model_pruning/compress.py", line 37, in compress
  ModelSpeedup(model, torch.randn(1, 3, 1024, 1024).to(device), masks_or_file=masks).speedup_model()
File "/home/miniconda3/lib/python3.8/site-packages/nni/compression/pytorch/speedup/v2/model_speedup.py", line 100, in __init__
  self.graph_module = graph_module if isinstance(graph_module, GraphModule) else concrete_trace(model, self.dummy_input)
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 1472, in concrete_trace
  graph = tracer.trace(root,
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 644, in trace
  fn, args, more_args, kwargs = self.create_args_for_root(fn, isinstance(root, torch.nn.Module), concrete_args)
File "/home/miniconda3/lib/python3.8/site-packages/nni/common/concrete_trace_utils/concrete_tracer.py", line 506, in create_args_for_root
  raise RuntimeError(f"Tracing expected {len(arg_names)} arguments but got {len(concrete_args)} concrete arguments")
RuntimeError: Tracing expected 0 arguments but got 1 concrete arguments

I have encount the same error "RuntimeError: Tracing expected 0 arguments but got 1 concrete arguments". and my error is caused by the definition of forward function. My definition is like this:

def forward_once(self, x):
    x = self.conv1(x)
    x = self.dw_conv1(x)
    x = self.blocks(x)
    x = self.conv2(x)
    x = self.linear7(x)        
    x = self.linear1(x)       
    x = x.view(x.size(0), -1)
    x = self.feature(x)
    x = F.normalize(x, p=2, dim=1)
    return x

def forward(self, *xs):
    if len(xs) == 1:
        y0 = self.forward_once(xs[0])
        return y0
    elif len(xs) == 2:
        y0 = self.forward_once(xs[0])
        y1 = self.forward_once(xs[1])
        return y0, y1
    elif len(xs) == 3:            
        y0 = self.forward_once(xs[0])
        y1 = self.forward_once(xs[1])
        y2 = self.forward_once(xs[2])
        return y0, y1, y2

Then I change the definition to this and the problem has been solved.

def forward(self, x):
    x = self.conv1(x)
    x = self.dw_conv1(x)
    x = self.blocks(x)
    x = self.conv2(x)
    x = self.linear7(x)        
    x = self.linear1(x)       
    x = x.view(x.size(0), -1)
    x = self.feature(x)
    x = F.normalize(x, p=2, dim=1)
    return x