(reparameterizable_model,
                                               functional_ckpt_vars)
return itertools.chain(
      zip(new_conv_kernel_vars, ckpt_conv_kernel_vars),
      zip(new_conv_bias_vars, ckpt_conv_bias_vars),
      zip(new_bn_scale_vars, ckpt_bn_scale_vars),
      zip(new_bn_offset_vars, ckpt_bn_offset_vars),
      zip(new_projection_vars, ckpt_projection_vars),
      zip(new_input_layer_vars, ckpt_input_layer_vars),
      zip(new_output_layer_vars, ckpt_output_layer_vars),
  )
Generate pairs of reparameterizable and functional variables.
Generate pairs of reparameterizable and functional variables.
Generate pairs of reparameterizable and functional variables .
def meta_dataset_functional_variable_generator(reparameterizable_model,
                                               functional_ckpt_vars):
  """Generate pairs of reparameterizable and functional variables."""
  new_vars = reparameterizable_model.trainable_variables

  # The previous meta-dataset implementation uses a projection layer in all
  # shortcut connections (corresponding to (C) in Table 3 of
  # https://arxiv.org/abs/1512.03385); cf. the TorchVision implementation
  # (https://github.com/pytorch/vision/blob/e61538cba036c42bab23ce8f9d205da9889977ae/torchvision/models/resnet.py#L183),
  # which uses a projection shortcut only if the dimensionality of the input and
  # output of the residual block are not the same (corresponding to (B) in Table
  # 3 of https://arxiv.org/abs/1512.03385).
  # (C) has slightly lower error rate but an increased parameter count over (B).
  # We adapt the old implementation using (C) to (B) by simply removing the
  # parameters of the shortcut projection layers unused by (B).
  ckpt_projection_vars = filter_and_sort_vars(functional_ckpt_vars,
                                              'projection')
  new_projection_vars = filter_and_sort_vars(new_vars, 'shortcut')
  functional_ckpt_vars = [
      x for x in functional_ckpt_vars if x not in ckpt_projection_vars
  ]
  new_vars = [x for x in new_vars if x not in new_projection_vars]

  if len(ckpt_projection_vars) != len(new_projection_vars):
    # pylint: disable=g-complex-comprehension
    ckpt_projection_vars = [
        x for x in ckpt_projection_vars if
        # TODO(eringrant): Remove this hack that is specific
        # to the WideResNet16V2 architecture from
        # meta-dataset-V2.
        ('wide_resnet' in x and 'block_0' in x) or
        # TODO(eringrant): Remove this hack that is specific
        # to the ResNet18V1 architecture from
        # meta-dataset-V2.
        ('resnet18' in x and 'conv2_x' not in x)
    ]
    # pylint: enable=g-complex-comprehension

  ckpt_output_layer_vars = filter_and_sort_vars(functional_ckpt_vars,
                                                'embedding')
  new_output_layer_vars = filter_and_sort_vars(new_vars, 'output_stage')
  functional_ckpt_vars = [
      x for x in functional_ckpt_vars if x not in ckpt_output_layer_vars
  ]
  new_vars = [x for x in new_vars if x not in new_output_layer_vars]

  ckpt_input_layer_vars = (
      list(filter_and_sort_vars(functional_ckpt_vars, 'resnet18/conv1')) +
      list(filter_and_sort_vars(functional_ckpt_vars, 'resnet/conv1')))
  new_input_layer_vars = filter_and_sort_vars(new_vars, 'input_stage')
  functional_ckpt_vars = [
      x for x in functional_ckpt_vars if x not in ckpt_input_layer_vars
  ]
  new_vars = [x for x in new_vars if x not in new_input_layer_vars]

  ckpt_conv_vars = filter_and_sort_vars(functional_ckpt_vars, '/conv')
  ckpt_conv_kernel_vars = filter_and_sort_vars(ckpt_conv_vars, 'weight')
  ckpt_conv_bias_vars = filter_and_sort_vars(ckpt_conv_vars, 'bias')

  new_conv_vars = filter_and_sort_vars(new_vars, '/conv')
  new_conv_kernel_vars = filter_and_sort_vars(new_conv_vars, 'kernel')
  new_conv_bias_vars = filter_and_sort_vars(new_conv_vars, 'bias')

  ckpt_bn_vars = filter_and_sort_vars(functional_ckpt_vars, 'batch_norm')
  ckpt_bn_scale_vars = filter_and_sort_vars(ckpt_bn_vars, 'scale')
  ckpt_bn_offset_vars = filter_and_sort_vars(ckpt_bn_vars, 'offset')

  new_bn_vars = filter_and_sort_vars(new_vars, 'batch_normalization')
  new_bn_scale_vars = filter_and_sort_vars(new_bn_vars, 'gamma')
  new_bn_offset_vars = filter_and_sort_vars(new_bn_vars, 'beta')

  return itertools.chain(
      zip(new_conv_kernel_vars, ckpt_conv_kernel_vars),
      zip(new_conv_bias_vars, ckpt_conv_bias_vars),
      zip(new_bn_scale_vars, ckpt_bn_scale_vars),
      zip(new_bn_offset_vars, ckpt_bn_offset_vars),
      zip(new_projection_vars, ckpt_projection_vars),
      zip(new_input_layer_vars, ckpt_input_layer_vars),
      zip(new_output_layer_vars, ckpt_output_layer_vars),
  )
