Model tracing not working

[axmav]

🐛 Bug

I am trying to trace SAMPNet and get following error:

Traceback (most recent call last):
  File "trace.py", line 38, in <module>
    traced_cell = torch.jit.trace(torch_model, (x, s))
  File "/home/alex/anaconda3/envs/torch/lib/python3.8/site-packages/torch/jit/_trace.py", line 741, in trace
    return trace_module(
  File "/home/alex/anaconda3/envs/torch/lib/python3.8/site-packages/torch/jit/_trace.py", line 958, in trace_module
    module._c._create_method_from_trace(
RuntimeError: 0INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1634272068694/work/torch/csrc/jit/ir/alias_analysis.cpp":584, please report a bug to PyTorch. We don't have an op for aten::fill_ but it isn't a special case.  Argument types: Tensor, bool, 

To Reproduce

Steps to reproduce the behavior:

1. Clone SAMPNet
2. Run code sample listed below
3. Fix first error TypeError: adaptive_max_pool1d(): argument 'output_size' (position 2) must be tuple of ints, not Tensor by providing output_size directly as constant output_size=784 in samp_module.py line 135
4. Run code sample again

from samp_net import EMDLoss, SAMPNet
import torch
from config import Config
import torch.nn.functional as F

# Load a pre-trained version of MobileNetV2
cfg = Config()
torch_model = SAMPNet(cfg,pretrained=False)
# Set the model in evaluation mode.
torch_model.eval()

x = torch.randn(2,3,224,224)
s = torch.randn(2,1,224,224)
weight, attribute, score = torch_model(x,s)
if weight is not None:
    print('weight', weight.shape, F.softmax(weight,dim=1))
if attribute is not None:
    print('attribute', attribute.shape, attribute)

print('score', score.shape, score)

traced_cell = torch.jit.trace(torch_model, (x, s))
print(traced_cell)

Expected behavior

As you can see the tensors x and s are actually able to go trough the model. I expected trace function to terminate without the error.

Environment

PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.11.0-38-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 495.29.05
cuDNN version: Probably one of the following:
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.4
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==1.10.0
[pip3] torchaudio==0.10.0
[pip3] torchvision==0.11.1
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.3.1 h2bc3f7f_2
[conda] ffmpeg 4.3 hf484d3e_0 pytorch
[conda] mkl 2021.4.0 h06a4308_640
[conda] mkl-service 2.4.0 py38h7f8727e_0
[conda] mkl_fft 1.3.1 py38hd3c417c_0
[conda] mkl_random 1.2.2 py38h51133e4_0
[conda] numpy 1.21.2 py38h20f2e39_0
[conda] numpy-base 1.21.2 py38h79a1101_0
[conda] pytorch 1.10.0 py3.8_cuda11.3_cudnn8.2.0_0 pytorch
[conda] pytorch-mutex 1.0 cuda pytorch
[conda] torchaudio 0.10.0 py38_cu113 pytorch
[conda] torchvision 0.11.1 py38_cu113 pytorch

[axmav]

More detailed error description:

RuntimeError: 0INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":611, please report a bug to PyTorch. We don't have an op for aten::fill_ but it isn't a special case.  Argument types: Tensor, bool, 

Candidates:
	aten::fill_.Scalar(Tensor(a!) self, Scalar value) -> (Tensor(a!))
	aten::fill_.Tensor(Tensor(a!) self, Tensor value) -> (Tensor(a!))

Is there a way to find out where the exactly the problem is in my code?

[axmav]

Hello! Here is the reduced code example to reproduce the bug:

import torch
import torch.nn.functional as F
import torch.nn as nn

class SurroundPattern(nn.Module):
    def __init__(self, crop_size=1./2):
        super(SurroundPattern, self).__init__()
        self.crop_size = crop_size

    def forward(self, x, s):
        H,W         = x.shape[2:]
        crop_h      = (int(H / 2 - self.crop_size / 2 * H), int(H / 2 + self.crop_size / 2 * H))
        crop_w      = (int(W / 2 - self.crop_size / 2 * W), int(W / 2 + self.crop_size / 2 * W))
        x_mask      = torch.zeros(H,W,device=x.device, dtype=torch.bool)
        x_mask[crop_h[0] : crop_h[1], crop_w[0] : crop_w[1]] = True

        inside_indices  = torch.where(x_mask)
        inside_part = x[:, :, inside_indices[0], inside_indices[1]]
        inside_feat = inside_part.mean(2)

        outside_indices = torch.where(~x_mask)
        outside_part    = x[:, :, outside_indices[0], outside_indices[1]]
        outside_feat    = outside_part.mean(2)
        fused = torch.stack([inside_feat, outside_feat], dim=2).unsqueeze(3)
        if s is None:
            return fused

        SH,SW       = s.shape[2:]
        crop_sh     = (int(SH / 2 - self.crop_size / 2 * SH), int(SH / 2 + self.crop_size / 2 * SH))
        crop_sw     = (int(SW / 2 - self.crop_size / 2 * SW), int(SW / 2 + self.crop_size / 2 * SW))
        s_mask      = torch.zeros(SH, SW, device=s.device, dtype=torch.bool)
        s_mask[crop_sh[0] : crop_sh[1], crop_sw[0] : crop_sw[1]] = True

        s_inside_indices = torch.where(s_mask)
        inside_sal  = s[:, :, s_inside_indices[0], s_inside_indices[1]].flatten(1)

        s_outside_indices = torch.where(~s_mask)
        outside_sal = s[:, :, s_outside_indices[0], s_outside_indices[1]].flatten(1)
        if outside_sal.shape != inside_sal.shape:
            outside_sal = F.adaptive_max_pool1d(outside_sal.unsqueeze(1), output_size=784)
            outside_sal = outside_sal.squeeze(1)
        fused_sal    = torch.stack([inside_sal, outside_sal], dim=2).unsqueeze(3)
        return fused, fused_sal

x = torch.randn(2, 512, 7, 7)
s = torch.randn(2, 1, 56, 56)

patt = SurroundPattern()

traced_cell = torch.jit.trace(patt, (x, s))
print(traced_cell)

[axmav]

I solved the problem by replacing Bool tensor type with Int:

x_mask= torch.zeros(H,W,device=x.device)
x_mask[crop_h[0] : crop_h[1], crop_w[0] : crop_w[1]] = 1

For more details see: https://stackoverflow.com/a/70167264/3879178
The issue is still opened because Bool support would be better.

[Ken1256]

https://github.com/cszn/SCUNet/blob/main/models/network_scunet.py
Same problem.

net = SCUNet()

x = torch.randn((1, 3, 128, 128))
o1 = net(x)
print(x.shape)

net_jit = torch.jit.trace(net, x)
o1_jit = net(x)
print(o1_jit.shape)

[Zalways]

same problem:

Traceback (most recent call last):
File "deploy/export_model00.py", line 227, in
exported_model = export_tracing(torch_model, sample_inputs)
File "deploy/export_model00.py", line 132, in export_tracing
torch.onnx.export(traceable_model, image, f, input_names=['input'], output_names=['output'], opset_version=11,dynamic_axes={'input':[1,2],'output':[1,2]})
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/init.py", line 316, in export
return utils.export(model, args, f, export_params, verbose, training,
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/utils.py", line 107, in export
_export(model, args, f, export_params, verbose, training, input_names, output_names,
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/utils.py", line 724, in _export
_model_to_graph(model, args, verbose, input_names,
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/utils.py", line 493, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args)
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/utils.py", line 437, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/onnx/utils.py", line 388, in _trace_and_get_graph_from_model
torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/jit/_trace.py", line 1166, in _get_trace_graph
outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/root/miniconda3/envs/deepsolo/lib/python3.8/site-packages/torch/jit/_trace.py", line 127, in forward
graph, out = torch._C.create_graph_by_tracing(
RuntimeError: 0INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":584, please report a bug to PyTorch. We don't have an op for aten::fill but it isn't a special case. Argument types: Tensor, bool,

[Zhaogaog]

@Zalways Same problem.Have u solved

[liaoqingPKU]

@Zalways Same problem.Have u solved

@Zalways Same problem.Have u solved
Have you solved the issue?

[Zhaogaog]

@liaoqingPKU Not yet.

[Zalways]

@liaoqingPKU Not yet.

try to export with different version of PyTorch and cuda, and also the python version. create different envs and you may find one of the envs works.