`Comparison exception: The values for attribute 'dtype' do not match: torch.float64 != torch.int64.` when pruning YoloV5-face

[vinhtq115]

Describe the issue: I'm trying to prune the pre-trained model yolov5n-0.5 from Yolov5-face. Here is the code I used:

from models.experimental import attempt_load
from nni.algorithms.compression.pytorch.pruning import LevelPruner
from nni.compression.pytorch import ModelSpeedup
import torch

if __name__ == '__main__':
    model = attempt_load('weights/yolov5n-0.5.pt', map_location=torch.device('cpu'))  # FP32
    model.eval()

    config_list = [{
        'sparsity': 0.5,
        'op_types': ['default'],
    }]

    pruner = LevelPruner(model, config_list)
    model = pruner.compress()

    pruner.export_model(model_path='pruned_yolov5n-0.5.pth', mask_path='mask_yolov5n-0.5.pth')

    m_speedup = ModelSpeedup(model, dummy_input=torch.rand((1, 3, 384, 640)), masks_file='mask_yolov5n-0.5.pth')
    m_speedup.speedup_model()

But it always throws this error:

...
ERROR: Tensor-valued Constant nodes differed in value across invocations. This often indicates that the tracer has encountered untraceable code.
    Node:
        %1693 : Tensor = prim::Constant[value={0.5}](), scope: __module.model.21 # /home/vinhtq115/PycharmProjects/yolov5-face/models/yolo.py:71:0
    Source Location:
        /home/vinhtq115/PycharmProjects/yolov5-face/models/yolo.py(71): forward
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/nn/modules/module.py(1090): _slow_forward
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl
        /home/vinhtq115/PycharmProjects/yolov5-face/models/yolo.py(167): forward_once
        /home/vinhtq115/PycharmProjects/yolov5-face/models/yolo.py(151): forward
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/nn/modules/module.py(1090): _slow_forward
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/jit/_trace.py(958): trace_module
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/jit/_trace.py(741): trace
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/common/graph_utils.py(78): _trace
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/common/graph_utils.py(66): __init__
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/common/graph_utils.py(252): __init__
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/common/graph_utils.py(24): build_module_graph
        /home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py(57): __init__
        prune2.py(21): <module>
    Comparison exception:   The values for attribute 'dtype' do not match: torch.float64 != torch.int64.

Full output: out.txt

I'm not sure what is wrong here because the input and weights are torch.float32, not torch.float64.

Environment:

• NNI version: 2.6
• Training service (local|remote|pai|aml|etc): N/A (using pretrained weights)
• Client OS: Ubuntu 20.04.3
• Server OS (for remote mode only): N/A
• Python version: 3.8.12
• PyTorch/TensorFlow version: PyTorch 1.10.0
• Is conda/virtualenv/venv used?: Conda
• Is running in Docker?: N/A

[vinhtq115]

Seems like adding model.model[-1].export = True fixed the issue with dtype.

from models.experimental import attempt_load
from nni.algorithms.compression.pytorch.pruning import LevelPruner
from nni.compression.pytorch import ModelSpeedup
import torch
import torch.nn as nn
import models
from utils.activations import Hardswish, SiLU

if __name__ == '__main__':
    model = attempt_load('weights/yolov5n-0.5.pt', map_location=torch.device('cpu'))  # FP32
    model.eval()

    # Update model
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        # elif isinstance(m, models.yolo.Detect):
        #     m.forward = m.forward_export  # assign forward (optional)
        if isinstance(m, models.common.ShuffleV2Block):  # shufflenet block nn.SiLU
            for i in range(len(m.branch1)):
                if isinstance(m.branch1[i], nn.SiLU):
                    m.branch1[i] = SiLU()
            for i in range(len(m.branch2)):
                if isinstance(m.branch2[i], nn.SiLU):
                    m.branch2[i] = SiLU()
    model.model[-1].export = True  # set Detect() layer export=True

    config_list = [{
        'sparsity': 0.5,
        'op_types': ['default'],
    }]

    pruner = LevelPruner(model, config_list)
    model = pruner.compress()

    pruner.export_model(model_path='pruned_yolov5n-0.5.pth', mask_path='mask_yolov5n-0.5.pth',
                        input_shape=[1, 3, 384, 640])

    m_speedup = ModelSpeedup(model, dummy_input=torch.rand((1, 3, 384, 640)),
                             masks_file='mask_yolov5n-0.5.pth')
    m_speedup.speedup_model()

