Exporting model to TorchScript fails once I load the saved torchscript file

Hello, I’m trying to do inference through a torchscript version of the Instance Segmentation model: mask_rcnn_R_50_FPN_3x without success. I’ve been following the examples found in the deployment documentation of detectron2, here.

I’ve been able to export the model to torchscript without problem. Saving the torchscript model works fine. The problem occurs when I try to load and infer the model in another place/script.

Instructions To Reproduce the 🐛 Bug:

1. Full runnable code or full changes you made: First I load the model and images:

# Imports
import torch
import torchvision
import numpy as np
import requests
from PIL import Image
from torch import Tensor
from io import BytesIO
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.export.torchscript import scripting_with_instances
from detectron2.export import TracingAdapter
from detectron2.structures import Boxes
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer

# Load the model
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
cfg.MODEL.DEVICE = 'cpu'
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
model = build_model(cfg).eval()
DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)

# Load an image
url_image = "https://vancouver.housing.ubc.ca/wp-content/uploads/2014/02/Room_detail_4BR_room_RU_1170x660.jpg"
response = requests.get(url_image)
pil_img = Image.open(BytesIO(response.content))
img = np.array(pil_img)[:,:,::-1] # BGR

If I perform inference now throught the model works fine:

out = model([{"image": torch.as_tensor(img.astype("float32").transpose(2, 0, 1))}])[0] # works

Now I try to export it to Torchscript and save it.

fields = {
    "proposal_boxes": Boxes,
    "objectness_logits": Tensor,
    "pred_boxes": Boxes,
    "scores": Tensor,
    "pred_classes": Tensor,
    "pred_masks": Tensor,
}
ts_model = scripting_with_instances(model, fields)
ts_out = ts_model.inference(
    [{"image": torch.as_tensor(img.astype("float32").transpose(2, 0, 1))}],
    do_postprocess=False
)[0] # works

torch.jit.save(ts_model, 'path_to_save_torchscript_model')

If in the same script than the one I save this torchscript model I load it and try to do inference, works:

