Repository: microsoft/nni

Issue:
aten::ScalarImplicit is not Supported
**Describe the issue**:



**Environment**:
- NNI version:2.10
- Training service (local|remote|pai|aml|etc):
- Client OS:
- Server OS (for remote mode only):
- Python version:
- PyTorch/TensorFlow version:
- Is conda/virtualenv/venv used?:
- Is running in Docker?:


**Configuration**:
 - Experiment config (remember to remove secrets!):
 - Search space:


**Log message**:
 - nnimanager.log:
 - dispatcher.log:
 - nnictl stdout and stderr:
 
<!--
Where can you find the log files:
LOG: https://github.com/microsoft/nni/blob/master/docs/en_US/Tutorial/HowToDebug.md#experiment-root-director
STDOUT/STDERR: https://nni.readthedocs.io/en/stable/reference/nnictl.html#nnictl-log-stdout
-->


**How to reproduce it?**:
I have an model:
```
class get_model(nn.Module):

    def __init__(self, args, num_channel=3, num_class=40, **kwargs):
        super(get_model, self).__init__()
        self.args = args
        self.bn1 = nn.BatchNorm2d(64)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.bn5 = nn.BatchNorm1d(args.emb_dims)

        self.conv1 = nn.Sequential(nn.Conv2d(num_channel*2, 64, kernel_size=1, bias=False),
                                   self.bn1,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                                   self.bn2,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                                   self.bn3,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                                   self.bn4,
                                   nn.LeakyReLU(negative_slope=0.2))
        self.conv5 = nn.Sequential(nn.Conv1d(512, args.emb_dims, kernel_size=1, bias=False),
                                   self.bn5,
                                   nn.LeakyReLU(negative_slope=0.2))

        self.linear1 = nn.Linear(args.emb_dims*2, 512, bias=False)
        self.bn6 = nn.BatchNorm1d(512)
        self.dp1 = nn.Dropout(p=args.dropout)
        self.linear2 = nn.Linear(512, 256)
        self.bn7 = nn.BatchNorm1d(256)
        self.dp2 = nn.Dropout(p=args.dropout)
        self.linear3 = nn.Linear(256, num_class)

    def forward(self, x):
        batch_size = x.size()[0]
        x = get_graph_feature(x, k=self.args.k)
        x = self.conv1(x)
        x1 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x1, k=self.args.k)
        x = self.conv2(x)
        x2 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x2, k=self.args.k)
        x = self.conv3(x)
        x3 = x.max(dim=-1, keepdim=False)[0]

        x = get_graph_feature(x3, k=self.args.k)
        x = self.conv4(x)
        x4 = x.max(dim=-1, keepdim=False)[0]

        x = torch.cat((x1, x2, x3, x4), dim=1)

        x = self.conv5(x)
        x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
        x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1)
        x = torch.cat((x1, x2), 1)

        x = F.leaky_relu(self.bn6(self.linear1(x)), negative_slope=0.2)
        x = self.dp1(x)
        x = F.leaky_relu(self.bn7(self.linear2(x)), negative_slope=0.2)
        x = self.dp2(x)
        x = self.linear3(x)
        return x
```
And try prune it and speedup:
```
prune_config_list = [{
    'sparsity_per_layer': 0.5,
    'op_types': ['Linear', 'Conv2d']
}, {
    'exclude': True,
    'op_names': ['conv1.0', 'linear3']
}]


from nni.algorithms.compression.v2.pytorch.pruning import L2NormPruner
from nni.compression.pytorch.speedup import ModelSpeedup

pruner = L2NormPruner(classifier, prune_config_list)
masked_model, masks = pruner.compress()
pruner.show_pruned_weights()

# need to unwrap the model, if the model is wrapped before speedup
pruner._unwrap_model()

# speedup the model, for more information about speedup, please refer :doc:`pruning_speedup`.


ModelSpeedup(classifier, torch.randn(1, 3, 1024, device='cuda'), masks).speedup_model()
print(classifier)
```
But get the error:
```
[2023-03-13 10:06:07] simulated prune conv2.0 remain/total: 32/64
[2023-03-13 10:06:07] simulated prune conv3.0 remain/total: 64/128
[2023-03-13 10:06:07] simulated prune conv4.0 remain/total: 128/256
[2023-03-13 10:06:07] simulated prune linear1 remain/total: 256/512
[2023-03-13 10:06:07] simulated prune linear2 remain/total: 128/256
[2023-03-13 10:06:08] start to speedup the model
[2023-03-13 10:06:08] infer module masks...
[2023-03-13 10:06:08] Update mask for .aten::size.22
[2023-03-13 10:06:08] Update mask for .aten::size.25
[2023-03-13 10:06:08] Update mask for .aten::size.30
[2023-03-13 10:06:08] Update mask for .aten::size.33
[2023-03-13 10:06:08] Update mask for .aten::Int.23
[2023-03-13 10:06:08] Update mask for .aten::Int.24
[2023-03-13 10:06:08] Update mask for .aten::Int.26
[2023-03-13 10:06:08] Update mask for .aten::Int.27
[2023-03-13 10:06:08] Update mask for .aten::ScalarImplicit.28
[2023-03-13 10:06:08] ERROR: aten::ScalarImplicit is not Supported! Please report an issue at https://github.com/microsoft/nni. Thanks~
[2023-03-13 10:06:08] Update mask for .aten::Int.29
[2023-03-13 10:06:08] Update mask for .aten::Int.31
[2023-03-13 10:06:08] Update mask for .aten::Int.32
[2023-03-13 10:06:08] Update mask for .aten::Int.34
[2023-03-13 10:06:08] Update mask for .aten::Int.35
[2023-03-13 10:06:08] Update mask for .aten::Int.36
[2023-03-13 10:06:08] Update mask for .aten::mul.57
[2023-03-13 10:06:08] Update mask for .aten::arange.50
Traceback (most recent call last):
  File "nni_optim.py", line 246, in <module>
    ModelSpeedup(classifier, torch.randn(1, 3, 1024, device='cuda'), masks).speedup_model()
  File "/home/zyfra-devbox/.conda/envs/nni/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 546, in speedup_model
    self.infer_modules_masks()
  File "/home/zyfra-devbox/.conda/envs/nni/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 383, in infer_modules_masks
    self.update_direct_sparsity(curnode)
  File "/home/zyfra-devbox/.conda/envs/nni/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 237, in update_direct_sparsity
    _auto_infer = AutoMaskInference(
  File "/home/zyfra-devbox/.conda/envs/nni/lib/python3.8/site-packages/nni/compression/pytorch/speedup/infer_mask.py", line 80, in __init__
    self.output = self.module(*dummy_input)
  File "/home/zyfra-devbox/.conda/envs/nni/lib/python3.8/site-packages/nni/compression/pytorch/speedup/jit_translate.py", line 227, in __call__
    assert len(args) >= len(self.undetermined)
AssertionError
```
I've checked several related issues:
https://github.com/microsoft/nni/issues/5097
https://github.com/microsoft/nni/issues/5090
But not find the workaround there was mentioned that  torch.full can cause the aten::ScalarImplicit. But I have no such operations.
What the workaround can be?

In addiction I've tried to prune with v2 api. I cloned and build the code from repo and tried:
```
from nni.compression.pytorch.speedup.v2 import ModelSpeedup
ModelSpeedup(classifier, torch.randn(1, 3, 1024, device='cuda'), masks).speedup_model()
```
But got:
```
[2023-03-13 08:12:57] simulated prune conv2.0 remain/total: 32/64
[2023-03-13 08:12:57] simulated prune conv3.0 remain/total: 64/128
[2023-03-13 08:12:57] simulated prune conv4.0 remain/total: 128/256
[2023-03-13 08:12:57] simulated prune linear1 remain/total: 256/512
[2023-03-13 08:12:57] simulated prune linear2 remain/total: 128/256
Traceback (most recent call last):
  File "nni_optim.py", line 248, in <module>
    ModelSpeedup(classifier, torch.randn(1, 3, 1024, device='cuda'), masks).speedup_model()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 93, in __init__
    self.graph_module = graph_module if isinstance(graph_module, GraphModule) else concrete_trace(model, self.dummy_input)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 1378, in concrete_trace
    graph = tracer.trace(root,
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 911, in trace
    (self.create_arg(OperatorPatcherContext.patch_run(fn, *args, *more_args, **kwargs)),),
  File "/root/devel/nni/nni/common/concrete_trace_utils/operator_patcher.py", line 270, in patch_run
    return new_func(*args, **kwargs)
  File "/root/devel/IAE/downstream_tasks/classification/models/dgcnn_clsft.py", line 98, in new_func
    x = F.leaky_relu(self.bn6(self.linear1(x)), negative_slope=0.2)
  File "/root/devel/nni/nni/common/concrete_trace_utils/operator_patcher.py", line 270, in patch_run
    return new_func(*args, **kwargs)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 659, in module_call_wrapper
    return self.create_proxy('call_module', module_qualified_name, args, kwargs)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 283, in create_proxy
    value_unwrapped = self.run_target(kind, target, args_unwrapped, kwargs_unwrapped)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 250, in run_target
    return OperatorPatcherContext.patch_run(mod, *args, **kwargs)
  File "/root/devel/nni/nni/common/concrete_trace_utils/operator_patcher.py", line 270, in patch_run
    return new_func(*args, **kwargs)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 639, in module_call_wrapper
    return _orig_module_call(mod, *args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/batchnorm.py", line 168, in forward
    return F.batch_norm(
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 1122, in func_wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 2280, in batch_norm
    _verify_batch_size(input.size())
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 1122, in func_wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 2248, in _verify_batch_size
    raise ValueError("Expected more than 1 value per channel when training, got input size {}".format(size))
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 512])
```
So I've tried to increase batch size. Tried different numbers but get cuda OOM even if I have batch size of 2:
```
******************************
******************************
[2023-03-13 08:14:57] simulated prune conv2.0 remain/total: 32/64
[2023-03-13 08:14:57] simulated prune conv3.0 remain/total: 64/128
[2023-03-13 08:14:57] simulated prune conv4.0 remain/total: 128/256
[2023-03-13 08:14:57] simulated prune linear1 remain/total: 256/512
[2023-03-13 08:14:57] simulated prune linear2 remain/total: 128/256
[2023-03-13 08:14:57] Start to speedup the model...
[2023-03-13 08:14:57] Resolve the mask conflict before mask propagate...
[2023-03-13 08:14:57] Infer module masks...
[2023-03-13 08:14:57] Propagate original variables
[2023-03-13 08:14:57] Propagate variables for placeholder: x
[2023-03-13 08:14:57] Propagate variables for call_method: size
[2023-03-13 08:14:57] Propagate variables for call_function: getitem
[2023-03-13 08:14:57] Propagate variables for call_method: size_1
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_1
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_2
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_3
[2023-03-13 08:14:57] Propagate variables for call_method: view
[2023-03-13 08:14:57] Propagate variables for call_method: transpose
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous
[2023-03-13 08:14:57] Propagate variables for call_function: matmul
[2023-03-13 08:14:57] Propagate variables for call_function: mul
[2023-03-13 08:14:57] Propagate variables for call_function: pow_1
[2023-03-13 08:14:57] Propagate variables for call_function: sum_1
[2023-03-13 08:14:57] Propagate variables for call_function: neg
[2023-03-13 08:14:57] Propagate variables for call_function: sub
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_1
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_1
[2023-03-13 08:14:57] Propagate variables for call_function: sub_1
[2023-03-13 08:14:57] Propagate variables for call_method: topk
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_4
[2023-03-13 08:14:57] Propagate variables for call_function: arange
[2023-03-13 08:14:57] Propagate variables for call_method: view_1
[2023-03-13 08:14:57] Propagate variables for call_function: mul_1
[2023-03-13 08:14:57] Propagate variables for call_function: iadd
[2023-03-13 08:14:57] Propagate variables for call_method: view_2
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_2
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_2
[2023-03-13 08:14:57] Propagate variables for call_function: mul_2
[2023-03-13 08:14:57] Propagate variables for call_method: view_3
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_5
[2023-03-13 08:14:57] Propagate variables for call_method: view_4
[2023-03-13 08:14:57] Propagate variables for call_method: view_5
[2023-03-13 08:14:57] Propagate variables for call_method: repeat
[2023-03-13 08:14:57] Propagate variables for call_function: sub_2
[2023-03-13 08:14:57] Propagate variables for call_function: cat
[2023-03-13 08:14:57] Propagate variables for call_method: permute
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_3
[2023-03-13 08:14:57] Propagate variables for call_module: conv1_0
[2023-03-13 08:14:57] Propagate variables for call_module: bn1
[2023-03-13 08:14:57] Propagate variables for call_module: conv1_2
[2023-03-13 08:14:57] Propagate variables for call_method: max_1
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_6
[2023-03-13 08:14:57] Propagate variables for call_method: size_2
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_7
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_8
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_9
[2023-03-13 08:14:57] Propagate variables for call_method: view_6
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_3
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_4
[2023-03-13 08:14:57] Propagate variables for call_function: matmul_1
[2023-03-13 08:14:57] Propagate variables for call_function: mul_3
[2023-03-13 08:14:57] Propagate variables for call_function: pow_2
[2023-03-13 08:14:57] Propagate variables for call_function: sum_2
[2023-03-13 08:14:57] Propagate variables for call_function: neg_1
[2023-03-13 08:14:57] Propagate variables for call_function: sub_3
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_4
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_5
[2023-03-13 08:14:57] Propagate variables for call_function: sub_4
[2023-03-13 08:14:57] Propagate variables for call_method: topk_1
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_10
[2023-03-13 08:14:57] Propagate variables for call_function: arange_1
[2023-03-13 08:14:57] Propagate variables for call_method: view_7
[2023-03-13 08:14:57] Propagate variables for call_function: mul_4
[2023-03-13 08:14:57] Propagate variables for call_function: iadd_1
[2023-03-13 08:14:57] Propagate variables for call_method: view_8
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_5
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_6
[2023-03-13 08:14:57] Propagate variables for call_function: mul_5
[2023-03-13 08:14:57] Propagate variables for call_method: view_9
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_11
[2023-03-13 08:14:57] Propagate variables for call_method: view_10
[2023-03-13 08:14:57] Propagate variables for call_method: view_11
[2023-03-13 08:14:57] Propagate variables for call_method: repeat_1
[2023-03-13 08:14:57] Propagate variables for call_function: sub_5
[2023-03-13 08:14:57] Propagate variables for call_function: cat_1
[2023-03-13 08:14:57] Propagate variables for call_method: permute_1
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_7
[2023-03-13 08:14:57] Propagate variables for call_module: conv2_0
[2023-03-13 08:14:57] Propagate variables for call_module: bn2
[2023-03-13 08:14:57] Propagate variables for call_module: conv2_2
[2023-03-13 08:14:57] Propagate variables for call_method: max_2
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_12
[2023-03-13 08:14:57] Propagate variables for call_method: size_3
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_13
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_14
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_15
[2023-03-13 08:14:57] Propagate variables for call_method: view_12
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_6
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_8
[2023-03-13 08:14:57] Propagate variables for call_function: matmul_2
[2023-03-13 08:14:57] Propagate variables for call_function: mul_6
[2023-03-13 08:14:57] Propagate variables for call_function: pow_3
[2023-03-13 08:14:57] Propagate variables for call_function: sum_3
[2023-03-13 08:14:57] Propagate variables for call_function: neg_2
[2023-03-13 08:14:57] Propagate variables for call_function: sub_6
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_7
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_9
[2023-03-13 08:14:57] Propagate variables for call_function: sub_7
[2023-03-13 08:14:57] Propagate variables for call_method: topk_2
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_16
[2023-03-13 08:14:57] Propagate variables for call_function: arange_2
[2023-03-13 08:14:57] Propagate variables for call_method: view_13
[2023-03-13 08:14:57] Propagate variables for call_function: mul_7
[2023-03-13 08:14:57] Propagate variables for call_function: iadd_2
[2023-03-13 08:14:57] Propagate variables for call_method: view_14
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_8
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_10
[2023-03-13 08:14:57] Propagate variables for call_function: mul_8
[2023-03-13 08:14:57] Propagate variables for call_method: view_15
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_17
[2023-03-13 08:14:57] Propagate variables for call_method: view_16
[2023-03-13 08:14:57] Propagate variables for call_method: view_17
[2023-03-13 08:14:57] Propagate variables for call_method: repeat_2
[2023-03-13 08:14:57] Propagate variables for call_function: sub_8
[2023-03-13 08:14:57] Propagate variables for call_function: cat_2
[2023-03-13 08:14:57] Propagate variables for call_method: permute_2
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_11
[2023-03-13 08:14:57] Propagate variables for call_module: conv3_0
[2023-03-13 08:14:57] Propagate variables for call_module: bn3
[2023-03-13 08:14:57] Propagate variables for call_module: conv3_2
[2023-03-13 08:14:57] Propagate variables for call_method: max_3
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_18
[2023-03-13 08:14:57] Propagate variables for call_method: size_4
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_19
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_20
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_21
[2023-03-13 08:14:57] Propagate variables for call_method: view_18
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_9
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_12
[2023-03-13 08:14:57] Propagate variables for call_function: matmul_3
[2023-03-13 08:14:57] Propagate variables for call_function: mul_9
[2023-03-13 08:14:57] Propagate variables for call_function: pow_4
[2023-03-13 08:14:57] Propagate variables for call_function: sum_4
[2023-03-13 08:14:57] Propagate variables for call_function: neg_3
[2023-03-13 08:14:57] Propagate variables for call_function: sub_9
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_10
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_13
[2023-03-13 08:14:57] Propagate variables for call_function: sub_10
[2023-03-13 08:14:57] Propagate variables for call_method: topk_3
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_22
[2023-03-13 08:14:57] Propagate variables for call_function: arange_3
[2023-03-13 08:14:57] Propagate variables for call_method: view_19
[2023-03-13 08:14:57] Propagate variables for call_function: mul_10
[2023-03-13 08:14:57] Propagate variables for call_function: iadd_3
[2023-03-13 08:14:57] Propagate variables for call_method: view_20
[2023-03-13 08:14:57] Propagate variables for call_method: transpose_11
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_14
[2023-03-13 08:14:57] Propagate variables for call_function: mul_11
[2023-03-13 08:14:57] Propagate variables for call_method: view_21
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_23
[2023-03-13 08:14:57] Propagate variables for call_method: view_22
[2023-03-13 08:14:57] Propagate variables for call_method: view_23
[2023-03-13 08:14:57] Propagate variables for call_method: repeat_3
[2023-03-13 08:14:57] Propagate variables for call_function: sub_11
[2023-03-13 08:14:57] Propagate variables for call_function: cat_3
[2023-03-13 08:14:57] Propagate variables for call_method: permute_3
[2023-03-13 08:14:57] Propagate variables for call_method: contiguous_15
[2023-03-13 08:14:57] Propagate variables for call_module: conv4_0
[2023-03-13 08:14:57] Propagate variables for call_module: bn4
[2023-03-13 08:14:57] Propagate variables for call_module: conv4_2
[2023-03-13 08:14:57] Propagate variables for call_method: max_4
[2023-03-13 08:14:57] Propagate variables for call_function: getitem_24
[2023-03-13 08:14:57] Propagate variables for call_function: cat_4
[2023-03-13 08:14:57] Propagate variables for call_module: conv5_0
[2023-03-13 08:14:57] Propagate variables for call_module: bn5
[2023-03-13 08:14:57] Propagate variables for call_module: conv5_2
[2023-03-13 08:14:57] Propagate variables for call_function: adaptive_max_pool1d
[2023-03-13 08:14:57] Propagate variables for call_method: view_24
[2023-03-13 08:14:57] Propagate variables for call_function: adaptive_avg_pool1d
[2023-03-13 08:14:57] Propagate variables for call_method: view_25
[2023-03-13 08:14:57] Propagate variables for call_function: cat_5
[2023-03-13 08:14:57] Propagate variables for call_module: linear1
[2023-03-13 08:14:57] Propagate variables for call_module: bn6
[2023-03-13 08:14:57] Propagate variables for call_function: leaky_relu
[2023-03-13 08:14:57] Propagate variables for call_module: dp1
[2023-03-13 08:14:57] Propagate variables for call_module: linear2
[2023-03-13 08:14:57] Propagate variables for call_module: bn7
[2023-03-13 08:14:57] Propagate variables for call_function: leaky_relu_1
[2023-03-13 08:14:57] Propagate variables for call_module: dp2
[2023-03-13 08:14:57] Propagate variables for call_module: linear3
[2023-03-13 08:14:57] Propagate variables for output: output
[2023-03-13 08:14:57] Update direct sparsity...
Traceback (most recent call last):
  File "nni_optim.py", line 248, in <module>
    ModelSpeedup(classifier, torch.randn(2, 3, 1024, device='cuda'), masks).speedup_model()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 383, in speedup_model
    self.update_direct_sparsity()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 237, in update_direct_sparsity
    self.node_infos[node].mask_updater.direct_update_preprocess(self, node)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/mask_updater.py", line 95, in direct_update_preprocess
    node_info.output_randomize = tree_map_zip(lambda t: randomize_if_tensor(t, batch_dim, batch_size), node_info.output_origin)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/utils.py", line 73, in tree_map_zip
    return tree_map(fn, pytrees[0])
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/_pytree.py", line 179, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/_pytree.py", line 179, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/mask_updater.py", line 95, in <lambda>
    node_info.output_randomize = tree_map_zip(lambda t: randomize_if_tensor(t, batch_dim, batch_size), node_info.output_origin)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/utils.py", line 55, in randomize_if_tensor
    new_obj = obj.clone().detach().contiguous()
RuntimeError: CUDA out of memory. Tried to allocate 80.00 MiB (GPU 0; 10.76 GiB total capacity; 9.16 GiB already allocated; 69.31 MiB free; 9.22 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```
It's too strange because I have the NVIDIA RTX2080 TI with 12 gb of memorry and I'm able to train the model with batch size 32.

