diff --git a/test/test_ops.py b/test/test_ops.py index e8844b47422..9f9c2eb9da3 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -162,6 +162,35 @@ def test_forward(self, device, contiguous, x_dtype, rois_dtype=None, determinist torch.testing.assert_close(gt_y.to(y), y, rtol=tol, atol=tol) + @pytest.mark.parametrize("device", cpu_and_cuda_and_mps()) + @pytest.mark.parametrize( + "coords", + ( + pytest.param([float("nan")] * 4, id="nan"), + pytest.param([float("inf"), float("-inf"), float("inf"), float("-inf")], id="inf"), + # Finite, but roi_end - roi_start overflows to inf, and 0 * inf is nan. + pytest.param([-3.4e38, -3.4e38, 3.4e38, 3.4e38], id="overflow"), + ), + ) + def test_non_finite_boxes(self, device, coords): + # A non-finite coordinate makes every comparison in the bounds check + # false, so the sample point used to be neither rejected nor clamped, + # and the index was then computed from (int)nan -- undefined behaviour, + # INT_MIN on x86 -- reading far outside the input buffer. + pool_size = 5 + # n_channels % (pool_size ** 2) == 0 required for PS operations. + n_channels = 2 * (pool_size**2) + x = torch.rand(1, n_channels, 10, 10, device=device, requires_grad=True) + rois = torch.tensor([[0.0] + coords], dtype=torch.float32, device=device) + + y = self.fn(x, rois, pool_size, pool_size, spatial_scale=1, sampling_ratio=1) + assert y.shape[0] == rois.shape[0] + assert not y.isinf().any() + + y.sum().backward() + assert x.grad.shape == x.shape + assert not x.grad.isinf().any() + @pytest.mark.parametrize("device", cpu_and_cuda()) def test_is_leaf_node(self, device): op_obj = self.make_obj(wrap=True).to(device=device) diff --git a/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp b/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp index dc1f380f36c..c2f12d55bb5 100644 --- a/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp @@ -23,7 +23,10 @@ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -166,7 +169,10 @@ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cpu/roi_align_common.h b/torchvision/csrc/ops/cpu/roi_align_common.h index 03e5084357f..173b2ac6314 100644 --- a/torchvision/csrc/ops/cpu/roi_align_common.h +++ b/torchvision/csrc/ops/cpu/roi_align_common.h @@ -57,7 +57,10 @@ void pre_calc_for_bilinear_interpolate( T x = xx; T y = yy; // deal with: inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty PreCalc pc; pc.pos1 = 0; diff --git a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp index 26dc7423853..d35cd3f550b 100644 --- a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp @@ -130,7 +130,10 @@ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu b/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu index 74f7e9d6004..97313250cfb 100644 --- a/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu +++ b/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu @@ -32,7 +32,10 @@ __device__ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -167,7 +170,10 @@ __device__ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cuda/roi_align_kernel.cu b/torchvision/csrc/ops/cuda/roi_align_kernel.cu index 346b6e4f0e2..99a05c5951c 100644 --- a/torchvision/csrc/ops/cuda/roi_align_kernel.cu +++ b/torchvision/csrc/ops/cuda/roi_align_kernel.cu @@ -32,7 +32,10 @@ __device__ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -170,7 +173,10 @@ __device__ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/mps/mps_stable_kernels.h b/torchvision/csrc/ops/mps/mps_stable_kernels.h index 9ced5a95fed..5156f2f924e 100644 --- a/torchvision/csrc/ops/mps/mps_stable_kernels.h +++ b/torchvision/csrc/ops/mps/mps_stable_kernels.h @@ -57,7 +57,10 @@ inline T bilinear_interpolate( T x, uint index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -164,7 +167,10 @@ inline void bilinear_interpolate_gradient( thread integer_t& y_high, uint index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1;