loaded_model = torch.jit.load('path_where_the_torchscript_model_is_saved')
_ = loaded_model.to('cpu')
loaded_out = loaded_model.inference(
    [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
    do_postprocess=False
)[0] # works

The problem occurs when I try to load and infer directly the saved torchscript file in another python file.

import torch
import torchvision
import numpy as np
import requests
from PIL import Image
from io import BytesIO

# Load an image
url_image = "https://vancouver.housing.ubc.ca/wp-content/uploads/2014/02/Room_detail_4BR_room_RU_1170x660.jpg"
response = requests.get(url_image)
pil_img = Image.open(BytesIO(response.content))
img = np.array(pil_img)[:,:,::-1] # BGR

# Load Torchscript model
loaded_model = torch.jit.load('path_where_the_torchscript_model_is_saved')
_ = loaded_model.to('cpu')
loaded_out = loaded_model.inference(
    [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
    do_postprocess=False
)[0] # raises and error

The error is the following:

<ipython-input-5-e7e91d1921ad>:1: UserWarning: floor_divide is deprecated, and will be removed 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'). (Triggered internally at  ../aten/src/ATen/native/BinaryOps.cpp:601.)
  loaded_model.inference(
<ipython-input-5-e7e91d1921ad>:1: 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.)
  loaded_model.inference(
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-e7e91d1921ad> in <module>
----> 1 loaded_out = loaded_model.inference(
      2     [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
      3     do_postprocess=False
      4 )[0]

RuntimeError: classTypeINTERNAL ASSERT FAILED at "../torch/csrc/jit/python/pybind_utils.h":780, please report a bug to PyTorch.

1. What exact command you run:

loaded_out = loaded_model.inference(
    [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
    do_postprocess=False
)[0]

1. Full logs or other relevant observations:

<ipython-input-5-e7e91d1921ad>:1: UserWarning: floor_divide is deprecated, and will be removed 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'). (Triggered internally at  ../aten/src/ATen/native/BinaryOps.cpp:601.)
  loaded_model.inference(
<ipython-input-5-e7e91d1921ad>:1: 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.)
  loaded_model.inference(
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-e7e91d1921ad> in <module>
----> 1 loaded_out = loaded_model.inference(
      2     [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
      3     do_postprocess=False
      4 )[0]

RuntimeError: classTypeINTERNAL ASSERT FAILED at "../torch/csrc/jit/python/pybind_utils.h":780, please report a bug to PyTorch

Expected behavior:

I expect to be able to load and infer the model in a different script from where the torchscript model is saved.

Environment:

Provide your environment information using the following command:

----------------------  ------------------------------------------------------------------------------------------
sys.platform            linux
Python                  3.8.10 (default, Nov 26 2021, 20:14:08) [GCC 9.3.0]
numpy                   1.19.5
detectron2              0.6 @/home/user/environments/research/lib/python3.8/site-packages/detectron2
Compiler                GCC 7.3
CUDA compiler           CUDA 10.2
detectron2 arch flags   3.7, 5.0, 5.2, 6.0, 6.1, 7.0, 7.5
DETECTRON2_ENV_MODULE   <not set>
PyTorch                 1.10.1+cu102 @/home/user/environments/research/lib/python3.8/site-packages/torch
PyTorch debug build     False
GPU available           Yes
GPU 0                   NVIDIA GeForce RTX 2080 Ti (arch=7.5)
Driver version          470.103.01
CUDA_HOME               /usr
Pillow                  8.1.2
torchvision             0.11.2+cu102 @/home/user/environments/research/lib/python3.8/site-packages/torchvision
torchvision arch flags  3.5, 5.0, 6.0, 7.0, 7.5
fvcore                  0.1.5.post20220212
iopath                  0.1.9
cv2                     4.5.3
----------------------  ------------------------------------------------------------------------------------------
PyTorch built with:
  - GCC 7.3
  - C++ Version: 201402
  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.2.3 (Git Hash 7336ca9f055cf1bfa13efb658fe15dc9b41f0740)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - LAPACK is enabled (usually provided by MKL)
  - NNPACK is enabled
  - CPU capability usage: AVX2
  - CUDA Runtime 10.2
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70
  - CuDNN 7.6.5
  - Magma 2.5.2
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=10.2, CUDNN_VERSION=7.6.5, CXX_COMPILER=/opt/rh/devtoolset-7/root/usr/bin/c++, CXX_FLAGS= -Wno-deprecated -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -fopenmp -DNDEBUG -DUSE_KINETO -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -DEDGE_PROFILER_USE_KINETO -O2 -fPIC -Wno-narrowing -Wall -Wextra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.10.1, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, 

I’m doing something wrong here? I can provide more details or code if it’s necessary.

Thanks in advance!!

[paudom]

@DirkKramer Because scripting does not allow post-process you need to do it yourself. With out being <detectron2.export.torchscript_patch1.ScriptedInstances1 object at 0x00000183272BB9E8> you can still access to their attributes such as pred_masks. After a lot of searching I found the functions related to post-process. in detectron2/layers/mask_ops.py to be more exact the function _do_paste_mask, then you can do the following:

out = loaded_model.inference(
    [{'image': torch.as_tensor(img.astype('float32').transpose(2, 0, 1))}],
    do_postprocess=False
)[0]
N = len(out.pred_boxes)
chunks = torch.chunk(torch.arange(N, device='cpu'), N)
img_masks = torch.zeros(
    N, img.shape[0], img.shape[1], device='cpu', dtype=torch.bool
)
masks = out.pred_masks[:,0,:,:]  # shape was [#instances, 1, 28, 28] --> [#instances, 28, 28]
for inds in chunks:
    masks_chunk, spatial_inds = _do_paste_mask(
        out.pred_masks[inds, :, :, :], out.pred_boxes.tensor[inds], img.shape[0], img.shape[1], skip_empty=True
    )

    masks_chunk = (masks_chunk >= 0.5).to(dtype=torch.bool)
    img_masks[(inds,) + spatial_inds] = masks_chunk
# if you want images:
final = Image.fromarray(np.array(np.uint8(img_masks[instance_idx])*255)) # where instance_idx is the instance detected index.

I really don’t know if there is any way more efficient than doing all these, but at least it works.

[willcray]

@paudom @DirkKramer

Thank you for this thread! Did you notice any degradation in model accuracy following this change? The model is substantially faster for me now, but there was a small (0.3%) degradation in accuracy. Do you have any idea that might occur?