pytorch
444 строки · 16.9 Кб
1
2
3
4
5
6from caffe2.python import core, workspace7from hypothesis import given, settings8
9import caffe2.python.hypothesis_test_util as hu10import caffe2.python.serialized_test.serialized_test_util as serial11import hypothesis.strategies as st12import numpy as np13import itertools as it14
15
16class TestReduceOps(serial.SerializedTestCase):17def run_reduce_op_test_impl(18self, op_name, X, axes, keepdims, ref_func, gc, dc, allow_broadcast_fastpath):19extra_args = dict(allow_broadcast_fastpath=True) if allow_broadcast_fastpath else {}20if axes is None:21op = core.CreateOperator(22op_name,23["X"],24["Y"],25keepdims=keepdims,26**extra_args,27)28else:29op = core.CreateOperator(30op_name,31["X"],32["Y"],33axes=axes,34keepdims=keepdims,35**extra_args,36)37
38def ref(X):39return [ref_func(40X, axis=None if axes is None else tuple(axes),41keepdims=keepdims)]42
43with self.set_disable_serialized_check(allow_broadcast_fastpath):44self.assertReferenceChecks(gc, op, [X], ref)45self.assertDeviceChecks(dc, op, [X], [0])46self.assertGradientChecks(gc, op, [X], 0, [0])47
48def run_reduce_op_test(49self, op_name, X, keepdims, num_axes, ref_func, gc, dc, allow_broadcast_fastpath=False):50self.run_reduce_op_test_impl(51op_name, X, None, keepdims, ref_func, gc, dc, allow_broadcast_fastpath)52
53num_dims = len(X.shape)54if num_dims < num_axes:55self.run_reduce_op_test_impl(56op_name, X, range(num_dims), keepdims, ref_func, gc, dc, allow_broadcast_fastpath)57else:58for axes in it.combinations(range(num_dims), num_axes):59self.run_reduce_op_test_impl(60op_name, X, axes, keepdims, ref_func, gc, dc, allow_broadcast_fastpath)61
62@serial.given(63X=hu.tensor(max_dim=3, dtype=np.float32),64keepdims=st.booleans(),65allow_broadcast_fastpath=st.booleans(),66num_axes=st.integers(1, 3), **hu.gcs)67def test_reduce_min(self, X, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):68X_dims = X.shape69X_size = X.size70X = np.arange(X_size, dtype=np.float32)71np.random.shuffle(X)72X = X.reshape(X_dims)73self.run_reduce_op_test(74"ReduceMin", X, keepdims, num_axes, np.min, gc, dc,75allow_broadcast_fastpath=allow_broadcast_fastpath)76
77@serial.given(78X=hu.tensor(max_dim=3, dtype=np.float32),79keepdims=st.booleans(),80allow_broadcast_fastpath=st.booleans(),81num_axes=st.integers(1, 3), **hu.gcs)82def test_reduce_max(self, X, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):83X_dims = X.shape84X_size = X.size85X = np.arange(X_size, dtype=np.float32)86np.random.shuffle(X)87X = X.reshape(X_dims)88self.run_reduce_op_test(89"ReduceMax", X, keepdims, num_axes, np.max, gc, dc,90allow_broadcast_fastpath=allow_broadcast_fastpath)91
92@given(n=st.integers(0, 5), m=st.integers(0, 5), k=st.integers(0, 5),93t=st.integers(0, 5), keepdims=st.booleans(),94allow_broadcast_fastpath=st.booleans(),95num_axes=st.integers(1, 3), **hu.gcs)96@settings(deadline=10000)97def test_reduce_sum(self, n, m, k, t, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):98X = np.random.randn(n, m, k, t).astype(np.float32)99self.run_reduce_op_test(100"ReduceSum", X, keepdims, num_axes, np.sum, gc, dc,101allow_broadcast_fastpath=allow_broadcast_fastpath)102
103@serial.given(X=hu.tensor(dtype=np.float32), keepdims=st.booleans(),104allow_broadcast_fastpath=st.booleans(),105num_axes=st.integers(1, 4), **hu.gcs)106def test_reduce_mean(self, X, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):107self.run_reduce_op_test(108"ReduceMean", X, keepdims, num_axes, np.mean, gc, dc,109allow_broadcast_fastpath=allow_broadcast_fastpath)110
111@given(n=st.integers(1, 3), m=st.integers(1, 3), k=st.integers(1, 3),112keepdims=st.booleans(), allow_broadcast_fastpath=st.booleans(),113num_axes=st.integers(1, 3), **hu.gcs_cpu_only)114@settings(deadline=10000)115def test_reduce_l1(self, n, m, k, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):116X = np.arange(n * m * k, dtype=np.float32) - 0.5117np.random.shuffle(X)118X = X.reshape((m, n, k))119self.run_reduce_op_test(120"ReduceL1", X, keepdims, num_axes, getNorm(1), gc, dc,121allow_broadcast_fastpath=allow_broadcast_fastpath)122
123@serial.given(n=st.integers(1, 5), m=st.integers(1, 5), k=st.integers(1, 5),124keepdims=st.booleans(), allow_broadcast_fastpath=st.booleans(),125num_axes=st.integers(1, 3), **hu.gcs_cpu_only)126def test_reduce_l2(self, n, m, k, keepdims, allow_broadcast_fastpath, num_axes, gc, dc):127X = np.random.randn(n, m, k).astype(np.float32)128self.run_reduce_op_test(129"ReduceL2", X, keepdims, num_axes, getNorm(2), gc, dc,130allow_broadcast_fastpath=allow_broadcast_fastpath)131
132
133def getNorm(p):134if p == 1:135def norm(X, axis, keepdims):136return np.sum(np.abs(X), axis=axis, keepdims=keepdims)137elif p == 2:138def norm(X, axis, keepdims):139return np.sqrt(np.sum(np.power(X, 2), axis=axis, keepdims=keepdims))140else:141raise RuntimeError("Only L1 and L2 norms supported")142return norm143
144
145class TestReduceFrontReductions(serial.SerializedTestCase):146def grad_variant_input_test(self, grad_op_name, X, ref, num_reduce_dim):147workspace.ResetWorkspace()148
149Y = np.array(ref(X)[0]).astype(np.float32)150dY = np.array(np.random.rand(*Y.shape)).astype(np.float32)151shape = np.array(X.shape).astype(np.int64)152
153workspace.FeedBlob("X", X)154workspace.FeedBlob("dY", dY)155workspace.FeedBlob("shape", shape)156
157grad_op = core.CreateOperator(158grad_op_name, ["dY", "X"], ["dX"], num_reduce_dim=num_reduce_dim)159
160grad_op1 = core.CreateOperator(161grad_op_name, ["dY", "shape"], ["dX1"],162num_reduce_dim=num_reduce_dim)163
164workspace.RunOperatorOnce(grad_op)165workspace.RunOperatorOnce(grad_op1)166
167dX = workspace.FetchBlob("dX")168dX1 = workspace.FetchBlob("dX1")169np.testing.assert_array_equal(dX, dX1)170
171def max_op_test(172self, op_name, num_reduce_dim, gc, dc, in_data, in_names, ref_max):173
174op = core.CreateOperator(175op_name,176in_names,177["outputs"],178num_reduce_dim=num_reduce_dim179)180
181self.assertReferenceChecks(182device_option=gc,183op=op,184inputs=in_data,185reference=ref_max,186)187
188# Skip gradient check because it is too unreliable with max.189# Just check CPU and CUDA have same results190Y = np.array(ref_max(*in_data)[0]).astype(np.float32)191dY = np.array(np.random.rand(*Y.shape)).astype(np.float32)192if len(in_data) == 2:193grad_in_names = ["dY", in_names[0], "Y", in_names[1]]194grad_in_data = [dY, in_data[0], Y, in_data[1]]195else:196grad_in_names = ["dY", in_names[0], "Y"]197grad_in_data = [dY, in_data[0], Y]198
199grad_op = core.CreateOperator(200op_name + "Gradient",201grad_in_names,202["dX"],203num_reduce_dim=num_reduce_dim204)205self.assertDeviceChecks(dc, grad_op, grad_in_data, [0])206
207def reduce_op_test(self, op_name, op_ref, in_data, in_names,208num_reduce_dims, device):209op = core.CreateOperator(210op_name,211in_names,212["outputs"],213num_reduce_dim=num_reduce_dims214)215
216self.assertReferenceChecks(217device_option=device,218op=op,219inputs=in_data,220reference=op_ref221)222
223self.assertGradientChecks(224device, op, in_data, 0, [0], stepsize=1e-2, threshold=1e-2)225
226@given(num_reduce_dim=st.integers(0, 4), **hu.gcs)227@settings(deadline=10000)228def test_reduce_front_sum(self, num_reduce_dim, gc, dc):229X = np.random.rand(7, 4, 3, 5).astype(np.float32)230
231def ref_sum(X):232return [np.sum(X, axis=(tuple(range(num_reduce_dim))))]233
234self.reduce_op_test(235"ReduceFrontSum", ref_sum, [X], ["input"], num_reduce_dim, gc)236self.grad_variant_input_test(237"ReduceFrontSumGradient", X, ref_sum, num_reduce_dim)238
239@given(num_reduce_dim=st.integers(0, 4), seed=st.integers(0, 4), **hu.gcs)240def test_reduce_front_sum_empty_batch(self, num_reduce_dim, seed, gc, dc):241np.random.seed(seed)242X = np.random.rand(0, 4, 3, 5).astype(np.float32)243
244def ref_sum(X):245return [np.sum(X, axis=(tuple(range(num_reduce_dim))))]246
247self.reduce_op_test(248"ReduceFrontSum", ref_sum, [X], ["input"], num_reduce_dim, gc)249self.grad_variant_input_test(250"ReduceFrontSumGradient", X, ref_sum, num_reduce_dim)251
252# test the second iteration253not_empty_X = np.random.rand(2, 4, 3, 5).astype(np.float32)254net = core.Net('test')255with core.DeviceScope(gc):256net.ReduceFrontSum(257['X'], ['output'],258num_reduce_dim=num_reduce_dim259)260workspace.CreateNet(net)261
262workspace.FeedBlob('X', not_empty_X)263workspace.RunNet(workspace.GetNetName(net))264output = workspace.FetchBlob('output')265np.testing.assert_allclose(266output, ref_sum(not_empty_X)[0], atol=1e-3)267
268workspace.FeedBlob('X', X)269workspace.RunNet(workspace.GetNetName(net))270output = workspace.FetchBlob('output')271np.testing.assert_allclose(output, ref_sum(X)[0], atol=1e-3)272
273@given(**hu.gcs)274@settings(deadline=None)275def test_reduce_front_sum_with_length(self, dc, gc):276num_reduce_dim = 1277X = np.random.rand(2, 3, 4, 5).astype(np.float32)278batch_size = int(np.prod([2, 3, 4, 5][num_reduce_dim:]))279d = 120 // batch_size280lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)281
282def ref_sum(X, lengths):283Y = X.reshape(d, lengths.size)284rv = np.zeros((lengths.size, 1)).astype(np.float32)285for ii in range(lengths.size):286rv[ii] = np.sum(Y[:lengths[ii], ii])287return [rv.reshape((2, 3, 4, 5)[num_reduce_dim:])]288
289self.reduce_op_test(290"ReduceFrontSum", ref_sum, [X, lengths], ["input", "lengths"],291num_reduce_dim, gc)292
293@given(num_reduce_dim=st.integers(0, 4), **hu.gcs)294@settings(deadline=10000)295def test_reduce_front_mean(self, num_reduce_dim, gc, dc):296X = np.random.rand(6, 7, 8, 2).astype(np.float32)297
298def ref_mean(X):299return [np.mean(X, axis=(tuple(range(num_reduce_dim))))]300
301self.reduce_op_test(302"ReduceFrontMean", ref_mean, [X], ["input"], num_reduce_dim, gc)303self.grad_variant_input_test(304"ReduceFrontMeanGradient", X, ref_mean, num_reduce_dim)305
306@given(**hu.gcs)307@settings(deadline=10000)308def test_reduce_front_mean_with_length(self, dc, gc):309num_reduce_dim = 1310X = np.random.rand(2, 3, 4, 5).astype(np.float32)311batch_size = int(np.prod([2, 3, 4, 5][num_reduce_dim:]))312d = 120 // batch_size313lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)314
315def ref_mean(X, lengths):316Y = X.reshape(d, lengths.size)317rv = np.zeros((lengths.size, 1)).astype(np.float32)318for ii in range(lengths.size):319rv[ii] = np.mean(Y[:lengths[ii], ii])320return [rv.reshape((2, 3, 4, 5)[num_reduce_dim:])]321
322self.reduce_op_test(323"ReduceFrontMean", ref_mean, [X, lengths], ["input", "lengths"],324num_reduce_dim, gc)325
326@serial.given(num_reduce_dim=st.integers(0, 4), **hu.gcs)327def test_reduce_front_max(self, num_reduce_dim, gc, dc):328X = np.random.rand(6, 7, 8, 2).astype(np.float32)329
330def ref_frontmax(X):331return [np.max(X, axis=(tuple(range(num_reduce_dim))))]332
333self.max_op_test(334"ReduceFrontMax", num_reduce_dim, gc, dc, [X], ["X"], ref_frontmax)335
336@given(**hu.gcs)337def test_reduce_front_max_with_length(self, dc, gc):338num_reduce_dim = 1339X = np.random.rand(2, 3, 4, 5).astype(np.float32)340batch_size = int(np.prod([2, 3, 4, 5][num_reduce_dim:]))341d = 120 // batch_size342lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)343
344def ref_max(X, lengths):345Y = X.reshape(d, lengths.size)346rv = np.zeros((lengths.size, 1)).astype(np.float32)347for ii in range(lengths.size):348rv[ii] = np.max(Y[:lengths[ii], ii])349return [rv.reshape((2, 3, 4, 5)[num_reduce_dim:])]350
351self.max_op_test(352"ReduceFrontMax", num_reduce_dim, gc, dc, [X, lengths],353["X", "lengths"], ref_max)354
355@serial.given(num_reduce_dim=st.integers(0, 4), **hu.gcs)356def test_reduce_back_max(self, num_reduce_dim, gc, dc):357X = np.random.rand(6, 7, 8, 2).astype(np.float32)358
359def ref_backmax(X):360return [np.max(X, axis=(0, 1, 2, 3)[4 - num_reduce_dim:])]361
362self.max_op_test(363"ReduceBackMax", num_reduce_dim, gc, dc, [X], ["X"], ref_backmax)364
365@given(**hu.gcs)366def test_reduce_back_max_with_length(self, gc, dc):367num_reduce_dim = 1368X = np.random.rand(2, 3, 4, 5).astype(np.float32)369batch_size = int(np.prod([2, 3, 4, 5][:4 - num_reduce_dim]))370d = 120 // batch_size371lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)372
373def ref_max(X, lengths):374Y = X.reshape(lengths.size, d)375rv = np.zeros((lengths.size, 1)).astype(np.float32)376for ii in range(lengths.size):377rv[ii] = np.max(Y[ii, :lengths[ii]])378return [rv.reshape((2, 3, 4, 5)[:4 - num_reduce_dim])]379
380self.max_op_test(381"ReduceBackMax", num_reduce_dim, gc, dc, [X, lengths],382["X", "lengths"], ref_max)383
384@given(**hu.gcs)385@settings(deadline=10000)386def test_reduce_back_sum(self, dc, gc):387num_reduce_dim = 1388X = np.random.rand(6, 7, 8, 2).astype(np.float32)389
390def ref_sum(X):391return [np.sum(X, axis=(0, 1, 2, 3)[4 - num_reduce_dim:])]392
393self.reduce_op_test(394"ReduceBackSum", ref_sum, [X], ["input"], num_reduce_dim, gc)395self.grad_variant_input_test(396"ReduceBackSumGradient", X, ref_sum, num_reduce_dim)397
398@given(**hu.gcs)399@settings(deadline=10000)400def test_reduce_back_sum_with_length(self, dc, gc):401num_reduce_dim = 1402X = np.random.rand(2, 3, 4, 5).astype(np.float32)403batch_size = int(np.prod([2, 3, 4, 5][:4 - num_reduce_dim]))404d = 120 // batch_size405lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)406
407def ref_sum(X, lengths):408Y = X.reshape(lengths.size, d)409rv = np.zeros((lengths.size, 1)).astype(np.float32)410for ii in range(lengths.size):411rv[ii] = np.sum(Y[ii, :lengths[ii]])412return [rv.reshape((2, 3, 4, 5)[:4 - num_reduce_dim])]413
414self.reduce_op_test(415"ReduceBackSum", ref_sum, [X, lengths], ["input", "lengths"],416num_reduce_dim, gc)417
418@given(num_reduce_dim=st.integers(0, 4), **hu.gcs)419@settings(deadline=10000)420def test_reduce_back_mean(self, num_reduce_dim, dc, gc):421X = np.random.rand(6, 7, 8, 2).astype(np.float32)422
423def ref_mean(X):424return [np.mean(X, axis=(0, 1, 2, 3)[4 - num_reduce_dim:])]425
426self.reduce_op_test(427"ReduceBackMean", ref_mean, [X], ["input"], num_reduce_dim, gc)428self.grad_variant_input_test(429"ReduceBackMeanGradient", X, ref_mean, num_reduce_dim)430
431@given(**hu.gcs)432@settings(deadline=None)433def test_reduce_back_mean_with_length(self, dc, gc):434num_reduce_dim = 1435X = np.random.rand(2, 3, 4, 5).astype(np.float32)436batch_size = int(np.prod([2, 3, 4, 5][:4 - num_reduce_dim]))437d = 120 // batch_size438lengths = np.random.randint(1, d, size=batch_size).astype(np.int32)439
440def ref_mean(X, lengths):441Y = X.reshape(lengths.size, d)442rv = np.zeros((lengths.size, 1)).astype(np.float32)443for ii in range(lengths.size):444rv[ii] = np.mean(Y[ii, :lengths[ii]])445return [rv.reshape((2, 3, 4, 5)[:4 - num_reduce_dim])]446
447self.reduce_op_test(448"ReduceBackMean", ref_mean, [X, lengths], ["input", "lengths"],449num_reduce_dim, gc)450