vision: Cannot load scripted Keypoint-RCNN model in C++

🐛 Bug

Hello!

I am trying to load a Keypoint-RCNN model into C++ via TorchScript scripting. Scripting in python itself works ok (with no errors), however loading to C++ throws an error:

terminate called after throwing an instance of 'torch::jit::ErrorReport'                                                                                                                                   [36/1871]
  what():
Unknown builtin op: torchvision::nms.
Could not find any similar ops to torchvision::nms. This op may not exist or may not be currently supported in TorchScript.
:
  File "/home/d.lysukhin/distr/anaconda3/envs/nightly/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 39
        by NMS, sorted in decreasing order of scores
    """
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
           ~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 91
    scores: Tensor,
    iou_threshold: float) -> Tensor:
  _42 = ops.torchvision.nms(boxes, scores, iou_threshold)
        ~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _42
'nms' is being compiled since it was called from 'batched_nms'
  File "/home/d.lysukhin/distr/anaconda3/envs/nightly/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 85
        offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
        boxes_for_nms = boxes + offsets[:, None]
        keep = nms(boxes_for_nms, scores, iou_threshold)
               ~~~ <--- HERE
        return keep
Serialized   File "code/__torch__/torchvision/ops/boxes.py", line 50
_18 = torch.slice(offsets, 0, 0, 9223372036854775807, 1)                                                                                                                                               [12/1871]
    boxes_for_nms = torch.add(boxes, torch.unsqueeze(_18, 1), alpha=1)
    keep = __torch__.torchvision.ops.boxes.nms(boxes_for_nms, scores, iou_threshold, )
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _11 = keep
  return _11
'batched_nms' is being compiled since it was called from 'RegionProposalNetwork.filter_proposals'
Serialized   File "code/__torch__/torchvision/models/detection/rpn.py", line 64
    _11 = __torch__.torchvision.ops.boxes.clip_boxes_to_image
    _12 = __torch__.torchvision.ops.boxes.remove_small_boxes
    _13 = __torch__.torchvision.ops.boxes.batched_nms
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    num_images = (torch.size(proposals))[0]
    device = ops.prim.device(proposals)
'RegionProposalNetwork.filter_proposals' is being compiled since it was called from 'RegionProposalNetwork.forward'
  File "/home/d.lysukhin/distr/anaconda3/envs/nightly/lib/python3.8/site-packages/torchvision/models/detection/rpn.py", line 493
        proposals = self.box_coder.decode(pred_bbox_deltas.detach(), anchors)
        proposals = proposals.view(num_images, -1, 4)
        boxes, scores = self.filter_proposals(proposals, objectness, images.image_sizes, num_anchors_per_level)
                        ~~~~~~~~~~~~~~~~~~~~~ <--- HERE

        losses = {}
Serialized   File "code/__torch__/torchvision/models/detection/rpn.py", line 37
    proposals = (self.box_coder).decode(torch.detach(pred_bbox_deltas0), anchors, )
    proposals0 = torch.view(proposals, [num_images, -1, 4])
    _8 = (self).filter_proposals(proposals0, objectness0, images.image_sizes, num_anchors_per_level, )
                                                                              ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    boxes, scores, = _8
    losses = annotate(Dict[str, Tensor], {})

Aborted

And yes, I am using nightly builds (torch ver. 1.7.0.dev20200929, torchvision ver. 0.8.0.dev20200929) (as offered here: https://discuss.pytorch.org/t/torchvision-ops-nms-in-torchscript/89286/2).

To Reproduce

Steps to reproduce the behavior:

1. Load and run scripting of torchvision kprcnn implementation:

import torch
import torchvision
from torchvision.models.detection.keypoint_rcnn import keypointrcnn_resnet50_fpn

print("torch ver.", torch.__version__) # yields 'torch ver. 1.7.0.dev20200929'
print("torchvision ver.", torchvision.__version__) # yields 'torchvision ver. 0.8.0.dev20200929'

def main():
    model = keypointrcnn_resnet50_fpn()
    scripted_module = torch.jit.script(model)
    
    with open("kprcnn.pt", "wb") as fp:
        torch.jit.save(scripted_module, fp)
        
if __name__ == "__main__":
    main()

1. Write C++ example (from here as well: https://pytorch.org/tutorials/advanced/cpp_export.html):

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::cout << "ok\n";
}

1. Write CMakeLists.txt as proposed in https://pytorch.org/tutorials/advanced/cpp_export.html:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(torchscript-example)

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

add_executable(torchscript-example torchscript-example.cpp)
target_link_libraries(torchscript-example "${TORCH_LIBRARIES}")
set_property(TARGET torchscript-example PROPERTY CXX_STANDARD 14)

1. Build

mkdir build; cd build; cmake -DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` ..; cmake --build . --config Release

1. Run ./torchscript-example ../kpcnn.pt

Environment

• PyTorch / torchvision Version (e.g., 1.0 / 0.4.0): 1.7.0.dev20200929 / 0.8.0.dev20200929
• OS (e.g., Linux): Ubuntu 18.04
• How you installed PyTorch / torchvision (conda, pip, source): conda
• Build command you used (if compiling from source):
• Python version: 3.8.5
• CUDA/cuDNN version: 10.2 / 8.0.4.30
• GPU models and configuration: 2080 Ti
• Any other relevant information:

Additional context

Same behaviour is met at stable release. Changing model to torchvision.models.resnet50() work just fine.