As I have 2 gpus I tried to run on parallel:
```
from nni.compression.pytorch.speedup.v2 import ModelSpeedup
classifier = torch.nn.DataParallel(classifier, device_ids=[0, 1])
ModelSpeedup(classifier, torch.randn(1, 3, 1024, device='cuda'), masks).speedup_model()
```
But got the next error:
```
  File "nni_optim.py", line 248, in <module>
    ModelSpeedup(classifier, torch.randn(2, 3, 1024, device='cuda'), masks).speedup_model()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 93, in __init__
    self.graph_module = graph_module if isinstance(graph_module, GraphModule) else concrete_trace(model, self.dummy_input)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 1378, in concrete_trace
    graph = tracer.trace(root,
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 589, in trace
    fn, args, more_args, kwargs = self.create_args_for_root(fn, isinstance(root, torch.nn.Module), concrete_args)
  File "/root/devel/nni/nni/common/concrete_trace_utils/concrete_tracer.py", line 464, 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
```
What is the best way to deal with this problem?


UPD:
I've run the code on cpu with code from built library from source and got the next error:
```
[2023-03-13 09:55:49] simulated prune conv1.0 remain/total: 32/64
[2023-03-13 09:55:49] simulated prune conv2.0 remain/total: 32/64
[2023-03-13 09:55:49] simulated prune conv3.0 remain/total: 64/128
[2023-03-13 09:55:49] simulated prune conv4.0 remain/total: 128/256
[2023-03-13 09:55:49] simulated prune linear1 remain/total: 256/512
[2023-03-13 09:55:49] simulated prune linear2 remain/total: 128/256
[2023-03-13 09:55:50] Start to speedup the model...
[2023-03-13 09:55:50] Resolve the mask conflict before mask propagate...
[2023-03-13 09:55:50] Infer module masks...
[2023-03-13 09:55:50] Propagate original variables
[2023-03-13 09:55:50] Propagate variables for placeholder: x
[2023-03-13 09:55:50] Propagate variables for call_method: size
[2023-03-13 09:55:50] Propagate variables for call_function: getitem
[2023-03-13 09:55:50] Propagate variables for call_method: size_1
[2023-03-13 09:55:50] Propagate variables for call_function: getitem_1
[2023-03-13 09:55:50] Propagate variables for call_function: getitem_2
[2023-03-13 09:55:50] Propagate variables for call_function: getitem_3
[2023-03-13 09:55:50] Propagate variables for call_method: view
[2023-03-13 09:55:50] Propagate variables for call_method: transpose
[2023-03-13 09:55:50] Propagate variables for call_method: contiguous
[2023-03-13 09:55:50] Propagate variables for call_function: matmul
[2023-03-13 09:55:50] Propagate variables for call_function: mul
[2023-03-13 09:55:50] Propagate variables for call_function: pow_1
[2023-03-13 09:55:50] Propagate variables for call_function: sum_1
[2023-03-13 09:55:50] Propagate variables for call_function: neg
[2023-03-13 09:55:50] Propagate variables for call_function: sub
[2023-03-13 09:55:50] Propagate variables for call_method: transpose_1
[2023-03-13 09:55:50] Propagate variables for call_method: contiguous_1
[2023-03-13 09:55:50] Propagate variables for call_function: sub_1
[2023-03-13 09:55:50] Propagate variables for call_method: topk
[2023-03-13 09:55:50] Propagate variables for call_function: getitem_4
[2023-03-13 09:55:50] Propagate variables for call_function: arange
[2023-03-13 09:55:50] Propagate variables for call_method: view_1
[2023-03-13 09:55:50] Propagate variables for call_function: mul_1
[2023-03-13 09:55:50] Propagate variables for call_function: iadd
[2023-03-13 09:55:50] Propagate variables for call_method: view_2
[2023-03-13 09:55:50] Propagate variables for call_method: transpose_2
[2023-03-13 09:55:50] Propagate variables for call_method: contiguous_2
[2023-03-13 09:55:50] Propagate variables for call_function: mul_2
[2023-03-13 09:55:50] Propagate variables for call_method: view_3
[2023-03-13 09:55:50] Propagate variables for call_function: getitem_5
[2023-03-13 09:55:50] Propagate variables for call_method: view_4
[2023-03-13 09:55:50] Propagate variables for call_method: view_5
[2023-03-13 09:55:50] Propagate variables for call_method: repeat
[2023-03-13 09:55:50] Propagate variables for call_function: sub_2
[2023-03-13 09:55:50] Propagate variables for call_function: cat
[2023-03-13 09:55:50] Propagate variables for call_method: permute
[2023-03-13 09:55:50] Propagate variables for call_method: contiguous_3
[2023-03-13 09:55:50] Propagate variables for call_module: conv1_0
[2023-03-13 09:55:51] Propagate variables for call_module: bn1
[2023-03-13 09:55:51] Propagate variables for call_module: conv1_2
[2023-03-13 09:55:51] Propagate variables for call_method: max_1
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_6
[2023-03-13 09:55:51] Propagate variables for call_method: size_2
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_7
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_8
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_9
[2023-03-13 09:55:51] Propagate variables for call_method: view_6
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_3
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_4
[2023-03-13 09:55:51] Propagate variables for call_function: matmul_1
[2023-03-13 09:55:51] Propagate variables for call_function: mul_3
[2023-03-13 09:55:51] Propagate variables for call_function: pow_2
[2023-03-13 09:55:51] Propagate variables for call_function: sum_2
[2023-03-13 09:55:51] Propagate variables for call_function: neg_1
[2023-03-13 09:55:51] Propagate variables for call_function: sub_3
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_4
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_5
[2023-03-13 09:55:51] Propagate variables for call_function: sub_4
[2023-03-13 09:55:51] Propagate variables for call_method: topk_1
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_10
[2023-03-13 09:55:51] Propagate variables for call_function: arange_1
[2023-03-13 09:55:51] Propagate variables for call_method: view_7
[2023-03-13 09:55:51] Propagate variables for call_function: mul_4
[2023-03-13 09:55:51] Propagate variables for call_function: iadd_1
[2023-03-13 09:55:51] Propagate variables for call_method: view_8
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_5
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_6
[2023-03-13 09:55:51] Propagate variables for call_function: mul_5
[2023-03-13 09:55:51] Propagate variables for call_method: view_9
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_11
[2023-03-13 09:55:51] Propagate variables for call_method: view_10
[2023-03-13 09:55:51] Propagate variables for call_method: view_11
[2023-03-13 09:55:51] Propagate variables for call_method: repeat_1
[2023-03-13 09:55:51] Propagate variables for call_function: sub_5
[2023-03-13 09:55:51] Propagate variables for call_function: cat_1
[2023-03-13 09:55:51] Propagate variables for call_method: permute_1
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_7
[2023-03-13 09:55:51] Propagate variables for call_module: conv2_0
[2023-03-13 09:55:51] Propagate variables for call_module: bn2
[2023-03-13 09:55:51] Propagate variables for call_module: conv2_2
[2023-03-13 09:55:51] Propagate variables for call_method: max_2
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_12
[2023-03-13 09:55:51] Propagate variables for call_method: size_3
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_13
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_14
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_15
[2023-03-13 09:55:51] Propagate variables for call_method: view_12
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_6
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_8
[2023-03-13 09:55:51] Propagate variables for call_function: matmul_2
[2023-03-13 09:55:51] Propagate variables for call_function: mul_6
[2023-03-13 09:55:51] Propagate variables for call_function: pow_3
[2023-03-13 09:55:51] Propagate variables for call_function: sum_3
[2023-03-13 09:55:51] Propagate variables for call_function: neg_2
[2023-03-13 09:55:51] Propagate variables for call_function: sub_6
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_7
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_9
[2023-03-13 09:55:51] Propagate variables for call_function: sub_7
[2023-03-13 09:55:51] Propagate variables for call_method: topk_2
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_16
[2023-03-13 09:55:51] Propagate variables for call_function: arange_2
[2023-03-13 09:55:51] Propagate variables for call_method: view_13
[2023-03-13 09:55:51] Propagate variables for call_function: mul_7
[2023-03-13 09:55:51] Propagate variables for call_function: iadd_2
[2023-03-13 09:55:51] Propagate variables for call_method: view_14
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_8
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_10
[2023-03-13 09:55:51] Propagate variables for call_function: mul_8
[2023-03-13 09:55:51] Propagate variables for call_method: view_15
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_17
[2023-03-13 09:55:51] Propagate variables for call_method: view_16
[2023-03-13 09:55:51] Propagate variables for call_method: view_17
[2023-03-13 09:55:51] Propagate variables for call_method: repeat_2
[2023-03-13 09:55:51] Propagate variables for call_function: sub_8
[2023-03-13 09:55:51] Propagate variables for call_function: cat_2
[2023-03-13 09:55:51] Propagate variables for call_method: permute_2
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_11
[2023-03-13 09:55:51] Propagate variables for call_module: conv3_0
[2023-03-13 09:55:51] Propagate variables for call_module: bn3
[2023-03-13 09:55:51] Propagate variables for call_module: conv3_2
[2023-03-13 09:55:51] Propagate variables for call_method: max_3
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_18
[2023-03-13 09:55:51] Propagate variables for call_method: size_4
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_19
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_20
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_21
[2023-03-13 09:55:51] Propagate variables for call_method: view_18
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_9
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_12
[2023-03-13 09:55:51] Propagate variables for call_function: matmul_3
[2023-03-13 09:55:51] Propagate variables for call_function: mul_9
[2023-03-13 09:55:51] Propagate variables for call_function: pow_4
[2023-03-13 09:55:51] Propagate variables for call_function: sum_4
[2023-03-13 09:55:51] Propagate variables for call_function: neg_3
[2023-03-13 09:55:51] Propagate variables for call_function: sub_9
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_10
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_13
[2023-03-13 09:55:51] Propagate variables for call_function: sub_10
[2023-03-13 09:55:51] Propagate variables for call_method: topk_3
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_22
[2023-03-13 09:55:51] Propagate variables for call_function: arange_3
[2023-03-13 09:55:51] Propagate variables for call_method: view_19
[2023-03-13 09:55:51] Propagate variables for call_function: mul_10
[2023-03-13 09:55:51] Propagate variables for call_function: iadd_3
[2023-03-13 09:55:51] Propagate variables for call_method: view_20
[2023-03-13 09:55:51] Propagate variables for call_method: transpose_11
[2023-03-13 09:55:51] Propagate variables for call_method: contiguous_14
[2023-03-13 09:55:51] Propagate variables for call_function: mul_11
[2023-03-13 09:55:51] Propagate variables for call_method: view_21
[2023-03-13 09:55:51] Propagate variables for call_function: getitem_23
[2023-03-13 09:55:51] Propagate variables for call_method: view_22
[2023-03-13 09:55:51] Propagate variables for call_method: view_23
[2023-03-13 09:55:51] Propagate variables for call_method: repeat_3
[2023-03-13 09:55:52] Propagate variables for call_function: sub_11
[2023-03-13 09:55:52] Propagate variables for call_function: cat_3
[2023-03-13 09:55:52] Propagate variables for call_method: permute_3
[2023-03-13 09:55:52] Propagate variables for call_method: contiguous_15
[2023-03-13 09:55:52] Propagate variables for call_module: conv4_0
[2023-03-13 09:55:52] Propagate variables for call_module: bn4
[2023-03-13 09:55:52] Propagate variables for call_module: conv4_2
[2023-03-13 09:55:52] Propagate variables for call_method: max_4
[2023-03-13 09:55:52] Propagate variables for call_function: getitem_24
[2023-03-13 09:55:52] Propagate variables for call_function: cat_4
[2023-03-13 09:55:52] Propagate variables for call_module: conv5_0
[2023-03-13 09:55:52] Propagate variables for call_module: bn5
[2023-03-13 09:55:52] Propagate variables for call_module: conv5_2
[2023-03-13 09:55:52] Propagate variables for call_function: adaptive_max_pool1d
[2023-03-13 09:55:52] Propagate variables for call_method: view_24
[2023-03-13 09:55:52] Propagate variables for call_function: adaptive_avg_pool1d
[2023-03-13 09:55:52] Propagate variables for call_method: view_25
[2023-03-13 09:55:52] Propagate variables for call_function: cat_5
[2023-03-13 09:55:52] Propagate variables for call_module: linear1
[2023-03-13 09:55:52] Propagate variables for call_module: bn6
[2023-03-13 09:55:52] Propagate variables for call_function: leaky_relu
[2023-03-13 09:55:52] Propagate variables for call_module: dp1
[2023-03-13 09:55:52] Propagate variables for call_module: linear2
[2023-03-13 09:55:52] Propagate variables for call_module: bn7
[2023-03-13 09:55:52] Propagate variables for call_function: leaky_relu_1
[2023-03-13 09:55:52] Propagate variables for call_module: dp2
[2023-03-13 09:55:52] Propagate variables for call_module: linear3
[2023-03-13 09:55:52] Propagate variables for output: output
[2023-03-13 09:55:52] Update direct sparsity...
[2023-03-13 09:55:56] Update direct mask for placeholder: x
[2023-03-13 09:55:56] Update direct mask for call_method: size
[2023-03-13 09:55:56] Update direct mask for call_function: getitem
[2023-03-13 09:55:56] Update direct mask for call_method: size_1
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_1
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_2
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_3
[2023-03-13 09:55:56] Update direct mask for call_method: view
[2023-03-13 09:55:56] Update direct mask for call_method: transpose
[2023-03-13 09:55:56] Update direct mask for call_method: contiguous
[2023-03-13 09:55:56] Update direct mask for call_function: matmul
[2023-03-13 09:55:56] Update direct mask for call_function: mul
[2023-03-13 09:55:56] Update direct mask for call_function: pow_1
[2023-03-13 09:55:56] Update direct mask for call_function: sum_1
[2023-03-13 09:55:56] Update direct mask for call_function: neg
[2023-03-13 09:55:56] Update direct mask for call_function: sub
[2023-03-13 09:55:56] Update direct mask for call_method: transpose_1
[2023-03-13 09:55:56] Update direct mask for call_method: contiguous_1
[2023-03-13 09:55:56] Update direct mask for call_function: sub_1
[2023-03-13 09:55:56] Update direct mask for call_method: topk
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_4
[2023-03-13 09:55:56] Update direct mask for call_function: arange
[2023-03-13 09:55:56] Update direct mask for call_method: view_1
[2023-03-13 09:55:56] Update direct mask for call_function: mul_1
[2023-03-13 09:55:56] Update direct mask for call_function: iadd
[2023-03-13 09:55:56] Update direct mask for call_method: view_2
[2023-03-13 09:55:56] Update direct mask for call_method: transpose_2
[2023-03-13 09:55:56] Update direct mask for call_method: contiguous_2
[2023-03-13 09:55:56] Update direct mask for call_function: mul_2
[2023-03-13 09:55:56] Update direct mask for call_method: view_3
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_5
[2023-03-13 09:55:56] Update direct mask for call_method: view_4
[2023-03-13 09:55:56] Update direct mask for call_method: view_5
[2023-03-13 09:55:56] Update direct mask for call_method: repeat
[2023-03-13 09:55:56] Update direct mask for call_function: sub_2
[2023-03-13 09:55:56] Update direct mask for call_function: cat
[2023-03-13 09:55:56] Update direct mask for call_method: permute
[2023-03-13 09:55:56] Update direct mask for call_method: contiguous_3
[2023-03-13 09:55:56] Update direct mask for call_module: conv1_0
[2023-03-13 09:55:56] Update direct mask for call_module: bn1
[2023-03-13 09:55:56] Update direct mask for call_module: conv1_2
[2023-03-13 09:55:56] Update direct mask for call_method: max_1
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_6
[2023-03-13 09:55:56] Update direct mask for call_method: size_2
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_7
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_8
[2023-03-13 09:55:56] Update direct mask for call_function: getitem_9
[2023-03-13 09:55:56] Update direct mask for call_method: view_6
[2023-03-13 09:55:56] Update direct mask for call_method: transpose_3
[2023-03-13 09:55:56] Update direct mask for call_method: contiguous_4
[2023-03-13 09:55:57] Update direct mask for call_function: matmul_1
[2023-03-13 09:55:57] Update direct mask for call_function: mul_3
[2023-03-13 09:55:57] Update direct mask for call_function: pow_2
[2023-03-13 09:55:57] Update direct mask for call_function: sum_2
[2023-03-13 09:55:57] Update direct mask for call_function: neg_1
[2023-03-13 09:55:57] Update direct mask for call_function: sub_3
[2023-03-13 09:55:57] Update direct mask for call_method: transpose_4
[2023-03-13 09:55:57] Update direct mask for call_method: contiguous_5
[2023-03-13 09:55:57] Update direct mask for call_function: sub_4
[2023-03-13 09:55:57] Update direct mask for call_method: topk_1
[2023-03-13 09:55:57] Update direct mask for call_function: getitem_10
[2023-03-13 09:55:57] Update direct mask for call_function: arange_1
[2023-03-13 09:55:57] Update direct mask for call_method: view_7
[2023-03-13 09:55:57] Update direct mask for call_function: mul_4
[2023-03-13 09:55:57] Update direct mask for call_function: iadd_1
[2023-03-13 09:55:57] Update direct mask for call_method: view_8
[2023-03-13 09:55:57] Update direct mask for call_method: transpose_5
[2023-03-13 09:55:57] Update direct mask for call_method: contiguous_6
[2023-03-13 09:55:57] Update direct mask for call_function: mul_5
[2023-03-13 09:55:57] Update direct mask for call_method: view_9
[2023-03-13 09:55:57] Update direct mask for call_function: getitem_11
[2023-03-13 09:55:57] Update direct mask for call_method: view_10
[2023-03-13 09:55:57] Update direct mask for call_method: view_11
[2023-03-13 09:55:57] Update direct mask for call_method: repeat_1
[2023-03-13 09:55:57] Update direct mask for call_function: sub_5
[2023-03-13 09:55:57] Update direct mask for call_function: cat_1
[2023-03-13 09:55:58] Update direct mask for call_method: permute_1
[2023-03-13 09:55:58] Update direct mask for call_method: contiguous_7
[2023-03-13 09:55:58] Update direct mask for call_module: conv2_0
[2023-03-13 09:55:58] Update direct mask for call_module: bn2
[2023-03-13 09:55:58] Update direct mask for call_module: conv2_2
[2023-03-13 09:55:58] Update direct mask for call_method: max_2
[2023-03-13 09:55:58] Update direct mask for call_function: getitem_12
[2023-03-13 09:55:58] Update direct mask for call_method: size_3
[2023-03-13 09:55:58] Update direct mask for call_function: getitem_13
[2023-03-13 09:55:58] Update direct mask for call_function: getitem_14
[2023-03-13 09:55:58] Update direct mask for call_function: getitem_15
[2023-03-13 09:55:58] Update direct mask for call_method: view_12
[2023-03-13 09:55:58] Update direct mask for call_method: transpose_6
[2023-03-13 09:55:58] Update direct mask for call_method: contiguous_8
[2023-03-13 09:55:58] Update direct mask for call_function: matmul_2
[2023-03-13 09:55:59] Update direct mask for call_function: mul_6
[2023-03-13 09:55:59] Update direct mask for call_function: pow_3
[2023-03-13 09:55:59] Update direct mask for call_function: sum_3
[2023-03-13 09:55:59] Update direct mask for call_function: neg_2
[2023-03-13 09:55:59] Update direct mask for call_function: sub_6
[2023-03-13 09:55:59] Update direct mask for call_method: transpose_7
[2023-03-13 09:55:59] Update direct mask for call_method: contiguous_9
[2023-03-13 09:55:59] Update direct mask for call_function: sub_7
[2023-03-13 09:55:59] Update direct mask for call_method: topk_2
[2023-03-13 09:55:59] Update direct mask for call_function: getitem_16
[2023-03-13 09:55:59] Update direct mask for call_function: arange_2
[2023-03-13 09:55:59] Update direct mask for call_method: view_13
[2023-03-13 09:55:59] Update direct mask for call_function: mul_7
[2023-03-13 09:55:59] Update direct mask for call_function: iadd_2
[2023-03-13 09:55:59] Update direct mask for call_method: view_14
[2023-03-13 09:55:59] Update direct mask for call_method: transpose_8
[2023-03-13 09:55:59] Update direct mask for call_method: contiguous_10
[2023-03-13 09:55:59] Update direct mask for call_function: mul_8
[2023-03-13 09:55:59] Update direct mask for call_method: view_15
[2023-03-13 09:55:59] Update direct mask for call_function: getitem_17
[2023-03-13 09:55:59] Update direct mask for call_method: view_16
[2023-03-13 09:55:59] Update direct mask for call_method: view_17
[2023-03-13 09:55:59] Update direct mask for call_method: repeat_2
[2023-03-13 09:55:59] Update direct mask for call_function: sub_8
[2023-03-13 09:55:59] Update direct mask for call_function: cat_2
[2023-03-13 09:56:00] Update direct mask for call_method: permute_2
[2023-03-13 09:56:00] Update direct mask for call_method: contiguous_11
[2023-03-13 09:56:00] Update direct mask for call_module: conv3_0
[2023-03-13 09:56:00] Update direct mask for call_module: bn3
[2023-03-13 09:56:00] Update direct mask for call_module: conv3_2
[2023-03-13 09:56:00] Update direct mask for call_method: max_3
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_18
[2023-03-13 09:56:01] Update direct mask for call_method: size_4
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_19
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_20
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_21
[2023-03-13 09:56:01] Update direct mask for call_method: view_18
[2023-03-13 09:56:01] Update direct mask for call_method: transpose_9
[2023-03-13 09:56:01] Update direct mask for call_method: contiguous_12
[2023-03-13 09:56:01] Update direct mask for call_function: matmul_3
[2023-03-13 09:56:01] Update direct mask for call_function: mul_9
[2023-03-13 09:56:01] Update direct mask for call_function: pow_4
[2023-03-13 09:56:01] Update direct mask for call_function: sum_4
[2023-03-13 09:56:01] Update direct mask for call_function: neg_3
[2023-03-13 09:56:01] Update direct mask for call_function: sub_9
[2023-03-13 09:56:01] Update direct mask for call_method: transpose_10
[2023-03-13 09:56:01] Update direct mask for call_method: contiguous_13
[2023-03-13 09:56:01] Update direct mask for call_function: sub_10
[2023-03-13 09:56:01] Update direct mask for call_method: topk_3
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_22
[2023-03-13 09:56:01] Update direct mask for call_function: arange_3
[2023-03-13 09:56:01] Update direct mask for call_method: view_19
[2023-03-13 09:56:01] Update direct mask for call_function: mul_10
[2023-03-13 09:56:01] Update direct mask for call_function: iadd_3
[2023-03-13 09:56:01] Update direct mask for call_method: view_20
[2023-03-13 09:56:01] Update direct mask for call_method: transpose_11
[2023-03-13 09:56:01] Update direct mask for call_method: contiguous_14
[2023-03-13 09:56:01] Update direct mask for call_function: mul_11
[2023-03-13 09:56:01] Update direct mask for call_method: view_21
[2023-03-13 09:56:01] Update direct mask for call_function: getitem_23
[2023-03-13 09:56:01] Update direct mask for call_method: view_22
[2023-03-13 09:56:01] Update direct mask for call_method: view_23
[2023-03-13 09:56:01] Update direct mask for call_method: repeat_3
[2023-03-13 09:56:02] Update direct mask for call_function: sub_11
[2023-03-13 09:56:02] Update direct mask for call_function: cat_3
[2023-03-13 09:56:02] Update direct mask for call_method: permute_3
[2023-03-13 09:56:03] Update direct mask for call_method: contiguous_15
[2023-03-13 09:56:04] Update direct mask for call_module: conv4_0
[2023-03-13 09:56:04] Update direct mask for call_module: bn4
[2023-03-13 09:56:04] Update direct mask for call_module: conv4_2
[2023-03-13 09:56:04] Update direct mask for call_method: max_4
[2023-03-13 09:56:04] Update direct mask for call_function: getitem_24
[2023-03-13 09:56:04] Update direct mask for call_function: cat_4
[2023-03-13 09:56:04] Update direct mask for call_module: conv5_0
[2023-03-13 09:56:05] Update direct mask for call_module: bn5
[2023-03-13 09:56:05] Update direct mask for call_module: conv5_2
[2023-03-13 09:56:05] Update direct mask for call_function: adaptive_max_pool1d
[2023-03-13 09:56:05] Update direct mask for call_method: view_24
[2023-03-13 09:56:05] Update direct mask for call_function: adaptive_avg_pool1d
[2023-03-13 09:56:05] Update direct mask for call_method: view_25
[2023-03-13 09:56:05] Update direct mask for call_function: cat_5
[2023-03-13 09:56:05] Update direct mask for call_module: linear1
[2023-03-13 09:56:05] Update direct mask for call_module: bn6
[2023-03-13 09:56:05] Update direct mask for call_function: leaky_relu
[2023-03-13 09:56:05] Update direct mask for call_module: dp1
[2023-03-13 09:56:05] Update direct mask for call_module: linear2
[2023-03-13 09:56:05] Update direct mask for call_module: bn7
[2023-03-13 09:56:05] Update direct mask for call_function: leaky_relu_1
[2023-03-13 09:56:05] Update direct mask for call_module: dp2
[2023-03-13 09:56:05] Update direct mask for call_module: linear3
[2023-03-13 09:56:05] Update direct mask for output: output
[2023-03-13 09:56:05] Update indirect sparsity...
[2023-03-13 09:56:05] Update indirect mask for output: output
[2023-03-13 09:56:05] Update indirect mask for call_module: linear3
[2023-03-13 09:56:05] Update indirect mask for call_module: dp2
[2023-03-13 09:56:05] Update indirect mask for call_function: leaky_relu_1
[2023-03-13 09:56:05] Update indirect mask for call_module: bn7
[2023-03-13 09:56:05] Update indirect mask for call_module: linear2
[2023-03-13 09:56:05] Update indirect mask for call_module: dp1
[2023-03-13 09:56:05] Update indirect mask for call_function: leaky_relu
[2023-03-13 09:56:05] Update indirect mask for call_module: bn6
[2023-03-13 09:56:05] Update indirect mask for call_module: linear1
[2023-03-13 09:56:05] Update indirect mask for call_function: cat_5
[2023-03-13 09:56:05] Update indirect mask for call_method: view_25
[2023-03-13 09:56:05] Update indirect mask for call_function: adaptive_avg_pool1d
[2023-03-13 09:56:05] Update indirect mask for call_method: view_24
[2023-03-13 09:56:05] Update indirect mask for call_function: adaptive_max_pool1d
[2023-03-13 09:56:05] Update indirect mask for call_module: conv5_2
[2023-03-13 09:56:05] Update indirect mask for call_module: bn5
[2023-03-13 09:56:05] Update indirect mask for call_module: conv5_0
[2023-03-13 09:56:05] Update indirect mask for call_function: cat_4
[2023-03-13 09:56:05] Update indirect mask for call_function: getitem_24
[2023-03-13 09:56:05] Update indirect mask for call_method: max_4
[2023-03-13 09:56:05] Update indirect mask for call_module: conv4_2
Traceback (most recent call last):
  File "nni_optim.py", line 246, in <module>
    ModelSpeedup(classifier, torch.randn(2, 3, 1024, device='cpu'), masks).speedup_model()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 384, in speedup_model
    self.update_indirect_sparsity()
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/model_speedup.py", line 259, in update_indirect_sparsity
    self.node_infos[node].mask_updater.indirect_update_process(self, node)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/mask_updater.py", line 438, in indirect_update_process
    indirect_fn(model_speedup, node)
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/mask_updater.py", line 382, in indirect_activation
    input_grad = tree_map_zip(lambda t, m: (t * m).type_as(t) if isinstance(m, torch.Tensor) else t, \
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/utils.py", line 81, in tree_map_zip
    return tree_unflatten([fn(*args) for args in zip(*flat_args_list)], spec_list[0])
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/utils.py", line 81, in <listcomp>
    return tree_unflatten([fn(*args) for args in zip(*flat_args_list)], spec_list[0])
  File "/root/devel/nni/nni/compression/pytorch/speedup/v2/mask_updater.py", line 382, in <lambda>
    input_grad = tree_map_zip(lambda t, m: (t * m).type_as(t) if isinstance(m, torch.Tensor) else t, \
TypeError: unsupported operand type(s) for *: 'NoneType' and 'Tensor'
```



