def test_keras(data_format):
# Unable to use `layer_test` as this layer has multiple inputs.
val_a, val_b = _create_test_data(data_format)
input_a = tf.keras.Input(shape=val_a.shape[1:])
input_b = tf.keras.Input(shape=val_b.shape[1:])
layer = CorrelationCost(
kernel_size=1,
max_displacement=2,
stride_1=1,
stride_2=2,
pad=4,
data_format=data_format,
)
expected_output_shape = tuple(
layer.compute_output_shape([input_a.shape, input_b.shape]))
x = [input_a, input_b]
y = layer(x)
model = tf.keras.models.Model(x, y)
actual_output = model([val_a, val_b])
expected_output_type = "float32"
assert tf.keras.backend.dtype(y[0]) == expected_output_type
assert actual_output.shape[1:] == expected_output_shape[0][1:]
def _forward(
self,
input_a,
input_b,
kernel_size,
max_displacement,
stride_1,
stride_2,
pad,
data_format,
):
input_a_op = tf.convert_to_tensor(input_a, dtype=tf.float32)
input_b_op = tf.convert_to_tensor(input_b, dtype=tf.float32)
output = CorrelationCost(
kernel_size=kernel_size,
max_displacement=max_displacement,
stride_1=stride_1,
stride_2=stride_2,
pad=pad,
data_format=data_format,
)([input_a_op, input_b_op])
return output
def correlation_fn(input_a, input_b):
return CorrelationCost(
kernel_size=kernel_size,
max_displacement=max_displacement,
stride_1=stride_1,
stride_2=stride_2,
pad=pad,
data_format=data_format)([input_a, input_b])
def _keras(self, data_format):
# Unable to use `layer_test` as this layer has multiple inputs.
with test_utils.use_gpu():
val_a = [[[[0, -6, 9, 5], [1, -5, 10, 3], [2, -4, 11, 1]],
[[3, -3, 12, -1], [4, -2, 13, -3], [5, -1, 14, -5]]],
[[[6, 0, 15, -7], [7, 1, 16, -9], [8, 2, 17, -11]],
[[9, 3, 18, -13], [10, 4, 19, -15], [11, 5, 20, -17]]]]
val_b = np.array(val_a).transpose(2, 3, 0, 1).reshape(2, 2, 3, 4)
# yapf: disable
input_a = tf.keras.Input(shape=(2, 3, 4,))
input_b = tf.keras.Input(shape=(2, 3, 4,))
layer = CorrelationCost(
kernel_size=1,
max_displacement=2,
stride_1=1,
stride_2=2,
pad=4,
data_format=data_format)
expected_output_shape = tuple(
layer.compute_output_shape([(2, 3, 4,), (2, 3, 4,)]))[1:]
# yapf: enable
x = [input_a, input_b]
y = layer(x)
model = tf.keras.models.Model(x, y)
actual_output = model.predict([val_a, val_b])
expected_output_type = 'float32'
if tf.keras.backend.dtype(y[0]) != expected_output_type:
raise AssertionError(
"Inferred output type %s does not equal "
"expected output type %s" %
(tf.keras.backend.dtype(y[0]), expected_output_type))
if actual_output[0].shape != expected_output_shape:
raise AssertionError(
"Expected shape %s does not equal output shape"
"%s" % (actual_output[0].shape, expected_output_shape))
def _keras(self, data_format, use_gpu=False):
# Unable to use `layer_test` as this layer has multiple inputs.
with test_utils.device(use_gpu):
val_a, val_b = _create_test_data(data_format)
# yapf: disable
input_a = tf.keras.Input(shape=val_a.shape[1:])
input_b = tf.keras.Input(shape=val_b.shape[1:])
layer = CorrelationCost(
kernel_size=1,
max_displacement=2,
stride_1=1,
stride_2=2,
pad=4,
data_format=data_format)
expected_output_shape = tuple(
layer.compute_output_shape([input_a.shape, input_b.shape]))
# yapf: enable
x = [input_a, input_b]
y = layer(x)
model = tf.keras.models.Model(x, y)
actual_output = model([val_a, val_b])
expected_output_type = "float32"
if tf.keras.backend.dtype(y[0]) != expected_output_type:
raise AssertionError(
"Inferred output type %s does not equal "
"expected output type %s"
% (tf.keras.backend.dtype(y[0]), expected_output_type)
)
if actual_output.shape[1:] != expected_output_shape[0][1:]:
raise AssertionError(
"Expected shape %s does not equal output shape"
"%s" % (actual_output.shape, expected_output_shape[0])
)