def gather_custom(x, mask, bsize, ksize, strides, padding):
x_shape = [int(ss) for ss in x.get_shape()]
block_params = calc_block_params(x_shape, bsize, ksize, strides, padding)
indices = convert_mask_to_indices_custom(mask, block_params, 0.0)
p = sbnet_module.sparse_gather(x,
indices.bin_counts,
indices.active_block_indices,
bsize=block_params.bsize,
boffset=block_params.boffset,
bstride=block_params.bstrides)
return p, indices
def _test_sparse_scatter(self, mask, x, w, out_shape, bsize, ksize,
strides, padding):
with tf.Session() as sess:
x = tf.constant(x)
w = tf.constant(w)
p, blk_indices = gather_tf(x, mask, bsize, ksize, strides, padding)
block_params = calc_block_params([int(ss) for ss in x.get_shape()],
bsize, ksize, strides, padding)
ind_custom = convert_mask_to_indices_custom(
mask, block_params, 0.0)
p_custom = sbnet_module.sparse_gather(
x,
ind_custom.bin_counts,
ind_custom.active_block_indices,
bsize=block_params.bsize,
bstride=block_params.bstrides,
boffset=block_params.boffset)
p_shape = [
int(x.get_shape()[0]), block_params.bsize[0],
block_params.bsize[1],
int(x.get_shape()[3])
]
q = tf.nn.conv2d(p, w, strides, 'VALID')
q_custom = tf.nn.conv2d(p_custom, w, strides, 'VALID')
y_tf = scatter_tf(q, blk_indices, out_shape)
q_shape = calc_out_size_4d_np(p_shape, ksize, strides, 'VALID')
bsize_out = [q_shape[1], q_shape[2]]
boffset = [0, 0]
y_custom = scatter_custom(q_custom, ind_custom, out_shape,
bsize_out, boffset,
block_params.bstrides)
p1, p2, q_val, y1, y2, active, num = sess.run([
p, p_custom, q, y_tf, y_custom,
ind_custom.active_block_indices, ind_custom.bin_counts
])
num = num[0]
sortIdx = active[:num].argsort()
p2 = p2[sortIdx]
# Make sure p's are the same.
np.testing.assert_array_equal(p1, p2)
# Check y's are the same.
np.testing.assert_array_equal(y1, y2)
def _test_sparse_scatter(self, mask, x, w, out_shape, bsize, ksize,
strides, padding):
with tf.Session() as sess:
x = tf.constant(x)
w = tf.constant(w)
p, blk_indices = gather_tf(x, mask, bsize, ksize, strides, padding)
block_params = calc_block_params([int(ss) for ss in x.get_shape()],
bsize, ksize, strides, padding)
ind_custom = convert_mask_to_indices_custom(
mask, block_params, 0.0)
p_custom = sbnet_module.sparse_gather(
x,
ind_custom.bin_counts,
ind_custom.active_block_indices,
dynamic_bsize=tf.constant(block_params.bsize, tf.int32),
dynamic_bstride=tf.constant(block_params.bstrides, tf.int32),
dynamic_boffset=tf.constant(block_params.boffset, tf.int32))
p_shape = [
int(x.get_shape()[0]), block_params.bsize[0],
block_params.bsize[1],
int(x.get_shape()[3])
]
q = tf.nn.conv2d(p, w, strides, 'VALID')
q_custom = tf.nn.conv2d(p_custom, w, strides, 'VALID')
y_tf = scatter_tf(q, blk_indices, out_shape)
q_shape = calc_out_size_4d_np(p_shape, ksize, strides, 'VALID')
bsize_out = [q_shape[1], q_shape[2]]
boffset = [0, 0]
y_custom = scatter_custom(q_custom, ind_custom, out_shape,
bsize_out, boffset,
block_params.bstrides)
p1, p2, q_val, y1, y2, active, num = sess.run([
p, p_custom, q, y_tf, y_custom,
ind_custom.active_block_indices, ind_custom.bin_counts
])
# Make sure p's are the same.
l1 = tuple([tuple(x) for x in p1.reshape(-1, 3).tolist()])
l2 = tuple([tuple(x) for x in p2.reshape(-1, 3).tolist()])
np.testing.assert_array_equal(set(l1), set(l2))
# Check y's are the same.
np.testing.assert_array_equal(y1, y2)