(model, imported_model)
return pairs
Pair variables for reparameterizable and BiT backbones.

  Args:
    model: A valid `reparameterizable_backbones.ReparameterizableBackbone`.
    imported_model: A `tf.AutoTrackable`.

  Returns:
    Pairs of corresponding `tf.Variable`s from `model` and `imported_model`.

  Raises:
    ValueError: If `model` is not a valid model.
Pair variables for reparameterizable and BiT backbones.
Pair variables for reparameterizable and BiT backbones .
def bit_variable_generator(model, imported_model):
  """Pair variables for reparameterizable and BiT backbones.

  Args:
    model: A valid `reparameterizable_backbones.ReparameterizableBackbone`.
    imported_model: A `tf.AutoTrackable`.

  Returns:
    Pairs of corresponding `tf.Variable`s from `model` and `imported_model`.

  Raises:
    ValueError: If `model` is not a valid model.
  """
  valid_models = (
      reparameterizable_backbones.ResNet18V2,
      reparameterizable_backbones.GNResNet18V2,
      reparameterizable_backbones.ResNet50V2,
      reparameterizable_backbones.GNResNet50V2,
  )
  if all(not isinstance(model, valid_model) for valid_model in valid_models):
    raise ValueError(
        'BiT variable mapping defined only for models among: {}'.format(
            valid_models))

  pairs = []
  name_to_variable = dict((v.name, v) for v in model.variables)
  imported_name_to_variable = dict(
      (v.name, v) for v in imported_model.variables)

  group_norm = isinstance(model, reparameterizable_backbones.GNResNet)

  # Assign root layer variables.
  # TODO(eringrant): Determine why 'input_stage' is not prepended to variables
  # in tensorflow Eager mode.
  input_stage_prefix = 'input_stage/' if not tf.executing_eagerly() else ''
  output_stage_prefix = 'output_stage/' if not tf.executing_eagerly() else ''

  name = '{}conv2d/kernel:0'.format(input_stage_prefix)
  imported_name = 'resnet/root_block/{}conv2d/kernel:0'.format(
      'standardized_' if group_norm else '')

  pairs += [(name_to_variable[name], imported_name_to_variable[imported_name])]

  # Assign normalization variables.
  for n in ('gamma', 'beta') + (() if group_norm else
                                ('moving_mean', 'moving_variance')):
    name = '{}{}_normalization_{}/{}:0'.format(
        output_stage_prefix, 'group' if group_norm else 'batch', 16 if
        (isinstance(model, reparameterizable_backbones.ResNet18V2) or
         isinstance(model, reparameterizable_backbones.GNResNet18V2)) else 48,
        n)
    imported_name = 'resnet/{}/{}:0'.format(
        'group_norm' if group_norm else 'batch_normalization', n)
    pairs += [(name_to_variable[name], imported_name_to_variable[imported_name])
             ]

  # Assign other layer variables.
  if (isinstance(model, reparameterizable_backbones.ResNet18V2) or
      isinstance(model, reparameterizable_backbones.GNResNet18V2)):
    for i, (stage, block, layer) in enumerate(
        itertools.product(range(4), range(2), range(2))):

      imported_base_path = 'resnet/block{}/unit0{}/{}/'.format(
          stage + 1, block + 1, {
              0: 'a',
              1: 'b'
          }[layer])

      for n in ('gamma', 'beta') + (() if group_norm else
                                    ('moving_mean', 'moving_variance')):
        name = 'stage_{}/block_{}/{}_normalization{}/{}:0'.format(
            stage, block, 'group' if group_norm else 'batch',
            '_{}'.format(i) if i > 0 else '', n)
        imported_name = (
            imported_base_path + '{}/{}:0'.format(
                'group_norm' if group_norm else 'batch_normalization', n))
        pairs += [(name_to_variable[name],
                   imported_name_to_variable[imported_name])]

      if stage > 0 and block == layer == 0:
        name = 'stage_{}/block_{}/shortcut_conv/kernel:0'.format(stage, block)
        imported_name = (
            imported_base_path + 'proj/{}conv2d/kernel:0'.format(
                'standardized_' if group_norm else ''))
        pairs += [(name_to_variable[name],
                   imported_name_to_variable[imported_name])]

      name = 'stage_{}/block_{}/conv2d_{}/kernel:0'.format(stage, block, i + 1)
      imported_name = (
          imported_base_path +
          '{}conv2d/kernel:0'.format('standardized_' if group_norm else ''))
      pairs += [(name_to_variable[name],
                 imported_name_to_variable[imported_name])]

  elif (isinstance(model, reparameterizable_backbones.ResNet50V2) or
        isinstance(model, reparameterizable_backbones.GNResNet50V2)):
    conv_id = 1
    batch_norm_id = 0
    for stage, num_blocks in enumerate((3, 4, 6, 3)):
      for block in range(num_blocks):
        for layer in range(3):
          imported_base_path = 'resnet/block{}/unit0{}/{}/'.format(
              stage + 1, block + 1, {
                  0: 'a',
                  1: 'b',
                  2: 'c'
              }[layer])

          for n in ('gamma', 'beta') + (() if group_norm else
                                        ('moving_mean', 'moving_variance')):
            name = 'stage_{}/block_{}/{}_normalization{}/{}:0'.format(
                stage, block, 'group' if group_norm else 'batch',
                '_{}'.format(batch_norm_id) if batch_norm_id > 0 else '', n)
            imported_name = (
                imported_base_path + '{}/{}:0'.format(
                    'group_norm' if group_norm else 'batch_normalization', n))
            pairs += [(name_to_variable[name],
                       imported_name_to_variable[imported_name])]
          batch_norm_id += 1

          name = 'stage_{}/block_{}/conv2d_{}/kernel:0'.format(
              stage, block, conv_id)
          imported_name = (
              imported_base_path +
              '{}conv2d/kernel:0'.format('standardized_' if group_norm else ''))
          pairs += [(name_to_variable[name],
                     imported_name_to_variable[imported_name])]
          conv_id += 1

        if block == 0:
          name = 'stage_{}/block_{}/conv2d{}/kernel:0'.format(
              stage, block, '_{}'.format(conv_id) if conv_id > 0 else '')
          imported_name = 'resnet/block{}/unit01/a/proj/{}conv2d/kernel:0'.format(
              stage + 1, 'standardized_' if group_norm else '')
          pairs += [(name_to_variable[name],
                     imported_name_to_variable[imported_name])]
          conv_id += 1

  else:
    raise NotImplementedError('BiT variable mappings are not defined for '
                              'models other than {}.'.format(valid_models))

  return pairs
