Added rois shape check in C++ (#2794)

* Added rois shape check in C++

* Fixes code formatting

* Remove accidental include

* - Updated code according to the review
- Replaced old AT_ASSERT/ERROR by new TORCH_CHECK

torchvision/csrc/DeformConv.h

@@ -32,7 +32,7 @@ at::Tensor DeformConv2d_forward(
         groups,
         offset_groups);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return DeformConv2d_forward_cpu(
@@ -72,7 +72,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> DeformConv2d_backward
         groups,
         offset_groups);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return DeformConv2d_backward_cpu(

torchvision/csrc/PSROIAlign.h

@@ -28,7 +28,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward(
         pooled_width,
         sampling_ratio);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return PSROIAlign_forward_cpu(
@@ -62,7 +62,7 @@ at::Tensor PSROIAlign_backward(
         height,
         width);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return PSROIAlign_backward_cpu(

torchvision/csrc/PSROIPool.h

@@ -20,7 +20,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIPool_forward(
     return PSROIPool_forward_cuda(
         input, rois, spatial_scale, pooled_height, pooled_width);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return PSROIPool_forward_cpu(
@@ -52,7 +52,7 @@ at::Tensor PSROIPool_backward(
         height,
         width);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return PSROIPool_backward_cpu(

torchvision/csrc/ROIPool.h

@@ -20,7 +20,7 @@ std::tuple<at::Tensor, at::Tensor> ROIPool_forward(
     return ROIPool_forward_cuda(
         input, rois, spatial_scale, pooled_height, pooled_width);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return ROIPool_forward_cpu(
@@ -52,7 +52,7 @@ at::Tensor ROIPool_backward(
         height,
         width);
 #else
-    AT_ERROR("Not compiled with GPU support");
+    TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
   }
   return ROIPool_backward_cpu(

torchvision/csrc/cpu/PSROIAlign_cpu.cpp

@@ -308,8 +308,10 @@ std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cpu(
     const int pooled_width,
     const int sampling_ratio) {
   // Check if input tensors are CPU tensors
-  AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -321,7 +323,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cpu(
   int height = input.size(2);
   int width = input.size(3);
 
-  AT_ASSERTM(
+  TORCH_CHECK(
       channels % (pooled_height * pooled_width) == 0,
       "input channels must be a multiple of pooling height * pooling width");
   int channels_out = channels / (pooled_height * pooled_width);
@@ -370,9 +372,9 @@ at::Tensor PSROIAlign_backward_cpu(
     const int height,
     const int width) {
   // Check if input tensors are CPU tensors
-  AT_ASSERTM(grad.device().is_cpu(), "grad must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
-  AT_ASSERTM(
+  TORCH_CHECK(grad.device().is_cpu(), "grad must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(
       channel_mapping.device().is_cpu(),
       "channel_mapping must be a CPU tensor");
 

torchvision/csrc/cpu/PSROIPool_cpu.cpp

@@ -150,8 +150,10 @@ std::tuple<at::Tensor, at::Tensor> PSROIPool_forward_cpu(
     const int pooled_height,
     const int pooled_width) {
   // Check if input tensors are CPU tensors
-  AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -163,7 +165,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIPool_forward_cpu(
   int height = input.size(2);
   int width = input.size(3);
 
-  AT_ASSERTM(
+  TORCH_CHECK(
       channels % (pooled_height * pooled_width) == 0,
       "input channels must be a multiple of pooling height * pooling width");
   int channels_out = channels / (pooled_height * pooled_width);
@@ -210,9 +212,9 @@ at::Tensor PSROIPool_backward_cpu(
     const int height,
     const int width) {
   // Check if input tensors are CPU tensors
-  AT_ASSERTM(grad.device().is_cpu(), "grad must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
-  AT_ASSERTM(
+  TORCH_CHECK(grad.device().is_cpu(), "grad must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(
       channel_mapping.device().is_cpu(),
       "channel_mapping must be a CPU tensor");
 

torchvision/csrc/cpu/ROIAlign_cpu.cpp

@@ -394,8 +394,9 @@ at::Tensor ROIAlign_forward_cpu(
     const int64_t pooled_width,
     const int64_t sampling_ratio,
     const bool aligned) {
-  AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -447,8 +448,8 @@ at::Tensor ROIAlign_backward_cpu(
     const int64_t width,
     const int64_t sampling_ratio,
     const bool aligned) {
-  AT_ASSERTM(grad.device().is_cpu(), "grad must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(grad.device().is_cpu(), "grad must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
 
   at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2};
 

torchvision/csrc/cpu/ROIPool_cpu.cpp

@@ -126,8 +126,8 @@ std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cpu(
     const float spatial_scale,
     const int pooled_height,
     const int pooled_width) {
-  AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -180,9 +180,11 @@ at::Tensor ROIPool_backward_cpu(
     const int height,
     const int width) {
   // Check if input tensors are CPU tensors
-  AT_ASSERTM(grad.device().is_cpu(), "grad must be a CPU tensor");
-  AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
-  AT_ASSERTM(argmax.device().is_cpu(), "argmax must be a CPU tensor");
+  TORCH_CHECK(grad.device().is_cpu(), "grad must be a CPU tensor");
+  TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
+  TORCH_CHECK(argmax.device().is_cpu(), "argmax must be a CPU tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2};
 

torchvision/csrc/cpu/image/readpng_cpu.cpp

@@ -7,7 +7,7 @@
 
 #if !PNG_FOUND
 torch::Tensor decodePNG(const torch::Tensor& data) {
-  AT_ERROR("decodePNG: torchvision not compiled with libPNG support");
+  TORCH_CHECK(false, "decodePNG: torchvision not compiled with libPNG support");
 }
 #else
 #include <png.h>

torchvision/csrc/cpu/nms_cpu.cpp

@@ -5,9 +5,9 @@ at::Tensor nms_cpu_kernel(
     const at::Tensor& dets,
     const at::Tensor& scores,
     const double iou_threshold) {
-  AT_ASSERTM(!dets.is_cuda(), "dets must be a CPU tensor");
-  AT_ASSERTM(!scores.is_cuda(), "scores must be a CPU tensor");
-  AT_ASSERTM(
+  TORCH_CHECK(!dets.is_cuda(), "dets must be a CPU tensor");
+  TORCH_CHECK(!scores.is_cuda(), "scores must be a CPU tensor");
+  TORCH_CHECK(
       dets.scalar_type() == scores.scalar_type(),
       "dets should have the same type as scores");
 

torchvision/csrc/cpu/video/Video.cpp

@@ -66,7 +66,8 @@ std::pair<std::string, ffmpeg::MediaType> const* _parse_type(
   if (device != types.end()) {
     return device;
   }
-  AT_ERROR("Expected one of [audio, video, subtitle, cc] ", stream_string);
+  TORCH_CHECK(
+      false, "Expected one of [audio, video, subtitle, cc] ", stream_string);
 }
 
 std::string parse_type_to_string(const std::string& stream_string) {
@@ -97,7 +98,8 @@ std::tuple<std::string, long> _parseStream(const std::string& streamString) {
     try {
       index_ = c10::stoi(match[2].str());
     } catch (const std::exception&) {
-      AT_ERROR(
+      TORCH_CHECK(
+          false,
           "Could not parse device index '",
           match[2].str(),
           "' in device string '",

torchvision/csrc/cuda/PSROIAlign_cuda.cu

@@ -302,8 +302,10 @@ std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cuda(
     const int pooled_width,
     const int sampling_ratio) {
   // Check if input tensors are CUDA tensors
-  AT_ASSERTM(input.is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -318,7 +320,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cuda(
   auto height = input.size(2);
   auto width = input.size(3);
 
-  AT_ASSERTM(
+  TORCH_CHECK(
       channels % (pooled_height * pooled_width) == 0,
       "input channels must be a multiple of pooling height * pooling width");
   int channels_out = channels / (pooled_height * pooled_width);
@@ -378,9 +380,9 @@ at::Tensor PSROIAlign_backward_cuda(
     const int height,
     const int width) {
   // Check if input tensors are CUDA tensors
-  AT_ASSERTM(grad.is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
-  AT_ASSERTM(
+  TORCH_CHECK(grad.is_cuda(), "grad must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
       channel_mapping.is_cuda(),
       "channel_mapping must be a CUDA tensor");
 

torchvision/csrc/cuda/PSROIPool_cuda.cu

@@ -139,8 +139,10 @@ std::tuple<at::Tensor, at::Tensor> PSROIPool_forward_cuda(
     const int pooled_height,
     const int pooled_width) {
   // Check if input tensors are CUDA tensors
-  AT_ASSERTM(input.is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -155,7 +157,7 @@ std::tuple<at::Tensor, at::Tensor> PSROIPool_forward_cuda(
   auto height = input.size(2);
   auto width = input.size(3);
 
-  AT_ASSERTM(
+  TORCH_CHECK(
       channels % (pooled_height * pooled_width) == 0,
       "input channels must be a multiple of pooling height * pooling width");
   int channels_out = channels / (pooled_height * pooled_width);
@@ -212,9 +214,9 @@ at::Tensor PSROIPool_backward_cuda(
     const int height,
     const int width) {
   // Check if input tensors are CUDA tensors
-  AT_ASSERTM(grad.is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
-  AT_ASSERTM(
+  TORCH_CHECK(grad.is_cuda(), "grad must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
       channel_mapping.is_cuda(),
       "channel_mapping must be a CUDA tensor");
 

torchvision/csrc/cuda/ROIAlign_cuda.cu

@@ -320,8 +320,10 @@ at::Tensor ROIAlign_forward_cuda(
     const int64_t pooled_width,
     const int64_t sampling_ratio,
     const bool aligned) {
-  AT_ASSERTM(input.is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "rois must have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -385,8 +387,8 @@ at::Tensor ROIAlign_backward_cuda(
     const int64_t width,
     const int64_t sampling_ratio,
     const bool aligned) {
-  AT_ASSERTM(grad.is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(grad.is_cuda(), "grad must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
 
   at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2};
 

torchvision/csrc/cuda/ROIPool_cuda.cu

@@ -121,8 +121,10 @@ std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cuda(
     const float spatial_scale,
     const int pooled_height,
     const int pooled_width) {
-  AT_ASSERTM(input.is_cuda(), "input must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(
+      rois.size(1) == 5, "Tensor rois should have shape as Tensor[K, 5]");
 
   at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
 
@@ -188,9 +190,9 @@ at::Tensor ROIPool_backward_cuda(
     const int height,
     const int width) {
   // Check if input tensors are CUDA tensors
-  AT_ASSERTM(grad.is_cuda(), "grad must be a CUDA tensor");
-  AT_ASSERTM(rois.is_cuda(), "rois must be a CUDA tensor");
-  AT_ASSERTM(argmax.is_cuda(), "argmax must be a CUDA tensor");
+  TORCH_CHECK(grad.is_cuda(), "grad must be a CUDA tensor");
+  TORCH_CHECK(rois.is_cuda(), "rois must be a CUDA tensor");
+  TORCH_CHECK(argmax.is_cuda(), "argmax must be a CUDA tensor");
 
   at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2},
       argmax_t{argmax, "argmax", 3};

torchvision/csrc/cuda/nms_cuda.cu

@@ -71,8 +71,8 @@ __global__ void nms_kernel(
 at::Tensor nms_cuda(const at::Tensor& dets,
     const at::Tensor& scores,
     const double iou_threshold) {
-  AT_ASSERTM(dets.is_cuda(), "dets must be a CUDA tensor");
-  AT_ASSERTM(scores.is_cuda(), "scores must be a CUDA tensor");
+  TORCH_CHECK(dets.is_cuda(), "dets must be a CUDA tensor");
+  TORCH_CHECK(scores.is_cuda(), "scores must be a CUDA tensor");
 
   TORCH_CHECK(
       dets.dim() == 2, "boxes should be a 2d tensor, got ", dets.dim(), "D");
@@ -96,7 +96,7 @@ at::Tensor nms_cuda(const at::Tensor& dets,
 #if defined(WITH_CUDA) || defined(WITH_HIP)
   at::cuda::CUDAGuard device_guard(dets.device());
 #else
-  AT_ERROR("Not compiled with GPU support");
+  TORCH_CHECK(false, "Not compiled with GPU support");
 #endif
 
   if (dets.numel() == 0) {

torchvision/ops/_utils.py

 import torch
 from torch import Tensor
-from torch.jit.annotations import List, Tuple
+from torch.jit.annotations import List
 
 
 def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor: