def test_shape(self, distribution): batch_size, height, width, num_channels = 8, 32, 32, 32 num_filters = 64 strides = 2 input_tensor = tf.random.normal( shape=[batch_size, height, width, num_channels]) with distribution.scope(): block = nn_blocks.DepthwiseSeparableConvBlock( num_filters, strides=strides) config_dict = block.get_config() recreate_block = nn_blocks.DepthwiseSeparableConvBlock(**config_dict) output_tensor = block(input_tensor) expected_output_shape = [ batch_size, height // strides, width // strides, num_filters ] self.assertEqual(output_tensor.shape.as_list(), expected_output_shape) output_tensor = recreate_block(input_tensor) self.assertEqual(output_tensor.shape.as_list(), expected_output_shape)
def _mobilenet_base( self, inputs: tf.Tensor) -> Tuple[tf.Tensor, Dict[str, tf.Tensor], int]: """Builds the base MobileNet architecture. Args: inputs: A `tf.Tensor` of shape `[batch_size, height, width, channels]`. Returns: A tuple of output Tensor and dictionary that collects endpoints. """ input_shape = inputs.get_shape().as_list() if len(input_shape) != 4: raise ValueError('Expected rank 4 input, was: %d' % len(input_shape)) # The current_stride variable keeps track of the output stride of the # activations, i.e., the running product of convolution strides up to the # current network layer. This allows us to invoke atrous convolution # whenever applying the next convolution would result in the activations # having output stride larger than the target output_stride. current_stride = 1 # The atrous convolution rate parameter. rate = 1 net = inputs endpoints = {} endpoint_level = 2 for i, block_def in enumerate(self._decoded_specs): block_name = 'block_group_{}_{}'.format(block_def.block_fn, i) # A small catch for gpooling block with None strides if not block_def.strides: block_def.strides = 1 if (self._output_stride is not None and current_stride == self._output_stride): # If we have reached the target output_stride, then we need to employ # atrous convolution with stride=1 and multiply the atrous rate by the # current unit's stride for use in subsequent layers. layer_stride = 1 layer_rate = rate rate *= block_def.strides else: layer_stride = block_def.strides layer_rate = 1 current_stride *= block_def.strides if block_def.block_fn == 'convbn': net = Conv2DBNBlock( filters=block_def.filters, kernel_size=block_def.kernel_size, strides=block_def.strides, activation=block_def.activation, use_bias=block_def.use_bias, use_normalization=block_def.use_normalization, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, use_sync_bn=self._use_sync_bn, norm_momentum=self._norm_momentum, norm_epsilon=self._norm_epsilon)(net) elif block_def.block_fn == 'depsepconv': net = nn_blocks.DepthwiseSeparableConvBlock( filters=block_def.filters, kernel_size=block_def.kernel_size, strides=block_def.strides, activation=block_def.activation, dilation_rate=layer_rate, regularize_depthwise=self._regularize_depthwise, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, use_sync_bn=self._use_sync_bn, norm_momentum=self._norm_momentum, norm_epsilon=self._norm_epsilon, )(net) elif block_def.block_fn == 'invertedbottleneck': use_rate = rate if layer_rate > 1 and block_def.kernel_size != 1: # We will apply atrous rate in the following cases: # 1) When kernel_size is not in params, the operation then uses # default kernel size 3x3. # 2) When kernel_size is in params, and if the kernel_size is not # equal to (1, 1) (there is no need to apply atrous convolution to # any 1x1 convolution). use_rate = layer_rate in_filters = net.shape.as_list()[-1] net = nn_blocks.InvertedBottleneckBlock( in_filters=in_filters, out_filters=block_def.filters, kernel_size=block_def.kernel_size, strides=layer_stride, expand_ratio=block_def.expand_ratio, se_ratio=block_def.se_ratio, expand_se_in_filters=True, se_gating_activation='hard_sigmoid', activation=block_def.activation, use_depthwise=block_def.use_depthwise, use_residual=block_def.use_residual, dilation_rate=use_rate, regularize_depthwise=self._regularize_depthwise, kernel_initializer=self._kernel_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, use_sync_bn=self._use_sync_bn, norm_momentum=self._norm_momentum, norm_epsilon=self._norm_epsilon, stochastic_depth_drop_rate=self. _stochastic_depth_drop_rate, divisible_by=self._get_divisible_by())(net) elif block_def.block_fn == 'gpooling': net = layers.GlobalAveragePooling2D()(net) net = layers.Reshape((1, 1, net.shape[1]))(net) else: raise ValueError('Unknown block type {} for layer {}'.format( block_def.block_fn, i)) net = tf.identity(net, name=block_name) if block_def.is_output: endpoints[str(endpoint_level)] = net endpoint_level += 1 return net, endpoints, endpoint_level