def test_simple_upsampling_block_refine_convs_bn_post(self):
block = encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv=False,
interp_method="bilinear",
skip_connection=True,
refine_convs=2,
refine_convs_filters=16,
refine_convs_use_bias=True,
refine_convs_kernel_size=3,
refine_convs_batch_norm=True,
refine_convs_batch_norm_before_activation=False,
refine_convs_activation="relu",
)
x_in = tf.keras.Input((8, 8, 1))
x = block.make_block(x_in)
model = tf.keras.Model(x_in, x)
self.assertEqual(len(model.layers), 8)
self.assertEqual(len(model.trainable_weights), 8)
self.assertEqual(model.count_params(), 2608)
self.assertAllEqual(model.output.shape, (None, 16, 16, 16))
self.assertIsInstance(model.layers[1], tf.keras.layers.UpSampling2D)
self.assertIsInstance(model.layers[2], tf.keras.layers.Conv2D)
self.assertIsInstance(model.layers[3], tf.keras.layers.Activation)
self.assertIsInstance(model.layers[4],
tf.keras.layers.BatchNormalization)
def decoder_stack(self) -> List[encoder_decoder.SimpleUpsamplingBlock]:
"""Define the decoder stack."""
blocks = []
for block in range(self.up_blocks):
block_filters_in = int(
self.filters *
(self.filters_rate
**(self.down_blocks + self.stem_blocks - 1 - block)))
if self.block_contraction:
block_filters_out = int(
self.filters *
(self.filters_rate
**(self.down_blocks + self.stem_blocks - 2 - block)))
else:
block_filters_out = block_filters_in
blocks.append(
encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv=(not self.up_interpolate),
transposed_conv_filters=block_filters_in,
transposed_conv_kernel_size=self.kernel_size,
transposed_conv_batch_norm=False,
interp_method="bilinear",
skip_connection=True,
skip_add=False,
refine_convs=self.convs_per_block,
refine_convs_first_filters=block_filters_in,
refine_convs_filters=block_filters_out,
refine_convs_kernel_size=self.kernel_size,
refine_convs_batch_norm=False,
))
return blocks
def test_simple_upsampling_block(self):
block = encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv=False,
interp_method="bilinear",
refine_convs=0,
)
x_in = tf.keras.Input((8, 8, 1))
x = block.make_block(x_in)
model = tf.keras.Model(x_in, x)
self.assertEqual(len(model.layers), 2)
self.assertEqual(len(model.trainable_weights), 0)
self.assertEqual(model.count_params(), 0)
self.assertAllEqual(model.output.shape, (None, 16, 16, 1))
self.assertIsInstance(model.layers[1], tf.keras.layers.UpSampling2D)
def test_simple_upsampling_block_skip_concat(self):
block = encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv=False,
interp_method="bilinear",
skip_connection=True,
skip_add=False,
refine_convs=0,
)
x_in = tf.keras.Input((8, 8, 1))
skip_src = tf.keras.Input((16, 16, 4))
x = block.make_block(x_in, skip_source=skip_src)
model = tf.keras.Model([x_in, skip_src], x)
self.assertEqual(len(model.layers), 4)
self.assertEqual(len(model.trainable_weights), 0)
self.assertEqual(model.count_params(), 0)
self.assertAllEqual(model.output.shape, (None, 16, 16, 5))
self.assertIsInstance(model.layers[2], tf.keras.layers.UpSampling2D)
self.assertIsInstance(model.layers[3], tf.keras.layers.Concatenate)
def test_simple_upsampling_block_skip_add(self):
block = encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv=False,
interp_method="bilinear",
skip_connection=True,
skip_add=True,
refine_convs=0,
)
x_in = tf.keras.Input((8, 8, 1))
skip_src = tf.ones((1, 16, 16, 1))
x = block.make_block(x_in, skip_source=skip_src)
model = tf.keras.Model(x_in, x)
self.assertEqual(len(model.layers), 3)
self.assertEqual(len(model.trainable_weights), 0)
self.assertEqual(model.count_params(), 0)
self.assertAllEqual(model.output.shape, (None, 16, 16, 1))
self.assertIsInstance(model.layers[1], tf.keras.layers.UpSampling2D)
self.assertTrue("add"
in model.layers[2].name.lower()) # tf_op_layer_AddV2
self.assertAllClose(model(tf.ones((1, 8, 8, 1))),
tf.ones((1, 16, 16, 1)) * 2)
def test_simple_upsampling_block_trans_conv_bn_post(self):
block = encoder_decoder.SimpleUpsamplingBlock(
upsampling_stride=2,
transposed_conv = True,
transposed_conv_filters = 8,
transposed_conv_kernel_size = 3,
transposed_conv_use_bias = True,
transposed_conv_batch_norm = True,
transposed_conv_batch_norm_before_activation = False,
transposed_conv_activation = "relu",
refine_convs = 0,
)
x_in = tf.keras.Input((8, 8, 1))
x = block.make_block(x_in)
model = tf.keras.Model(x_in, x)
self.assertEqual(len(model.layers), 1 + 3)
self.assertEqual(len(model.trainable_weights), 4)
self.assertEqual(model.count_params(), 112)
self.assertAllEqual(model.output.shape, (None, 16, 16, 8))
self.assertIsInstance(model.layers[1], tf.keras.layers.Conv2DTranspose)
self.assertIsInstance(model.layers[2], tf.keras.layers.Activation)
self.assertIsInstance(model.layers[3], tf.keras.layers.BatchNormalization)