Esempio n. 1
0
    def __init__(self,
                 feat_level,
                 target_num_channels,
                 apply_bn=False,
                 is_training_bn=None,
                 conv_after_downsample=False,
                 strategy=None,
                 data_format=None,
                 pooling_type=None,
                 upsampling_type=None,
                 model_optimizations=None,
                 name='resample_p0'):
        super().__init__(name=name)
        self.apply_bn = apply_bn
        self.is_training_bn = is_training_bn
        self.data_format = data_format
        self.target_num_channels = target_num_channels
        self.feat_level = feat_level
        self.strategy = strategy
        self.conv_after_downsample = conv_after_downsample
        self.pooling_type = pooling_type or 'max'
        self.upsampling_type = upsampling_type or 'nearest'

        self.conv2d = tf.keras.layers.Conv2D(self.target_num_channels, (1, 1),
                                             padding='same',
                                             data_format=self.data_format,
                                             name='conv2d')
        if model_optimizations:
            for method in model_optimizations.keys():
                self.conv2d = tfmot.get_method(method)(self.conv2d)
        self.bn = util_keras.build_batch_norm(
            is_training_bn=self.is_training_bn,
            data_format=self.data_format,
            strategy=self.strategy,
            name='bn')
  def __init__(self,
               is_training_bn,
               conv_bn_act_pattern,
               separable_conv,
               fpn_num_filters,
               act_type,
               data_format,
               strategy,
               model_optimizations,
               name='op_after_combine'):
    super().__init__(name=name)
    self.conv_bn_act_pattern = conv_bn_act_pattern
    self.separable_conv = separable_conv
    self.fpn_num_filters = fpn_num_filters
    self.act_type = act_type
    self.data_format = data_format
    self.strategy = strategy
    self.is_training_bn = is_training_bn
    if self.separable_conv:
      conv2d_layer = functools.partial(
          tf.keras.layers.SeparableConv2D, depth_multiplier=1)
    else:
      conv2d_layer = tf.keras.layers.Conv2D

    self.conv_op = conv2d_layer(
        filters=fpn_num_filters,
        kernel_size=(3, 3),
        padding='same',
        use_bias=not self.conv_bn_act_pattern,
        data_format=self.data_format,
        name='conv')
    if model_optimizations:
      for method in model_optimizations.keys():
        self.conv_op = (
            tfmot.get_method(method)(self.conv_op))
    self.bn = util_keras.build_batch_norm(
        is_training_bn=self.is_training_bn,
        data_format=self.data_format,
        strategy=self.strategy,
        name='bn')