def bit_variable_generator ( model , imported_model ) : valid_models = ( reparameterizable_backbones . ResNet18V2 , reparameterizable_backbones . GNResNet18V2 , reparameterizable_backbones . ResNet50V2 , reparameterizable_backbones . GNResNet50V2 , ) if all ( not isinstance ( model , valid_model ) for valid_model in valid_models ) : raise ValueError ( 'BiT variable mapping defined only for models among: {}' . format ( valid_models ) ) pairs = [ ] name_to_variable = dict ( ( v . name , v ) for v in model . variables ) imported_name_to_variable = dict ( ( v . name , v ) for v in imported_model . variables ) group_norm = isinstance ( model , reparameterizable_backbones . GNResNet ) # Assign root layer variables. # TODO(eringrant): Determine why 'input_stage' is not prepended to variables # in tensorflow Eager mode. input_stage_prefix = 'input_stage/' if not tf . executing_eagerly ( ) else '' output_stage_prefix = 'output_stage/' if not tf . executing_eagerly ( ) else '' name = '{}conv2d/kernel:0' . format ( input_stage_prefix ) imported_name = 'resnet/root_block/{}conv2d/kernel:0' . format ( 'standardized_' if group_norm else '' ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] # Assign normalization variables. for n in ( 'gamma' , 'beta' ) + ( ( ) if group_norm else ( 'moving_mean' , 'moving_variance' ) ) : name = '{}{}_normalization_{}/{}:0' . format ( output_stage_prefix , 'group' if group_norm else 'batch' , 16 if ( isinstance ( model , reparameterizable_backbones . ResNet18V2 ) or isinstance ( model , reparameterizable_backbones . GNResNet18V2 ) ) else 48 , n ) imported_name = 'resnet/{}/{}:0' . format ( 'group_norm' if group_norm else 'batch_normalization' , n ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] # Assign other layer variables. if ( isinstance ( model , reparameterizable_backbones . ResNet18V2 ) or isinstance ( model , reparameterizable_backbones . GNResNet18V2 ) ) : for i , ( stage , block , layer ) in enumerate ( itertools . product ( range ( 4 ) , range ( 2 ) , range ( 2 ) ) ) : imported_base_path = 'resnet/block{}/unit0{}/{}/' . format ( stage + 1 , block + 1 , { 0 : 'a' , 1 : 'b' } [ layer ] ) for n in ( 'gamma' , 'beta' ) + ( ( ) if group_norm else ( 'moving_mean' , 'moving_variance' ) ) : name = 'stage_{}/block_{}/{}_normalization{}/{}:0' . format ( stage , block , 'group' if group_norm else 'batch' , '_{}' . format ( i ) if i > 0 else '' , n ) imported_name = ( imported_base_path + '{}/{}:0' . format ( 'group_norm' if group_norm else 'batch_normalization' , n ) ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] if stage > 0 and block == layer == 0 : name = 'stage_{}/block_{}/shortcut_conv/kernel:0' . format ( stage , block ) imported_name = ( imported_base_path + 'proj/{}conv2d/kernel:0' . format ( 'standardized_' if group_norm else '' ) ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] name = 'stage_{}/block_{}/conv2d_{}/kernel:0' . format ( stage , block , i + 1 ) imported_name = ( imported_base_path + '{}conv2d/kernel:0' . format ( 'standardized_' if group_norm else '' ) ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] elif ( isinstance ( model , reparameterizable_backbones . ResNet50V2 ) or isinstance ( model , reparameterizable_backbones . GNResNet50V2 ) ) : conv_id = 1 batch_norm_id = 0 for stage , num_blocks in enumerate ( ( 3 , 4 , 6 , 3 ) ) : for block in range ( num_blocks ) : for layer in range ( 3 ) : imported_base_path = 'resnet/block{}/unit0{}/{}/' . format ( stage + 1 , block + 1 , { 0 : 'a' , 1 : 'b' , 2 : 'c' } [ layer ] ) for n in ( 'gamma' , 'beta' ) + ( ( ) if group_norm else ( 'moving_mean' , 'moving_variance' ) ) : name = 'stage_{}/block_{}/{}_normalization{}/{}:0' . format ( stage , block , 'group' if group_norm else 'batch' , '_{}' . format ( batch_norm_id ) if batch_norm_id > 0 else '' , n ) imported_name = ( imported_base_path + '{}/{}:0' . format ( 'group_norm' if group_norm else 'batch_normalization' , n ) ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] batch_norm_id += 1 name = 'stage_{}/block_{}/conv2d_{}/kernel:0' . format ( stage , block , conv_id ) imported_name = ( imported_base_path + '{}conv2d/kernel:0' . format ( 'standardized_' if group_norm else '' ) ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] conv_id += 1 if block == 0 : name = 'stage_{}/block_{}/conv2d{}/kernel:0' . format ( stage , block , '_{}' . format ( conv_id ) if conv_id > 0 else '' ) imported_name = 'resnet/block{}/unit01/a/proj/{}conv2d/kernel:0' . format ( stage + 1 , 'standardized_' if group_norm else '' ) pairs += [ ( name_to_variable [ name ] , imported_name_to_variable [ imported_name ] ) ] conv_id += 1 else : raise NotImplementedError ( 'BiT variable mappings are not defined for ' 'models other than {}.' . format ( valid_models ) ) return pairs