/usr/local/lib64/python3.6/site-packages/caffe2/python/operator_test
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/operator_test/gather_ops_test.py (9216B)
import numpy as np
from caffe2.python import core, workspace
from hypothesis import given, settings
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.serialized_test.serialized_test_util as serial
import hypothesis.strategies as st
import hypothesis.extra.numpy as hnp
# Basic implementation of gather for axis == 0, shich is lookup of indices
# in the outer dimension. Keeping it for reference here, although is similar
# to more general function below.
def ref_gather_axis0():
def inner(data, ind):
if ind.size == 0 or data.shape[0] == 0:
return [np.zeros((0, 10, 20)).astype(np.float32)]
output = [data[i] for i in ind]
return [output]
return inner
# Returns axis-based lookup. We just use numpy take() which handles different
# axis values as we want.
def ref_gather(axis):
def inner(data, ind):
if ind.size == 0 or data.shape[axis] == 0:
shape = list(data.shape)
shape[0] = 0
return [np.zeros(tuple(shape)).astype(np.float32)]
# np.take() does axis lookup same as gather
output = data.take(ind, axis).astype(np.float32)
return [output]
return inner
# Gather(..., match_outer==True)
def ref_gather_match_outer(axis=1):
def inner(data, ind):
if ind.size == 0 or data.shape[axis] == 0:
shape = list(data.shape)
shape[0] = 0
return [np.zeros(tuple(shape)).astype(np.float32)]
input_shape = list(data.shape)
output_shape = input_shape[:axis] + list(ind.shape[axis:]) + input_shape[axis + 1:]
output = np.zeros(tuple(output_shape)).astype(np.float32)
if axis == 1:
for i in range(data.shape[0]):
output[i] = data[i, ind[i], ]
elif axis == 2:
for i in range(data.shape[0]):
for j in range(data.shape[1]):
output[i, j] = data[i, j, ind[i, j], ]
else:
raise NotImplementedError
return [output]
return inner
class TestGatherOps(serial.SerializedTestCase):
@given(rows_num=st.integers(0, 10000),
index_num=st.integers(0, 5000),
**hu.gcs)
@settings(deadline=10000)
def test_gather_ops(self, rows_num, index_num, gc, dc):
data = np.random.random((rows_num, 10, 20)).astype(np.float32)
if rows_num > 0:
ind = np.random.randint(rows_num, size=(index_num, )).astype('int32')
else:
ind = np.random.randint(10, size=(index_num, )).astype('int32')
op = core.CreateOperator(
'Gather',
['data', 'ind'],
['output'])
self.assertReferenceChecks(gc, op, [data, ind], ref_gather_axis0())
self.assertDeviceChecks(dc, op, [data, ind], [0])
return
# Test axis == 2, this keeps outer dimension but will replace data
# within axis by lookup of index array (repeated for each outer entry)
@given(batch_num=st.integers(1, 4000),
rows_num=st.integers(1, 6),
index_num=st.integers(1, 20),
**hu.gcs)
def test_gather_ops_axis2(self, batch_num, rows_num, index_num, gc, dc):
data = np.random.random((batch_num, rows_num, 5)).astype(np.float32)
ind = np.random.randint(5, size=(index_num, )).astype('int32')
op = core.CreateOperator(
'Gather',
['data', 'ind'],
['output'],
axis=2)
self.assertReferenceChecks(gc, op, [data, ind], ref_gather(axis=2))
self.assertDeviceChecks(dc, op, [data, ind], [0])
return
# Test match_outer == true, the indices has the same outer dimensions as data
@given(batch_num=st.integers(1, 40),
rows_num=st.integers(1, 6),
index_num=st.integers(1, 20),
**hu.gcs_cpu_only)
@settings(deadline=10000)
def test_gather_ops_match_outer(self, batch_num, rows_num, index_num, gc, dc):
data = np.random.random((batch_num, rows_num, 5)).astype(np.float32)
ind = np.random.randint(rows_num, size=(batch_num, index_num)).astype('int32')
op = core.CreateOperator(
'Gather',
['data', 'ind'],
['output'],
axis=1,
match_outer=True)
self.assertReferenceChecks(gc, op, [data, ind], ref_gather_match_outer())
self.assertDeviceChecks(dc, op, [data, ind], [0])
self.assertGradientChecks(gc, op, [data, ind], 0, [0])
return
# Test BatchGather with match_outer == true, the indices has the same outer dimensions as data
# Note BatchGather is equivalent to Gather(..., axis=1)
@given(batch_num=st.integers(1, 40),
rows_num=st.integers(1, 6),
index_num=st.integers(1, 20),
**hu.gcs_cpu_only)
@settings(deadline=10000)
def test_batch_gather_op_match_outer(self, batch_num, rows_num, index_num, gc, dc):
data = np.random.random((batch_num, rows_num, 5)).astype(np.float32)
ind = np.random.randint(rows_num, size=(batch_num, index_num)).astype('int32')
op = core.CreateOperator(
'BatchGather',
['data', 'ind'],
['output'],
match_outer=True)
self.assertReferenceChecks(gc, op, [data, ind], ref_gather_match_outer())
self.assertDeviceChecks(dc, op, [data, ind], [0])
self.assertGradientChecks(gc, op, [data, ind], 0, [0])
return
# when the data is larger,
# this test sometimes passes, sometimes fails,
# test log here: https://fb.quip.com/SeiyAVWQXvsN (second run failed)
# after some digging, this turns out to be numerical error,
# the failed run has max|grad - estimated_grad| = 0.009
# so here we changed the gradient checking threshold to 0.02 for this test to pass
@given(batch_num=st.integers(1, 30),
rows_num=st.integers(1, 6),
index_num=st.integers(1, 10),
index_num2=st.integers(1, 10),
axis2_num=st.integers(1, 10),
**hu.gcs_cpu_only)
@settings(deadline=None, max_examples=50)
def test_gather_op_match_outer_axis2_data4D_ind4D(
self, batch_num, rows_num, axis2_num, index_num, index_num2, gc, dc
):
data = np.random.random((batch_num, rows_num, axis2_num, 5)).astype(np.float32)
ind = np.random.randint(axis2_num, size=(batch_num, rows_num, index_num, index_num2)).astype('int32')
op = core.CreateOperator(
'Gather',
['data', 'ind'],
['output'],
axis=2,
match_outer=True)
self.assertReferenceChecks(gc, op, [data, ind], ref_gather_match_outer(axis=2))
self.assertDeviceChecks(dc, op, [data, ind], [0])
self.assertGradientChecks(gc, op, [data, ind], 0, [0], threshold=0.02)
return
# Generates data arrays of max dims 10x100x2 and indexing array up to rows_num
@st.composite
def _inputs(draw):
batch_size = draw(st.integers(2, 10))
rows_num = draw(st.integers(1, 100))
block_size = draw(st.integers(1, 2))
index_num = draw(st.integers(1, 10))
return (
draw(hnp.arrays(
np.float32,
(batch_size, rows_num, block_size),
elements=hu.floats(-10.0, 10.0),
)),
draw(hnp.arrays(
np.int32,
(index_num, 1),
elements=st.integers(0, rows_num - 1),
)),
)
class TestBatchGatherOps(hu.HypothesisTestCase):
@given(inputs=_inputs(),
**hu.gcs)
@settings(deadline=10000)
def test_batch_gather_ops(self, inputs, gc, dc):
data, ind = inputs
op = core.CreateOperator(
'BatchGather',
['data', 'ind'],
['output'])
self.assertReferenceChecks(gc, op, [data, ind], ref_gather(axis=1))
self.assertGradientChecks(gc, op, [data, ind], 0, [0])
class TestGatherFused8BitRowwise(hu.HypothesisTestCase):
@given(rows_num=st.integers(1, 10000),
cols_num=st.integers(1, 128),
index_num=st.integers(0, 5000),
**hu.gcs)
@settings(deadline=10000)
def test_batch_gather_ops(self, rows_num, cols_num, index_num, gc, dc):
data = np.random.random((rows_num, cols_num)).astype(np.float32)
ind = np.random.randint(rows_num, size=(index_num, )).astype('int32')
net = core.Net("bench")
quantized_data = net.FloatToFused8BitRowwiseQuantized(
'data', 'quantized_data')
dequantized_data = net.Fused8BitRowwiseQuantizedToFloat(
quantized_data, 'dequantized_data')
net.Gather(
[dequantized_data, 'ind'], 'gather_reference')
net.GatherFused8BitRowwise(
[quantized_data, 'ind'], 'gather_quantized')
workspace.FeedBlob('data', data)
workspace.FeedBlob('ind', ind)
workspace.CreateNet(net)
workspace.RunNetOnce(net)
gather_reference = workspace.FetchBlob('gather_reference')
gather_quantized = workspace.FetchBlob('gather_quantized')
np.testing.assert_array_almost_equal(gather_reference, gather_quantized)
if __name__ == "__main__":
import unittest
unittest.main()