def meta_dataset_functional_variable_generator ( reparameterizable_model , functional_ckpt_vars ) : new_vars = reparameterizable_model . trainable_variables # The previous meta-dataset implementation uses a projection layer in all # shortcut connections (corresponding to (C) in Table 3 of # https://arxiv.org/abs/1512.03385); cf. the TorchVision implementation # (https://github.com/pytorch/vision/blob/e61538cba036c42bab23ce8f9d205da9889977ae/torchvision/models/resnet.py#L183), # which uses a projection shortcut only if the dimensionality of the input and # output of the residual block are not the same (corresponding to (B) in Table # 3 of https://arxiv.org/abs/1512.03385). # (C) has slightly lower error rate but an increased parameter count over (B). # We adapt the old implementation using (C) to (B) by simply removing the # parameters of the shortcut projection layers unused by (B). ckpt_projection_vars = filter_and_sort_vars ( functional_ckpt_vars , 'projection' ) new_projection_vars = filter_and_sort_vars ( new_vars , 'shortcut' ) functional_ckpt_vars = [ x for x in functional_ckpt_vars if x not in ckpt_projection_vars ] new_vars = [ x for x in new_vars if x not in new_projection_vars ] if len ( ckpt_projection_vars ) != len ( new_projection_vars ) : # pylint: disable=g-complex-comprehension ckpt_projection_vars = [ x for x in ckpt_projection_vars if # TODO(eringrant): Remove this hack that is specific # to the WideResNet16V2 architecture from # meta-dataset-V2. ( 'wide_resnet' in x and 'block_0' in x ) or # TODO(eringrant): Remove this hack that is specific # to the ResNet18V1 architecture from # meta-dataset-V2. ( 'resnet18' in x and 'conv2_x' not in x ) ] # pylint: enable=g-complex-comprehension ckpt_output_layer_vars = filter_and_sort_vars ( functional_ckpt_vars , 'embedding' ) new_output_layer_vars = filter_and_sort_vars ( new_vars , 'output_stage' ) functional_ckpt_vars = [ x for x in functional_ckpt_vars if x not in ckpt_output_layer_vars ] new_vars = [ x for x in new_vars if x not in new_output_layer_vars ] ckpt_input_layer_vars = ( list ( filter_and_sort_vars ( functional_ckpt_vars , 'resnet18/conv1' ) ) + list ( filter_and_sort_vars ( functional_ckpt_vars , 'resnet/conv1' ) ) ) new_input_layer_vars = filter_and_sort_vars ( new_vars , 'input_stage' ) functional_ckpt_vars = [ x for x in functional_ckpt_vars if x not in ckpt_input_layer_vars ] new_vars = [ x for x in new_vars if x not in new_input_layer_vars ] ckpt_conv_vars = filter_and_sort_vars ( functional_ckpt_vars , '/conv' ) ckpt_conv_kernel_vars = filter_and_sort_vars ( ckpt_conv_vars , 'weight' ) ckpt_conv_bias_vars = filter_and_sort_vars ( ckpt_conv_vars , 'bias' ) new_conv_vars = filter_and_sort_vars ( new_vars , '/conv' ) new_conv_kernel_vars = filter_and_sort_vars ( new_conv_vars , 'kernel' ) new_conv_bias_vars = filter_and_sort_vars ( new_conv_vars , 'bias' ) ckpt_bn_vars = filter_and_sort_vars ( functional_ckpt_vars , 'batch_norm' ) ckpt_bn_scale_vars = filter_and_sort_vars ( ckpt_bn_vars , 'scale' ) ckpt_bn_offset_vars = filter_and_sort_vars ( ckpt_bn_vars , 'offset' ) new_bn_vars = filter_and_sort_vars ( new_vars , 'batch_normalization' ) new_bn_scale_vars = filter_and_sort_vars ( new_bn_vars , 'gamma' ) new_bn_offset_vars = filter_and_sort_vars ( new_bn_vars , 'beta' ) return itertools . chain ( zip ( new_conv_kernel_vars , ckpt_conv_kernel_vars ) , zip ( new_conv_bias_vars , ckpt_conv_bias_vars ) , zip ( new_bn_scale_vars , ckpt_bn_scale_vars ) , zip ( new_bn_offset_vars , ckpt_bn_offset_vars ) , zip ( new_projection_vars , ckpt_projection_vars ) , zip ( new_input_layer_vars , ckpt_input_layer_vars ) , zip ( new_output_layer_vars , ckpt_output_layer_vars ) , )