(simclr_episode_fraction, *episode)
return tf.cond(
      tf.random_uniform([], minval=0, maxval=1,
                        dtype=tf.float32) > simclr_episode_fraction,
      lambda: episode, lambda: convert_to_simclr_episode(*episode))
Convert simclr_episode_fraction of episodes into SimCLR Episodes.
Convert simclr_episode_fraction of episodes into SimCLR Episodes.
Convert simclr_episode_fraction of episodes into SimCLR Episodes .
def add_simclr_episodes(simclr_episode_fraction, *episode):
  """Convert simclr_episode_fraction of episodes into SimCLR Episodes."""

  def convert_to_simclr_episode(support_images=None,
                                support_labels=None,
                                support_class_ids=None,
                                query_images=None,
                                query_labels=None,
                                query_class_ids=None):
    """Convert a single episode into a SimCLR Episode."""

    # If there were k query examples of class c, keep the first k support
    # examples of class c as 'simclr' queries.  We do this by assigning an
    # id for each image in the query set, implemented as label*1e5+x+1, where
    # x is the number of images of the same label with a lower index within
    # the query set.  We do the same for the support set, which gives us a
    # mapping between query and support images which is injective (as long
    # as there's enough support-set images of each class).
    #
    # note: assumes max support label is 10000 - max_images_per_class
    query_idx_within_class = tf.cast(
        tf.equal(query_labels[tf.newaxis, :], query_labels[:, tf.newaxis]),
        tf.int32)
    query_idx_within_class = tf.linalg.diag_part(
        tf.cumsum(query_idx_within_class, axis=1))
    query_uid = query_labels * 10000 + query_idx_within_class
    support_idx_within_class = tf.cast(
        tf.equal(support_labels[tf.newaxis, :], support_labels[:, tf.newaxis]),
        tf.int32)
    support_idx_within_class = tf.linalg.diag_part(
        tf.cumsum(support_idx_within_class, axis=1))
    support_uid = support_labels * 10000 + support_idx_within_class

    # compute which support-set images have matches in the query set, and
    # discard the rest to produce the new query set.
    support_keep = tf.reduce_any(
        tf.equal(support_uid[:, tf.newaxis], query_uid[tf.newaxis, :]), axis=1)
    query_images = tf.boolean_mask(support_images, support_keep)

    support_labels = tf.range(
        tf.shape(support_labels)[0], dtype=support_labels.dtype)
    query_labels = tf.boolean_mask(support_labels, support_keep)
    query_class_ids = tf.boolean_mask(support_class_ids, support_keep)

    # Finally, apply SimCLR augmentation to all images.
    # Note simclr only blurs one image.
    query_images = simclr_augment(query_images, blur=True)
    support_images = simclr_augment(support_images)

    return (support_images, support_labels, support_class_ids, query_images,
            query_labels, query_class_ids)

  return tf.cond(
      tf.random_uniform([], minval=0, maxval=1,
                        dtype=tf.float32) > simclr_episode_fraction,
      lambda: episode, lambda: convert_to_simclr_episode(*episode))
def add_simclr_episodes ( simclr_episode_fraction , * episode ) : def convert_to_simclr_episode ( support_images = None , support_labels = None , support_class_ids = None , query_images = None , query_labels = None , query_class_ids = None ) : """Convert a single episode into a SimCLR Episode.""" # If there were k query examples of class c, keep the first k support # examples of class c as 'simclr' queries.  We do this by assigning an # id for each image in the query set, implemented as label*1e5+x+1, where # x is the number of images of the same label with a lower index within # the query set.  We do the same for the support set, which gives us a # mapping between query and support images which is injective (as long # as there's enough support-set images of each class). # # note: assumes max support label is 10000 - max_images_per_class query_idx_within_class = tf . cast ( tf . equal ( query_labels [ tf . newaxis , : ] , query_labels [ : , tf . newaxis ] ) , tf . int32 ) query_idx_within_class = tf . linalg . diag_part ( tf . cumsum ( query_idx_within_class , axis = 1 ) ) query_uid = query_labels * 10000 + query_idx_within_class support_idx_within_class = tf . cast ( tf . equal ( support_labels [ tf . newaxis , : ] , support_labels [ : , tf . newaxis ] ) , tf . int32 ) support_idx_within_class = tf . linalg . diag_part ( tf . cumsum ( support_idx_within_class , axis = 1 ) ) support_uid = support_labels * 10000 + support_idx_within_class # compute which support-set images have matches in the query set, and # discard the rest to produce the new query set. support_keep = tf . reduce_any ( tf . equal ( support_uid [ : , tf . newaxis ] , query_uid [ tf . newaxis , : ] ) , axis = 1 ) query_images = tf . boolean_mask ( support_images , support_keep ) support_labels = tf . range ( tf . shape ( support_labels ) [ 0 ] , dtype = support_labels . dtype ) query_labels = tf . boolean_mask ( support_labels , support_keep ) query_class_ids = tf . boolean_mask ( support_class_ids , support_keep ) # Finally, apply SimCLR augmentation to all images. # Note simclr only blurs one image. query_images = simclr_augment ( query_images , blur = True ) support_images = simclr_augment ( support_images ) return ( support_images , support_labels , support_class_ids , query_images , query_labels , query_class_ids ) return tf . cond ( tf . random_uniform ( [ ] , minval = 0 , maxval = 1 , dtype = tf . float32 ) > simclr_episode_fraction , lambda : episode , lambda : convert_to_simclr_episode ( * episode ) )