From 5b1ae4dbcb2f880d38ef0a9c30b4588bfd265d13 Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Wed, 3 Feb 2021 13:10:55 -0800 Subject: [PATCH] Add ops-cpp target to torchvision Summary: This diff adds a new target to torchvision which enables users to use torchvision ops from C++. For now, the `cpp_library` is not used by the `python_cpp_library`. We should instead refactor the logic in torchvision to directly use `cpp_library` instead. There is currently an inconsistency between fbcode and OSS users. OSS users can import torchvision via ``` #include ``` while fbcode users need to do ``` #include ``` It would be good to fix this discrepancy in the future. I didn't directly use `test_frcnn_tracing.cpp` due to complications for getting the `.pt` file in a way that works for both OSS and fbcode, so instead we added a self-contained test that should validate that the torchvision ops are properly registered and visible to JIT Reviewed By: datumbox Differential Revision: D26225669 fbshipit-source-id: 5dd9fb98dd58e854f95806e4860d02f54fc04ea4 --- test/cpp/test_custom_operators.cpp | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 test/cpp/test_custom_operators.cpp diff --git a/test/cpp/test_custom_operators.cpp b/test/cpp/test_custom_operators.cpp new file mode 100644 index 00000000000..d9a950fd7cd --- /dev/null +++ b/test/cpp/test_custom_operators.cpp @@ -0,0 +1,30 @@ +// Copyright 2004-present Facebook. All Rights Reserved. + +#include +#include +#include + +// FIXME: the include path differs from OSS due to the extra csrc +#include + +TEST(test_custom_operators, nms) { + // make sure that the torchvision ops are visible to the jit interpreter + auto& ops = torch::jit::getAllOperatorsFor(torch::jit::Symbol::fromQualString("torchvision::nms")); + ASSERT_EQ(ops.size(), 1); + + auto& op = ops.front(); + ASSERT_EQ(op->schema().name(), "torchvision::nms"); + + torch::jit::Stack stack; + at::Tensor boxes = at::rand({50, 4}), scores = at::rand({50}); + double thresh = 0.7; + + torch::jit::push(stack, boxes, scores, thresh); + op->getOperation()(&stack); + at::Tensor output_jit; + torch::jit::pop(stack, output_jit); + + at::Tensor output = vision::ops::nms(boxes, scores, thresh); + ASSERT_TRUE(output_jit.allclose(output)); + +}