Discussion:
I ran into the previous error as well, but the last one seems to be solvable by replacing the lambda with a regular function
Thanks for tips. Do you mean that lambda should not be in the code of the model?  Or  one should change lambda to function in code of library in module mask_updater.py":
```
        input_grad = tree_map_zip(lambda t, m: (t * m).type_as(t) if isinstance(m, torch.Tensor) else t, \
            model_speedup.node_infos[node].output_grad, model_speedup.node_infos[node].output_masks)
```
> Thanks for tips. Do you mean that lambda should not be in the code of the model? Or one should change lambda to function in code of library in module mask_updater.py":
> 
> ```
>         input_grad = tree_map_zip(lambda t, m: (t * m).type_as(t) if isinstance(m, torch.Tensor) else t, \
>             model_speedup.node_infos[node].output_grad, model_speedup.node_infos[node].output_masks)
> ```

I think that lambda should not be in the code of the model.
Ok. thanks. Double checked and I have no lambdas in my code, so the issue is in another place
Hi, may I know what is the function `get_graph_feature()`?
Sure, I forgot to include:
```
def knn(x, k):
    inner = -2 * torch.matmul(x.transpose(2, 1).contiguous(), x)
    xx = torch.sum(x ** 2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1).contiguous()
    idx = pairwise_distance.topk(k=k, dim=-1)[1]
    return idx


def get_graph_feature(x, k=20, idx=None, extra_dim=False):

    batch_size, num_dims, num_points = x.size()
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        if extra_dim is False:
            idx = knn(x, k=k)
        else:
            idx = knn(x[:, 6:], k=k)  # idx = knn(x[:, :3], k=k)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
    idx += idx_base
    idx = idx.view(-1)

    x = x.transpose(2, 1).contiguous()
    feature = x.view(batch_size*num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims)
    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
    feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2).contiguous()

    return feature  # (batch_size, 2 * num_dims, num_points, k)
```
I tried this model recently. The latest nni speedup works well on CPU. However, the output model was still unable to run because of incorrect shapes.

