diff --git a/modules/detectron/upsample_nearest_op.h b/modules/detectron/upsample_nearest_op.h index ba5890400a9d5..a891de420cb9d 100644 --- a/modules/detectron/upsample_nearest_op.h +++ b/modules/detectron/upsample_nearest_op.h @@ -50,7 +50,7 @@ class UpsampleNearestOp final : public Operator { d2 = Y->dim32(1); d3 = Y->dim32(2); } else { - d1 = Y->dim32(1); + d1 = Y->dim32(0) * Y->dim32(1); d2 = Y->dim32(2); d3 = Y->dim32(3); } diff --git a/modules/detectron/upsample_nearest_op_test.py b/modules/detectron/upsample_nearest_op_test.py new file mode 100644 index 0000000000000..5eef675029ab9 --- /dev/null +++ b/modules/detectron/upsample_nearest_op_test.py @@ -0,0 +1,43 @@ +from __future__ import absolute_import, division, print_function, unicode_literals + +import unittest + +import caffe2.python.hypothesis_test_util as hu +import hypothesis.strategies as st +import numpy as np +from caffe2.python import core, dyndep +from hypothesis import given + + +dyndep.InitOpsLibrary("@/caffe2/modules/detectron:detectron_ops") + + +class TestUpsampleNearestOp(hu.HypothesisTestCase): + @given( + N=st.integers(1, 3), + H=st.integers(10, 300), + W=st.integers(10, 300), + scale=st.integers(1, 3), + **hu.gcs + ) + def test_upsample_nearest_op(self, N, H, W, scale, gc, dc): + C = 32 + X = np.random.randn(N, C, H, W).astype(np.float32) + op = core.CreateOperator("UpsampleNearest", ["X"], ["Y"], scale=scale) + + def ref(X): + outH = H * scale + outW = W * scale + outH_idxs, outW_idxs = np.meshgrid( + np.arange(outH), np.arange(outW), indexing="ij" + ) + inH_idxs = (outH_idxs / scale).astype(np.int32) + inW_idxs = (outW_idxs / scale).astype(np.int32) + Y = X[:, :, inH_idxs, inW_idxs] + return [Y] + + self.assertReferenceChecks(device_option=gc, op=op, inputs=[X], reference=ref) + + +if __name__ == "__main__": + unittest.main()