diff --git a/test/test_ops.py b/test/test_ops.py index 8f8756fe675..8b2aa32aba9 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -474,6 +474,31 @@ def expected_fn( def test_boxes_shape(self): self._helper_boxes_shape(ops.roi_align) + @pytest.mark.parametrize("aligned", (True, False)) + @pytest.mark.parametrize("sampling_ratio", (-1, 1)) + def test_nan_roi_coordinates_do_not_crash(self, aligned, sampling_ratio): + # Non-regression test: a ROI with NaN coordinates used to bypass the + # out-of-bounds check in the CPU kernel (every comparison against NaN + # is False), get truncated to an undefined int index, and read out of + # bounds from the input feature map (can segfault under ASAN). + x = torch.zeros(2, 10, 15, 19, dtype=torch.float32) + rois = torch.tensor( + [[0.0, float("nan"), 0.0, float("nan"), 1.0]], + dtype=torch.float32, + ) + + y = ops.roi_align( + x, + rois, + output_size=(7, 7), + spatial_scale=0.1, + sampling_ratio=sampling_ratio, + aligned=aligned, + ) + + assert y.shape == (1, 10, 7, 7) + assert torch.isfinite(y).all() or torch.isnan(y).all() + @pytest.mark.parametrize("aligned", (True, False)) @pytest.mark.parametrize("device", cpu_and_cuda_and_mps()) @pytest.mark.parametrize("x_dtype", (torch.float16, torch.float32, torch.float64)) # , ids=str) diff --git a/torchvision/csrc/ops/cpu/roi_align_common.h b/torchvision/csrc/ops/cpu/roi_align_common.h index e10c67b5b79..1c7bc318846 100644 --- a/torchvision/csrc/ops/cpu/roi_align_common.h +++ b/torchvision/csrc/ops/cpu/roi_align_common.h @@ -1,6 +1,7 @@ #pragma once #include +#include namespace vision { namespace ops { @@ -57,7 +58,13 @@ 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) { + // NaN coordinates must be treated as out of bounds too, since every + // comparison against NaN is false and would otherwise fall through + // this check, get truncated to an undefined int value below, and + // produce an out-of-bounds pos1/pos2/pos3/pos4 index. + if (std::isnan(static_cast(x)) || + std::isnan(static_cast(y)) || 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 e0185da45df..ac1cb00cef8 100644 --- a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp @@ -1,6 +1,8 @@ #include #include +#include + #include "./roi_align_common.h" namespace vision { @@ -123,7 +125,12 @@ 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) { + // NaN coordinates must be treated as out of bounds too, since every + // comparison against NaN is false and would otherwise fall through this + // check and get truncated to an undefined int value below. + if (std::isnan(static_cast(x)) || + std::isnan(static_cast(y)) || 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;