mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
1bea5fc3ad
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13002 Batch dim wasn't handled in the CPU impl (will fail for inputs with N > 1). Fixing that here. Differential Revision: D10515159 fbshipit-source-id: ee7e4f489d2d4de793f550b31db7c0e2ba3651e8
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
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()
|