but now it shows this error:

/home/vinhtq115/miniconda3/envs/tensorrt/bin/python /home/vinhtq115/PycharmProjects/yolov5-face/prune2.py
Fusing layers... 
/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
[2022-02-21 10:36:33] INFO (utils.torch_utils/MainThread) Model Summary: 305 layers, 446376 parameters, 0 gradients, 1.5 GFLOPS
[2022-02-21 10:36:33] INFO (nni.compression.pytorch.compressor/MainThread) Model state_dict saved to pruned_yolov5n-0.5.pth
[2022-02-21 10:36:33] INFO (nni.compression.pytorch.compressor/MainThread) Mask dict saved to mask_yolov5n-0.5.pth
/home/vinhtq115/PycharmProjects/yolov5-face/models/common.py:23: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  channels_per_group = num_channels // groups
[2022-02-21 10:36:36] INFO (nni.compression.pytorch.speedup.compressor/MainThread) start to speed up the model
[2022-02-21 10:36:39] INFO (FixMaskConflict/MainThread) {'model.0.stem_1.conv.module': 1, 'model.0.stem_2a.conv.module': 1, 'model.0.stem_2b.conv.module': 1, 'model.0.stem_3.conv.module': 1, 'model.1.branch1.0.module': 1, 'model.1.branch1.2.module': 1, 'model.1.branch2.0.module': 1, 'model.1.branch2.3.module': 1, 'model.1.branch2.5.module': 1, 'model.2.0.branch2.0.module': 1, 'model.2.0.branch2.3.module': 1, 'model.2.0.branch2.5.module': 1, 'model.2.1.branch2.0.module': 1, 'model.2.1.branch2.3.module': 1, 'model.2.1.branch2.5.module': 1, 'model.2.2.branch2.0.module': 1, 'model.2.2.branch2.3.module': 1, 'model.2.2.branch2.5.module': 1, 'model.3.branch1.0.module': 1, 'model.3.branch1.2.module': 1, 'model.3.branch2.0.module': 1, 'model.3.branch2.3.module': 1, 'model.3.branch2.5.module': 1, 'model.4.0.branch2.0.module': 1, 'model.4.0.branch2.3.module': 1, 'model.4.0.branch2.5.module': 1, 'model.4.1.branch2.0.module': 1, 'model.4.1.branch2.3.module': 1, 'model.4.1.branch2.5.module': 1, 'model.4.2.branch2.0.module': 1, 'model.4.2.branch2.3.module': 1, 'model.4.2.branch2.5.module': 1, 'model.4.3.branch2.0.module': 1, 'model.4.3.branch2.3.module': 1, 'model.4.3.branch2.5.module': 1, 'model.4.4.branch2.0.module': 1, 'model.4.4.branch2.3.module': 1, 'model.4.4.branch2.5.module': 1, 'model.4.5.branch2.0.module': 1, 'model.4.5.branch2.3.module': 1, 'model.4.5.branch2.5.module': 1, 'model.4.6.branch2.0.module': 1, 'model.4.6.branch2.3.module': 1, 'model.4.6.branch2.5.module': 1, 'model.5.branch1.0.module': 1, 'model.5.branch1.2.module': 1, 'model.5.branch2.0.module': 1, 'model.5.branch2.3.module': 1, 'model.5.branch2.5.module': 1, 'model.6.0.branch2.0.module': 1, 'model.6.0.branch2.3.module': 1, 'model.6.0.branch2.5.module': 1, 'model.6.1.branch2.0.module': 1, 'model.6.1.branch2.3.module': 1, 'model.6.1.branch2.5.module': 1, 'model.6.2.branch2.0.module': 1, 'model.6.2.branch2.3.module': 1, 'model.6.2.branch2.5.module': 1, 'model.7.conv.module': 1, 'model.10.cv1.conv.module': 1, 'model.10.m.0.cv1.conv.module': 1, 'model.10.m.0.cv2.conv.module': 1, 'model.10.cv2.conv.module': 1, 'model.10.cv3.conv.module': 1, 'model.11.conv.module': 1, 'model.14.cv1.conv.module': 1, 'model.14.m.0.cv1.conv.module': 1, 'model.14.m.0.cv2.conv.module': 1, 'model.14.cv2.conv.module': 1, 'model.14.cv3.conv.module': 1, 'model.15.conv.module': 1, 'model.17.cv1.conv.module': 1, 'model.17.m.0.cv1.conv.module': 1, 'model.17.m.0.cv2.conv.module': 1, 'model.17.cv2.conv.module': 1, 'model.17.cv3.conv.module': 1, 'model.18.conv.module': 1, 'model.20.cv1.conv.module': 1, 'model.20.m.0.cv1.conv.module': 1, 'model.20.m.0.cv2.conv.module': 1, 'model.20.cv2.conv.module': 1, 'model.20.cv3.conv.module': 1, 'model.21.m.0.module': 1, 'model.21.m.1.module': 1, 'model.21.m.2.module': 1}
[2022-02-21 10:36:39] WARNING (FixMaskConflict/MainThread) no multi-dimension masks found.
[2022-02-21 10:36:39] INFO (FixMaskConflict/MainThread) Dectected conv prune dim" 0
[2022-02-21 10:36:39] INFO (nni.compression.pytorch.speedup.compressor/MainThread) infer module masks...
[2022-02-21 10:36:39] INFO (nni.compression.pytorch.speedup.compressor/MainThread) Update mask for model.0.stem_1.conv.module
[2022-02-21 10:36:40] INFO (nni.compression.pytorch.speedup.compressor/MainThread) Update mask for model.1.aten::size.212
Traceback (most recent call last):
  File "/home/vinhtq115/PycharmProjects/yolov5-face/prune2.py", line 46, in <module>
    m_speedup.speedup_model()
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 509, in speedup_model
    self.infer_modules_masks()
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 355, in infer_modules_masks
    self.update_direct_sparsity(curnode)
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 216, in update_direct_sparsity
    _auto_infer = AutoMaskInference(
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/infer_mask.py", line 80, in __init__
    self.output = self.module(*dummy_input)
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() missing 1 required positional argument: 'x'

[vinhtq115]

I tried generating config with not_safe_to_prune which seems to work better but it generated AssertionError (similar to #4160 this time.

from models.experimental import attempt_load
from nni.algorithms.compression.pytorch.pruning import LevelPruner
from nni.compression.pytorch import ModelSpeedup, apply_compression_results
import torch
import torch.nn as nn
import models
from utils.activations import Hardswish, SiLU
from nni.compression.pytorch.utils import not_safe_to_prune

if __name__ == '__main__':
    model = attempt_load('weights/yolov5n-0.5.pt', map_location=torch.device('cpu'))  # FP32
    model.eval()

    # Update model
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        # elif isinstance(m, models.yolo.Detect):
        #     m.forward = m.forward_export  # assign forward (optional)
        if isinstance(m, models.common.ShuffleV2Block):  # shufflenet block nn.SiLU
            for i in range(len(m.branch1)):
                if isinstance(m.branch1[i], nn.SiLU):
                    m.branch1[i] = SiLU()
            for i in range(len(m.branch2)):
                if isinstance(m.branch2[i], nn.SiLU):
                    m.branch2[i] = SiLU()
    model.model[-1].export = True  # set Detect() layer export=True

    dummy_input = torch.rand(1, 3, 384, 640)
    model(dummy_input)

    not_safe = not_safe_to_prune(model, dummy_input)
    cfg_list = []
    for name, module in model.named_modules():
        if name in not_safe:
            continue
        if isinstance(module, torch.nn.Conv2d):
            cfg_list.append({'op_types': ['Conv2d'], 'sparsity': 0.5, 'op_names': [name]})

    pruner = LevelPruner(model, cfg_list)
    pruner.compress()
    pruner.export_model(model_path='/tmp/ramdisk/pruned_yolov5n-0.5.pth', mask_path='/tmp/ramdisk/mask_yolov5n-0.5.pth')
    pruner._unwrap_model()

    apply_compression_results(model, '/tmp/ramdisk/mask_yolov5n-0.5.pth', 'cpu')

    m_speedup = ModelSpeedup(model, dummy_input=dummy_input, masks_file='/tmp/ramdisk/mask_yolov5n-0.5.pth')
    m_speedup.speedup_model()

Output:

[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) Layers model.20.cv2.conv using fine-grained pruning
[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) Layers model.20.cv3.conv using fine-grained pruning
[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) dim0 sparsity: 0.000000
[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) dim1 sparsity: 0.000601
[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) Dectected conv prune dim" 1
[2022-02-21 11:38:02] INFO (FixMaskConflict/MainThread) Fine-grianed mask detected
Traceback (most recent call last):
  File "/home/vinhtq115/PycharmProjects/yolov5-face/prune2.py", line 53, in <module>
    m_speedup.speedup_model()
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/speedup/compressor.py", line 506, in speedup_model
    fix_mask_conflict(self.masks, self.bound_model, self.dummy_input)
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/utils/mask_conflict.py", line 54, in fix_mask_conflict
    masks = fix_channel_mask.fix_mask()
  File "/home/vinhtq115/miniconda3/envs/tensorrt/lib/python3.8/site-packages/nni/compression/pytorch/utils/mask_conflict.py", line 263, in fix_mask
    assert len(set(num_channels_list)) == 1
AssertionError

[J-shang]

hello @vinhtq115 , LevelPruner generates fine-grained masks, and does not recommend speed up. speed up is usually combined with coarse-grained masks.

TypeError: forward() missing 1 required positional argument: 'x' this error is usually caused by model containing a branch structure, and normal inference cannot touch the corresponding branch. You could find which layer this forward function belongs to, and if this layer can be removed before use speed up, then add back after speed up.

[scarlett2018]

hello @vinhtq115 , LevelPruner generates fine-grained masks, and does not recommend speed up. speed up is usually combined with coarse-grained masks.

TypeError: forward() missing 1 required positional argument: 'x' this error is usually caused by model containing a branch structure, and normal inference cannot touch the corresponding branch. You could find which layer this forward function belongs to, and if this layer can be removed before use speed up, then add back after speed up.

@vinhtq115 - could you please report back how things going with the suggestions? we are closing resolved issues and wondering whether this issue is resolved or not. thanks.

[vinhtq115]

hello @vinhtq115 , LevelPruner generates fine-grained masks, and does not recommend speed up. speed up is usually combined with coarse-grained masks. TypeError: forward() missing 1 required positional argument: 'x' this error is usually caused by model containing a branch structure, and normal inference cannot touch the corresponding branch. You could find which layer this forward function belongs to, and if this layer can be removed before use speed up, then add back after speed up.

@vinhtq115 - could you please report back how things going with the suggestions? we are closing resolved issues and wondering whether this issue is resolved or not. thanks.

I tried L1FilterPruner and still have the same issue. Seems like it's model-related issue so I will close the issue.