def test_output_shapes(self, groups, input_channel, output_channel): l = custom_layers.GroupConv2D(output_channel, 3, groups=groups, use_bias=False, padding='same') outputs = l(_get_random_inputs(input_shape=[2, 32, 32, input_channel])) self.assertListEqual(outputs.get_shape().as_list(), [2, 32, 32, output_channel])
def test_kernel_shapes(self, groups, input_channel, output_channel): l = custom_layers.GroupConv2D(output_channel, 3, groups=groups, use_bias=False) _ = l(_get_random_inputs(input_shape=(1, 32, 32, input_channel))) expected_kernel_shapes = [(3, 3, int(input_channel / groups), int(output_channel / groups)) for _ in range(groups)] kernel_shapes = [ l.trainable_weights[i].get_shape() for i in range(len(l.trainable_weights)) ] self.assertListEqual(kernel_shapes, expected_kernel_shapes)
def test_construction(self, groups, input_channel, output_channel, use_batch_norm): batch_norm_layer = BATCH_NORM_LAYER if use_batch_norm else None l = custom_layers.GroupConv2D(output_channel, 3, groups=groups, use_bias=True, batch_norm_layer=batch_norm_layer) inputs = _get_random_inputs(input_shape=(1, 4, 4, output_channel)) _ = l(inputs) # kernel and bias for each group. When using batch norm, 2 additional # trainable weights per group for batchnorm layers: gamma and beta. expected_num_trainable_weights = groups * (2 + 2 * use_batch_norm) self.assertLen(l.trainable_weights, expected_num_trainable_weights)
def test_equivalence(self, groups, input_channel, output_channel, use_batch_norm, activation): batch_norm_layer = BATCH_NORM_LAYER if use_batch_norm else None kwargs = dict(filters=output_channel, groups=groups, kernel_size=1, use_bias=False, batch_norm_layer=batch_norm_layer, activation=activation) gc_layer = tf.keras.Sequential([custom_layers.GroupConv2D(**kwargs)]) gc_model = custom_layers.GroupConv2DKerasModel(**kwargs) gc_layer.build(input_shape=(None, 3, 3, input_channel)) gc_model.build(input_shape=(None, 3, 3, input_channel)) inputs = _get_random_inputs((2, 3, 3, input_channel)) gc_layer.set_weights(gc_model.get_weights()) self.assertAllEqual(gc_layer(inputs), gc_model(inputs))
def groupconv2d_block(conv_filters: Optional[int], config: ModelConfig, kernel_size: Any = (1, 1), strides: Any = (1, 1), group_size: Optional[int] = None, use_batch_norm: bool = True, use_bias: bool = False, activation: Any = None, name: Optional[str] = None) -> tf.keras.layers.Layer: """2D group convolution with batchnorm and activation.""" batch_norm = common_modules.get_batch_norm(config.batch_norm) bn_momentum = config.bn_momentum bn_epsilon = config.bn_epsilon data_format = tf.keras.backend.image_data_format() weight_decay = config.weight_decay if group_size is None: group_size = config.group_base_size name = name or '' # Compute the # of groups if conv_filters % group_size != 0: raise ValueError( f'Number of filters: {conv_filters} is not divisible by ' f'size of the groups: {group_size}') groups = int(conv_filters / group_size) # Collect args based on what kind of groupconv2d block is desired init_kwargs = { 'kernel_size': kernel_size, 'strides': strides, 'use_bias': use_bias, 'padding': 'same', 'name': name + '_groupconv2d', 'kernel_regularizer': tf.keras.regularizers.l2(weight_decay), 'bias_regularizer': tf.keras.regularizers.l2(weight_decay), 'filters': conv_filters, 'groups': groups, 'batch_norm_layer': batch_norm if use_batch_norm else None, 'bn_epsilon': bn_epsilon, 'bn_momentum': bn_momentum, 'activation': activation, 'data_format': data_format, } return custom_layers.GroupConv2D(**init_kwargs)
def test_serialization_deserialization(self, groups, use_batch_norm, activation): batch_norm_layer = BATCH_NORM_LAYER if use_batch_norm else None l = custom_layers.GroupConv2D(filters=8, kernel_size=1, groups=groups, use_bias=False, padding='same', batch_norm_layer=batch_norm_layer, activation=activation) config = l.get_config() # New layer from config new_l = custom_layers.GroupConv2D.from_config(config) # Copy the weights too. l.build(input_shape=(1, 1, 4)) new_l.build(input_shape=(1, 1, 4)) new_l.set_weights(l.get_weights()) inputs = _get_random_inputs((1, 1, 1, 4)) self.assertNotEqual(l, new_l) self.assertAllEqual(l(inputs), new_l(inputs))