For the 5 conv_bn blocks, you should write it this way because it helps nni detect the bn module better.

```python
self.conv1 = nn.Sequential(nn.Conv2d(num_channel*2, 64, kernel_size=1, bias=False),
                           nn.BatchNorm2d(64),
                           nn.LeakyReLU(negative_slope=0.2))
self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
                           nn.BatchNorm2d(64),
                           nn.LeakyReLU(negative_slope=0.2))
self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False),
                           nn.BatchNorm2d(128),
                           nn.LeakyReLU(negative_slope=0.2))
self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False),
                           nn.BatchNorm2d(256),
                           nn.LeakyReLU(negative_slope=0.2))
self.conv5 = nn.Sequential(nn.Conv1d(512, args.emb_dims, kernel_size=1, bias=False),
                           nn.BatchNorm1d(args.emb_dims),
                           nn.LeakyReLU(negative_slope=0.2))
```

However, since we haven't got a replacer for ``Conv1d``, it is hard to prune this structure.

```python
x = torch.cat((x1, x2, x3, x4), dim=1)
x = self.conv5(x)
```

Let me discuss this with @J-shang.

Resolving patch:
diff --git a/nni/compression/pytorch/speedup/compress_modules.py b/nni/compression/pytorch/speedup/compress_modules.py
--- a/nni/compression/pytorch/speedup/compress_modules.py
+++ b/nni/compression/pytorch/speedup/compress_modules.py
@@ -14,6 +14,7 @@
     'BatchNorm1d': lambda module, masks: replace_batchnorm1d(module, masks),
     'InstanceNorm2d': lambda module, masks: replace_instancenorm2d(module, masks),
     'Conv2d': lambda module, masks: replace_conv2d(module, masks),
+    'Conv1d': lambda module, masks: replace_conv1d(module, masks),
     'Linear': lambda module, masks: replace_linear(module, masks),
     'MaxPool2d': lambda module, masks: no_replace(module, masks),
     'AvgPool2d': lambda module, masks: no_replace(module, masks),
@@ -57,6 +58,7 @@
 def convert_to_coarse_mask(t_mask, dim):
     """
     Convert the mask tensor to the coarse-grained mask tensor.
+
     Parameters
     ---------
     t_mask: torch.Tensor
@@ -576,6 +578,139 @@ def replace_conv2d(conv, masks):
     return new_conv
 
 
+def replace_conv1d(conv, masks):
+    """
+    Replace the original conv with a new one according to the infered
+    masks, the function support the fine-grained sparsity and coarse-grained
+    sparsity. In the fine-grained scenario, this replace function will replace
+    the filters that happen to be totally coverd by the fine-grained sparsity.
+
+    Parameters
+    ----------
+    conv : torch.nn.Conv1d
+        The conv1d module to be replaced
+    masks : Tuple of the input masks, output masks and weight masks
+        Tuple of the masks, for example
+        ([input_m1, input_m2], [output_m], {'weight':weight_m})
+
+    Returns
+    -------
+    torch.nn.Conv1d
+        The new conv1d module
+    """
+    in_masks, output_mask, weight_masks = masks
+    assert isinstance(conv, torch.nn.Conv1d)
+    if len(in_masks) != 1:
+        raise InputsNumberError()
+
+    in_mask = in_masks[0]
+    weight_mask = weight_masks['weight']
+    pruned_in, remained_in = convert_to_coarse_mask(in_mask, 1)
+    pruned_out, remained_out = convert_to_coarse_mask(output_mask, 1)
+
+    n_remained_in = weight_mask.size(1) * conv.groups - pruned_in.size(0)
+    n_remained_out = weight_mask.size(0) - pruned_out.size(0)
+
+    if n_remained_in != remained_in.size(0) or n_remained_out != remained_out.size(0):
+        raise ShapeMisMatchError()
+
+    k_size = conv.kernel_size[0]
+    # Note: We should resolve the group dependency of the conv layers before
+    # run into here.
+    # check if the mask tensor meets the group dependency and calculate the
+    # new number of the groups after pruning
+    # the original step size of the input channel for each group
+    ori_inchannel_step = int(conv.in_channels / conv.groups)
+    # the original step size of the output channel for each group
+    ori_outchannel_step = int(conv.out_channels / conv.groups)
+    # calculate the new_in_channel_step and new_outchannel_step first
+    new_inchannel_step = new_outchannel_step = None
+    for groupid in range(conv.groups):
+        in_start = groupid * ori_inchannel_step
+        in_end = in_start + ori_inchannel_step
+        out_start = groupid * ori_outchannel_step
+        out_end = out_start + ori_outchannel_step
+
+        new_inchannel_step: int = torch.logical_and(
+            in_start <= remained_in,
+            remained_in < in_end
+        ).sum().item()
+        new_outchannel_step: int = torch.logical_and(
+            out_start <= remained_out,
+            remained_out < out_end
+        ).sum().item()
+        # remap the global index to the group index
+        if new_inchannel_step == 0:
+            # if the whole group are pruned
+            continue
+        else:
+            break
+    tmp_weight = torch.ones(
+        n_remained_out, new_inchannel_step, k_size)
+    tmp_weight = tmp_weight.to(conv.weight.device)
+    if new_inchannel_step == 0 or new_outchannel_step == 0:
+        raise EmptyLayerError()
+    if n_remained_in % new_inchannel_step != 0 or n_remained_out % new_outchannel_step != 0:
+        raise UnBalancedGroupError()
+
+    new_groups = 0
+    for groupid in range(conv.groups):
+        in_start = groupid * ori_inchannel_step
+        in_end = in_start + ori_inchannel_step
+        out_start = groupid * ori_outchannel_step
+        out_end = out_start + ori_outchannel_step
+
+        current_input_mask = torch.logical_and(in_start <= remained_in, remained_in < in_end)
+        current_input_index = remained_in[current_input_mask]
+
+        current_output_mask = torch.logical_and(out_start <= remained_out, remained_out < out_end)
+        current_output_index = remained_out[current_output_mask]
+
+        # remap the global index to the group index
+        current_input_index = current_input_index - in_start
+        if len(current_input_index) == 0:
+            # if the whole group are pruned
+            assert len(current_output_index) == 0
+            continue
+        # check if the number of remained channel of each group are the same
+        if len(current_input_index) != new_inchannel_step or len(current_output_index) != new_outchannel_step:
+            raise UnBalancedGroupError()
+
+        # copy the weight into tmp_weight
+        new_out_start = new_outchannel_step * new_groups
+        new_out_end = new_out_start + new_outchannel_step
+        tmp_weight[new_out_start:new_out_end] = torch.index_select(
+            conv.weight[current_output_index], 1, torch.as_tensor(current_input_index, dtype=torch.long).to(conv.weight.device))
+        new_groups += 1
+
+    _logger.info("replace conv1d with in_channels: %d, out_channels: %d",
+                  n_remained_in, n_remained_out)
+
+    # need_bias is a flag that indicates that if a conv layer need
+    # bias, if the original conv doesn't have a bias and there is
+    # no constant need to be folded into the bias, the need_bias is False.
+    need_bias = conv.bias is not None
+    new_conv = torch.nn.Conv1d(in_channels=n_remained_in,
+                               out_channels=n_remained_out,
+                               kernel_size=conv.kernel_size,
+                               stride=conv.stride,
+                               padding=conv.padding,
+                               dilation=conv.dilation,
+                               groups=new_groups,
+                               bias=need_bias,
+                               padding_mode=conv.padding_mode)
+
+    new_conv.to(conv.weight.device)
+    new_conv.weight.copy_(tmp_weight)
+
+    # copy the bias data
+    if conv.bias is not None:
+        new_conv.bias.data.copy_(torch.index_select(
+            conv.bias.data, 0, remained_out))
+
+    return new_conv
+
+
 def replace_convtranspose2d(convtrans, masks):
     """
     We need anothor replace function for


Test patch: