<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2021 The TensorFlow Similarity Authors.
Step1:  TensorFlow Similarity Self-Supervised Learning Hello World
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  Dataset Preperation
The following section
Step3:  Load The Raw Data
We are going to load the CIFAR10 dataset. This dataset is often used in the self-supervised papers, enabling us to reproduce the published results.
Step4:  Create Data Splits
The TensorFlow Datasets' CIFAR10 dataset provides a test and train split. 
* Train
Step5:  Self-Supervised Training Setup
This following section
Step6:  Training Parameter Setup
Step8:  Augmented View Configuration
Self-supervised networks require at least two augmented "views" of each example. This can be created using a DataSet and an augmentation function. The DataSet treats each example in the batch as its own class and then the augment function produces two separate views for each example. 
This means the resulting batch will yield tuples containing the two views, i.e., Tuple[(BATCH_SIZE, 32, 32, 3), (BATCH_SIZE, 32, 32, 3)]. TODO make this clearer, "each example in the batch is of this type"
TensorFlow Similarity provides several random augmentation functions, and here we combine augmenters from the simCLR module to replicate the augmentations used in simsiam.
Step9:  Visualize Augmentations
The following cell plots the pairs of augmented views side by side. This can be a useful sanity check as many augmentation functions are set up for the larger ImageNet examples and can be overly aggressive for smaller images found in CIFAR.
Step10:  Contrastive Model Setup
The following section creates the sub-models used by the different algorithms. There are various architectures for building self-supervised models which may include some of the following
Step11:  Projector Model
This MLP is common to all the self-supervised models and is typically a stack of 3 layers of the same size. However, SimSiam only uses 2 layers for the smaller CIFAR images. Having too much capacity in the models can make it difficult for the loss to stabilize and converge.
Additionally, the SimSiam paper found that disabling the center and scale parameters can lead to a small boost in the final loss.
<div class="alert alert-block alert-info"> <b>NOTE</b> This is the model output that is returned by `ContrastiveModel.predict()` and represents the distance based embedding. This embedding can be used for the KNN lookups and matching classification metrics. However, when using the pre-train model for downstream tasks, only the `ContrastiveModel.backbone` is used. </div>
Step12:  Predictor model
The predictor model is used by BYOL and SimSiam, and is an additional 2 layer MLP containing a bottleneck in the hidden layer.
Step13:  Self-Supervised Algorithms
The following section builds the ContrastiveModel based on the ALGORITHM set at the start of the Notebook.
The model training is very sensitive to the learning rate decay and weight decay.
* SimSiam
Step14:  Compile
Step15:  Callbacks
We track the training using several callbacks.
EvalCallback creates an index at the end of each epoch and provides a proxy for the nearest neighbor matching classification using binary_accuracy.
TensordBoard and ModelCheckpoint are provided for tracking the training progress.
Step16:  Model Training
The model training provides a number of metrics.
- loss
Step17:  Save and Reload
The ContrastiveModel contains a set of sub-models and custom train and test steps. Consequently, the ContrastiveModel implements a custom save function that performs the following
Step18:  Evaluation
This final section trains two different classifiers.
No Pre-training
Step19:  No Pretrain
Step20:  Pretrained
Step21:  Comparison
<ASSISTANT_TASK:>
Python Code:
# @title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2021 The TensorFlow Similarity Authors.
End of explanation
import gc
import os
import random
import time
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from tabulate import tabulate
# INFO messages are not printed.
# This must be run before loading other modules.
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"
# install TF addons if needed
try:
    import tensorflow_addons as tfa  # main package
except ModuleNotFoundError:
    !pip install tensorflow-addons
    import tensorflow_addons as tfa
# install TF similarity if needed
try:
    import tensorflow_similarity as tfsim  # main package
except ModuleNotFoundError:
    !pip install tensorflow_similarity
    import tensorflow_similarity as tfsim
import tensorflow as tf
import tensorflow_datasets as tfds
tfsim.utils.tf_cap_memory()  # Avoid GPU memory blow up
# Clear out any old model state.
gc.collect()
tf.keras.backend.clear_session()
print("TensorFlow:", tf.__version__)
print("TensorFlow Similarity", tfsim.__version__)
Explanation: TensorFlow Similarity Self-Supervised Learning Hello World
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/similarity/blob/master/examples/unsupervised_hello_world.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/similarity/blob/master/examples/unsupervised_hello_world.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>
TensorFlow Similarity is a python package focused on making similarity learning and self-supervised learning quick and easy.
Self-supervised learning is an approach to pre-training models using unlabeled data. This approach drastically increases accuracy when you have very few labeled examples but a lot of unlabelled data. The key insight is that you can train a self-supervised model to learn data representations by contrasting multiple augmented views of the same example. These learned representations capture data invariants, e.g., object translation, color jitter, noise, etc. Training a simple linear classifier on top of the frozen representations is easier and requires fewer labels because the pre-trained model already produces meaningful and generally useful features. Overall, self-supervised pre-training learns representations which are more generic and robust than other approaches to augmented training and pre-training. 
Tensorflow Similarity provides a set of network architectures, losses, and data augmentations that are common across a number of self-supervised learning techniques. The Tensorflow Similarity package attempts to provide a consistent terminology across these techniques; however, this leads to slightly different naming conventions, as many papers use different terms to describe the same components. 
The main terms used by Tensorflow Similarity are:
View: A view represents an augmented example.
Backbone: Refers to the model that learns the Representation that we will use for downstream tasks.
Projector: Is an MLP model that projects the backbone representation of a view to an Embedding that is contrasted with other views using a specialized contrastive loss.
Predictor: Is an optional MLP model that is used, in conjunction with gradient stopping, in some recent architectures to further improve the representation quality.
Stop Gradient: Is used by some algorithms to ensure that we only propagate the update from the main view and not the contrasting view.
Notebook goal
This notebook demonstrates how to use Tensorflow Similarity to boost classification accuracy by pre-training a ResNet18 model using contrastive learning on the cifar10 dataset. As you will see, the pre-trained model achieves about ~1.8x the accuracy of the model trained without pre-training. For example, using SimSiam pre-training, you can achieve 87%-91% accuracy versus 50% accuracy when training the same architecture from scratch.
To do this, this notebook walks through how to:
Create a tf.data.Dataset that will generate two augmented views for each example in a batch.
Create basic versions of Backbone, Projector, and Predictor networks to construct the ContrastiveModel()
Train a ContrastiveModel() using one of the following contrastive learning algorithms: SimCLR, SimSiam, and Barlow Twins.
Compare the classification performance of the pre-trained model versus a model trained from scratch.
Things to try
You can explore the following things that will affect the model performance:
- Try different self-supervised learning algorithms.
- Pre-train on a different dataset, e.g., cifar100.
- Restrict supervised training to use a subset of the data, e.g., 100 labels per class.
- Use a larger embedding by increasing the size of the projection and predictor layers.
- Try using different augmentation profiles.
End of explanation
DATA_PATH = Path("tfsim_contrastive_model")
if not DATA_PATH.exists():
    DATA_PATH.mkdir(parents=True)
Explanation: Dataset Preperation
The following section:
* Loads the CIFAR10 data from tensorflow datasets.
* Creates the train, val, test, and query/index splits.
End of explanation
((x_raw_train, y_raw_train), (x_test, y_test)), ds_info = tfds.load(
    "cifar10",
    split=["train", "test"],
    batch_size=-1,
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)
Explanation: Load The Raw Data
We are going to load the CIFAR10 dataset. This dataset is often used in the self-supervised papers, enabling us to reproduce the published results.
End of explanation
# Compute the indicies for query, index, val, and train splits
query_idxs, index_idxs, val_idxs, train_idxs = [], [], [], []
for cid in range(ds_info.features["label"].num_classes):
    idxs = tf.random.shuffle(tf.where(y_raw_train == cid))
    idxs = tf.reshape(idxs, (-1,))
    query_idxs.extend(idxs[:200])  # 200 query examples per class
    index_idxs.extend(idxs[200:400])  # 200 index examples per class
    val_idxs.extend(idxs[400:500])  # 100 validation examples per class
    train_idxs.extend(idxs[500:])  # The remaining are used for training
random.shuffle(query_idxs)
random.shuffle(index_idxs)
random.shuffle(val_idxs)
random.shuffle(train_idxs)
def create_split(idxs: list) -> tuple:
    x, y = [], []
    for idx in idxs:
        x.append(x_raw_train[int(idx)])
        y.append(y_raw_train[int(idx)])
    return tf.convert_to_tensor(np.array(x)), tf.convert_to_tensor(np.array(y))
x_query, y_query = create_split(query_idxs)
x_index, y_index = create_split(index_idxs)
x_val, y_val = create_split(val_idxs)
x_train, y_train = create_split(train_idxs)
print(
    tabulate(
        [
            ["train", x_train.shape, y_train.shape],
            ["val", x_val.shape, y_val.shape],
            ["query", x_query.shape, y_query.shape],
            ["index", x_index.shape, y_index.shape],
            ["test", x_test.shape, y_test.shape],
        ],
        headers=["Examples", "Labels"],
    )
)
Explanation: Create Data Splits
The TensorFlow Datasets' CIFAR10 dataset provides a test and train split. 
* Train: Used for self-supervised pre-training and for training the classifiers.
* Test: Reserved for the classifier evaluation.
However, we are going to partition the train data into the following additional splits:
Validation: Data used for validation metrics during the pre-training phase. 
Query and Index: Data used to compute matching metrics. The query data is used to retrieve the nearest indexed examples.
In particular, the Query and Index split allows us to track the matching classification performance during training. 
<div class="alert alert-block alert-info"> <b>NOTE </b> An increasing match accuracy is a strong indication that the model is learning useful features, however, it does require that we have labeled data. If a dataset only has a small number of labeled examples, they can be passed as the query and index to help monitor the potential matching classification performance during training.</div>
End of explanation
ALGORITHM = "simsiam"  # @param ["barlow", "simsiam", "simclr"]
Explanation: Self-Supervised Training Setup
This following section:
* Sets the training parameters used for building the different architectures.
* Creates a train and test tf.data.Dataset.
* Visualizes the pairs of augmented views from a single batch.
Self-Supervised Algorithm Selection
TensorFlow Similarity currently supports three different self-supervised models.
* SimCLR: Only requires the Backbone and the projector and uses a contrastive cross-entropy loss.
* SimSiam: Requires the Backbone, projector, and predictor and only compares the cosine distance between augmented views from the same example.
* Barlow Twins: Only requires the Backbone and the projector and uses a loss that compares the feature covariance instead of contrasting the views.
The ALGORITHM parameter is used throughout this notebook to set up the various architectures and the parameters defined below are set up to reproduce the results published in the papers.
End of explanation
CIFAR_IMG_SIZE = 32
if ALGORITHM == "simsiam":
    BATCH_SIZE = 512
    PRE_TRAIN_EPOCHS = 800
    PRE_TRAIN_STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE
    VAL_STEPS_PER_EPOCH = 20
    WEIGHT_DECAY = 5e-4
    INIT_LR = 3e-2 * int(BATCH_SIZE / 256)
    WARMUP_LR = 0.0
    WARMUP_STEPS = 0
    DIM = 2048  # The layer size for the projector and predictor models.
elif ALGORITHM == "barlow":
    BATCH_SIZE = 512
    PRE_TRAIN_EPOCHS = 800
    PRE_TRAIN_STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE
    VAL_STEPS_PER_EPOCH = 20
    WEIGHT_DECAY = 5e-4
    INIT_LR = 1e-3  # Initial LR for the learning rate schedule.
    WARMUP_LR = 0.0
    WARMUP_STEPS = 1000
    DIM = 2048  # The layer size for the projector and predictor models.
elif ALGORITHM == "simclr":
    BATCH_SIZE = 512
    PRE_TRAIN_EPOCHS = 800
    PRE_TRAIN_STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE
    VAL_STEPS_PER_EPOCH = 20
    WEIGHT_DECAY = 5e-4
    INIT_LR = 1e-3  # Initial LR for the learning rate schedule, see section B.1 in the paper.
    TEMPERATURE = 0.5  # Tuned for CIFAR10, see section B.9 in the paper.
    DIM = 2048  # The layer size for the projector and predictor models.
Explanation: Training Parameter Setup
End of explanation
def img_scaling(img):
    return tf.keras.applications.imagenet_utils.preprocess_input(
        img, 
        data_format=None, 
        mode='torch')
@tf.function
def simsiam_augmenter(img, blur=True, area_range=(0.2, 1.0)):
    SimSiam augmenter.
    The SimSiam augmentations are based on the SimCLR augmentations, but have
    some important differences.
    * The crop area lower bound is 20% instead of 8%.
    * The color jitter and grayscale are applied separately instead of together.
    * The color jitter ranges are much smaller.
    * Blur is not applied for the cifar10 dataset.
    args:
        img: Single image tensor of shape (H, W, C)
        blur: If true, apply blur. Should be disabled for cifar10.
        area_range: The upper and lower bound of the random crop percentage.
    returns:
        A single image tensor of shape (H, W, C) with values between 0.0 and 1.0.
    
    # random resize and crop. Increase the size before we crop.
    img = tfsim.augmenters.augmentation_utils.cropping.crop_and_resize(
        img, CIFAR_IMG_SIZE, CIFAR_IMG_SIZE, area_range=area_range
    )
    
    # The following transforms expect the data to be [0, 1]
    img /= 255.
    
    # random color jitter
    def _jitter_transform(x):
        return tfsim.augmenters.augmentation_utils.color_jitter.color_jitter_rand(
            x,
            np.random.uniform(0.0, 0.4),
            np.random.uniform(0.0, 0.4),
            np.random.uniform(0.0, 0.4),
            np.random.uniform(0.0, 0.1),
            "multiplicative",
        )
    img = tfsim.augmenters.augmentation_utils.random_apply.random_apply(_jitter_transform, p=0.8, x=img)
    # # random grayscale
    def _grascayle_transform(x):
        return tfsim.augmenters.augmentation_utils.color_jitter.to_grayscale(x)
    img = tfsim.augmenters.augmentation_utils.random_apply.random_apply(_grascayle_transform, p=0.2, x=img)
    # optional random gaussian blur
    if blur:
        img = tfsim.augmenters.augmentation_utils.blur.random_blur(img, p=0.5)
    # random horizontal flip
    img = tf.image.random_flip_left_right(img)
    
    # scale the data back to [0, 255]
    img = img * 255.
    img = tf.clip_by_value(img, 0., 255.)
    return img
@tf.function()
def process(img):
    view1 = simsiam_augmenter(img, blur=False)
    view1 = img_scaling(view1)
    view2 = simsiam_augmenter(img, blur=False)
    view2 = img_scaling(view2)
    return (view1, view2)
train_ds = tf.data.Dataset.from_tensor_slices(x_train)
train_ds = train_ds.repeat()
train_ds = train_ds.shuffle(1024)
train_ds = train_ds.map(process, num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.batch(BATCH_SIZE)
train_ds = train_ds.prefetch(tf.data.AUTOTUNE)
val_ds = tf.data.Dataset.from_tensor_slices(x_val)
val_ds = val_ds.repeat()
val_ds = val_ds.shuffle(1024)
val_ds = val_ds.map(process, num_parallel_calls=tf.data.AUTOTUNE)
val_ds = val_ds.batch(BATCH_SIZE)
val_ds = val_ds.prefetch(tf.data.AUTOTUNE)
Explanation: Augmented View Configuration
Self-supervised networks require at least two augmented "views" of each example. This can be created using a DataSet and an augmentation function. The DataSet treats each example in the batch as its own class and then the augment function produces two separate views for each example. 
This means the resulting batch will yield tuples containing the two views, i.e., Tuple[(BATCH_SIZE, 32, 32, 3), (BATCH_SIZE, 32, 32, 3)]. TODO make this clearer, "each example in the batch is of this type"
TensorFlow Similarity provides several random augmentation functions, and here we combine augmenters from the simCLR module to replicate the augmentations used in simsiam.
End of explanation
display_imgs = next(train_ds.as_numpy_iterator())
max_pixel = np.max([display_imgs[0].max(), display_imgs[1].max()])
min_pixel = np.min([display_imgs[0].min(), display_imgs[1].min()])
tfsim.visualization.visualize_views(
    views=display_imgs,
    num_imgs=16,
    views_per_col=8,
    max_pixel_value=max_pixel,
    min_pixel_value=min_pixel,
)
Explanation: Visualize Augmentations
The following cell plots the pairs of augmented views side by side. This can be a useful sanity check as many augmentation functions are set up for the larger ImageNet examples and can be overly aggressive for smaller images found in CIFAR.
End of explanation
def get_backbone(img_size, activation="relu", preproc_mode="torch"):
    input_shape = (img_size, img_size, 3)
    backbone = tfsim.architectures.ResNet18Sim(
        input_shape,
        include_top=False,  # Take the pooling layer as the output.
        pooling="avg",
    )
    return backbone
backbone = get_backbone(CIFAR_IMG_SIZE)
backbone.summary()
Explanation: Contrastive Model Setup
The following section creates the sub-models used by the different algorithms. There are various architectures for building self-supervised models which may include some of the following:
* Backbone: This is the base model and is typically an existing architecture like ResNet or EfficientNet.
* Projector: This is a small multi-layer Neural Net and provides the embedding features at the end of training.
* Predictor: This model is used by BYOL and SimSiam and provides an additional small multi-layer Neural Net.
Typically, the projector and predictor networks are only 2 or 3 layers with batch normalization. While some papers claim that batch normalization is not required, the SimSiam paper shows empirically that the batch normalization is critical to prevent the model from collapsing to a degenerate solution.
Additionally, many papers show a single encoder block, but this often contains both the Backbone and the Projector network.
The diagram above shows three self-supervised architectures supported by TensorFlow Similarity. As you can see, they all share a common structure:
* Processing multiple views of the same example.
* Using a backbone model for learning the representation output. 
* Using a projector for the embedding output.
* Additionally, note that the loss is symmetric, so we compute it twice during each step. First for view 1 and then a second time for view 2. These two losses are then summed up to compute the final aggregate loss.
However, there are some key differences including:
The loss is different in each of the three architectures.
SimCLR uses a contrastive cross-entropy loss where each example's pair of augmented views should be close to 1 and all other views should be close to zero. Including the negative examples in the loss is part of the reason that SimCLR benefits from larger batch sizes.
SimSiam only compares the cosine distance between the pairs of augmented views. This means there are no negative examples in the loss, and consequently, the batch size can be much smaller.
Barlow Twins attempts to learn the cross-correlation between the embedding features, where similar feature activations between two views will produce small values along the diagonal and a positive penalty is applied for positive or negative off-diagonal correlations. In this way the diagonal represents features that are invariant to augmented views, while the off-diagonal values encourage the features to be independent of each other.
The SimSiam architecture requires an additional Predictor network.
The SimSiam architecture also uses a stop-gradient to ensure that we only flow the gradient from the target view.
TensorFlow Similarity offers a common set of modules to construct the common parts, and implement the parts which are specific to popular self-supervised algorithms on top of it. Attempting to offer flexibility to add your own algorithm or reuse proven architectures that the results we have reproduced. 
Backbone Model
The backbone uses a custom version of ResNet18 in order to reproduce the SimSiam CIFAR10 results.
<div class="alert alert-block alert-info"> <b>NOTE</b> The ResNet models provided in `tf.keras.applications` use larger `[(1x1), (3x3), (1x1)]` blocks that can't be used to reproduce the SimSiam CIFAR10 results. </div>
End of explanation
def get_projector(input_dim, dim, activation="relu", num_layers: int = 3):
    inputs = tf.keras.layers.Input((input_dim,), name="projector_input")
    x = inputs
    for i in range(num_layers - 1):
        x = tf.keras.layers.Dense(
            dim,
            use_bias=False,
            kernel_initializer=tf.keras.initializers.LecunUniform(),
            name=f"projector_layer_{i}",
        )(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1.001e-5, name=f"batch_normalization_{i}")(x)
        x = tf.keras.layers.Activation(activation, name=f"{activation}_activation_{i}")(x)
    x = tf.keras.layers.Dense(
        dim,
        use_bias=False,
        kernel_initializer=tf.keras.initializers.LecunUniform(),
        name="projector_output",
    )(x)
    x = tf.keras.layers.BatchNormalization(
        epsilon=1.001e-5,
        center=False,  # Page:5, Paragraph:2 of SimSiam paper
        scale=False,  # Page:5, Paragraph:2 of SimSiam paper
        name=f"batch_normalization_ouput",
    )(x)
    # Metric Logging layer. Monitors the std of the layer activations.
    # Degnerate solutions colapse to 0 while valid solutions will move
    # towards something like 0.0220. The actual number will depend on the layer size.
    o = tfsim.layers.ActivationStdLoggingLayer(name="proj_std")(x)
    projector = tf.keras.Model(inputs, o, name="projector")
    return projector
projector = get_projector(input_dim=backbone.output.shape[-1], dim=DIM, num_layers=2)
projector.summary()
Explanation: Projector Model
This MLP is common to all the self-supervised models and is typically a stack of 3 layers of the same size. However, SimSiam only uses 2 layers for the smaller CIFAR images. Having too much capacity in the models can make it difficult for the loss to stabilize and converge.
Additionally, the SimSiam paper found that disabling the center and scale parameters can lead to a small boost in the final loss.
<div class="alert alert-block alert-info"> <b>NOTE</b> This is the model output that is returned by `ContrastiveModel.predict()` and represents the distance based embedding. This embedding can be used for the KNN lookups and matching classification metrics. However, when using the pre-train model for downstream tasks, only the `ContrastiveModel.backbone` is used. </div>
End of explanation
def get_predictor(input_dim, hidden_dim=512, activation="relu"):
    inputs = tf.keras.layers.Input(shape=(input_dim,), name="predictor_input")
    x = inputs
    x = tf.keras.layers.Dense(
        hidden_dim,
        use_bias=False,
        kernel_initializer=tf.keras.initializers.LecunUniform(),
        name="predictor_layer_0",
    )(x)
    x = tf.keras.layers.BatchNormalization(epsilon=1.001e-5, name="batch_normalization_0")(x)
    x = tf.keras.layers.Activation(activation, name=f"{activation}_activation_0")(x)
    x = tf.keras.layers.Dense(
        input_dim,
        kernel_initializer=tf.keras.initializers.LecunUniform(),
        name="predictor_output",
    )(x)
    # Metric Logging layer. Monitors the std of the layer activations.
    # Degnerate solutions colapse to 0 while valid solutions will move
    # towards something like 0.0220. The actual number will depend on the layer size.
    o = tfsim.layers.ActivationStdLoggingLayer(name="pred_std")(x)
    predictor = tf.keras.Model(inputs, o, name="predictor")
    return predictor
predictor = get_predictor(input_dim=DIM, hidden_dim=512)
predictor.summary()
Explanation: Predictor model
The predictor model is used by BYOL and SimSiam, and is an additional 2 layer MLP containing a bottleneck in the hidden layer.
End of explanation
if ALGORITHM == "simsiam":
    loss = tfsim.losses.SimSiamLoss(projection_type="cosine_distance", name=ALGORITHM)
    contrastive_model = tfsim.models.ContrastiveModel(
        backbone=backbone,
        projector=projector,
        predictor=predictor,  # NOTE: simiam requires predictor model.
        algorithm=ALGORITHM,
        name=ALGORITHM,
    )
    lr_decayed_fn = tf.keras.optimizers.schedules.CosineDecay(
        initial_learning_rate=INIT_LR,
        decay_steps=PRE_TRAIN_EPOCHS * PRE_TRAIN_STEPS_PER_EPOCH,
    )
    wd_decayed_fn = tf.keras.optimizers.schedules.CosineDecay(
        initial_learning_rate=WEIGHT_DECAY,
        decay_steps=PRE_TRAIN_EPOCHS * PRE_TRAIN_STEPS_PER_EPOCH,
    )
    optimizer = tfa.optimizers.SGDW(learning_rate=lr_decayed_fn, weight_decay=wd_decayed_fn, momentum=0.9)
elif ALGORITHM == "barlow":
    loss = tfsim.losses.Barlow(name=ALGORITHM)
    contrastive_model = tfsim.models.ContrastiveModel(
        backbone=backbone,
        projector=projector,
        algorithm="barlow",
    )
    optimizer = tfa.optimizers.LAMB(learning_rate=INIT_LR)
elif ALGORITHM == "simclr":
    loss = tfsim.losses.SimCLRLoss(name=ALGORITHM, temperature=TEMPERATURE)
    contrastive_model = tfsim.models.ContrastiveModel(
        backbone=backbone,
        projector=projector,
        algorithm=ALGORITHM,
    )
    optimizer = tfa.optimizers.LAMB(learning_rate=INIT_LR)
else:
    raise ValueError(f"{ALGORITHM} is not supported.")
Explanation: Self-Supervised Algorithms
The following section builds the ContrastiveModel based on the ALGORITHM set at the start of the Notebook.
The model training is very sensitive to the learning rate decay and weight decay.
* SimSiam: Requires using SGD with weight decay from TF Addons. Adding weight decay as a kernel_regularizer doesn't seem to be able to reproduce the published results in the paper.
* Barlow Twins: We can use LAMB and avoid the need for the learning rate schedule. Lamb is similar to the LARS optimizer used in the Barlow paper, but includes the use of ADAM. Alternatively, we can use SGD but the optimizer requires a warm up period, otherwise the loss explodes.
* SimCLR: We can also use LAMB as the original paper uses LARS. However, LAMB seems to require smaller learning rates than shown for LARS in the original paper.
End of explanation
contrastive_model.compile(
    optimizer=optimizer,
    loss=loss,
)
Explanation: Compile
End of explanation
log_dir = DATA_PATH / "models" / "logs" / f"{loss.name}_{time.time()}"
chkpt_dir = DATA_PATH / "models" / "checkpoints" / f"{loss.name}_{time.time()}"
evb = tfsim.callbacks.EvalCallback(
    img_scaling(tf.cast(x_query, tf.float32)),
    y_query,
    img_scaling(tf.cast(x_index, tf.float32)),
    y_index,
    metrics=["binary_accuracy"],
    k=1,
    tb_logdir=log_dir,
)
tbc = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    histogram_freq=1,
    update_freq=100,
)
mcp = tf.keras.callbacks.ModelCheckpoint(
    filepath=chkpt_dir,
    monitor="loss",
    mode="min",
    save_best_only=True,
    save_weights_only=True,
)
Explanation: Callbacks
We track the training using several callbacks.
EvalCallback creates an index at the end of each epoch and provides a proxy for the nearest neighbor matching classification using binary_accuracy.
TensordBoard and ModelCheckpoint are provided for tracking the training progress.
End of explanation
history = contrastive_model.fit(
    train_ds,
    epochs=PRE_TRAIN_EPOCHS,
    steps_per_epoch=PRE_TRAIN_STEPS_PER_EPOCH,
    validation_data=val_ds,
    validation_steps=VAL_STEPS_PER_EPOCH,
    callbacks=[evb, tbc, mcp],
)
plt.figure(figsize=(15, 4))
plt.subplot(1, 3, 1)
plt.plot(history.history["loss"])
plt.grid()
plt.title(f"{loss.name} - loss")
plt.subplot(1, 3, 2)
plt.plot(history.history["proj_std"], label="proj")
if "pred_std" in history.history:
    plt.plot(history.history["pred_std"], label="pred")
plt.grid()
plt.title(f"{loss.name} - std metrics")
plt.legend()
plt.subplot(1, 3, 3)
plt.plot(history.history["binary_accuracy"], label="acc")
plt.grid()
plt.title(f"{loss.name} - match metrics")
plt.legend()
plt.show()
Explanation: Model Training
The model training provides a number of metrics.
- loss: This represents the total loss over the contrastive batch. Separate contrastive and regularization losses will also be shown if there are trainable variables in the model layers.
- proj_std and pred_std: These are added as metric logging layers in the model and show the std of the activations of the final layer in the projector or predictor models.
- binary_accuracy: This is the nearest neighbor matching classification accuracy. A new index is built at the end of each epoch and the accuracy is computed using the query and index examples.
Notes: Move this to the evaluation section in the paper.
- SGDW or LAMB seems to be required to prevent the loss from becoming unstable.
- Per layer kernel_regularization with weight decay doesn't work.
- The kernel initialization seems very critical. Smaller, uniform initialization schemes like Lecunn seem to work best.
- The projector is a more stable output for the KNN match metrics. The predictor output is typically worse and higher variance.
- The LR and weight decay schedule seem important. The loss becomes unstable if the updates are too large later in training.
- Too much capacity in the combined model and the training won't converge... same goes for too little capacity.
- Augmentation is critical, but too much augmentation and the model won't converge.
End of explanation
# You can also load the trained weights from the model checkpoints
# contrastive_model.load_weights(DATA_PATH / 'models' / 'checkpoints' / f"{loss.name}_{CHECKPOINT_TIMESTAMP}")
contrastive_model.save(DATA_PATH / "models" / f"trained_model")
del contrastive_model
contrastive_model = tfsim.models.contrastive_model.load_model(DATA_PATH / "models" / f"trained_model")
Explanation: Save and Reload
The ContrastiveModel contains a set of sub-models and custom train and test steps. Consequently, the ContrastiveModel implements a custom save function that performs the following:
* Saves each of the sub models, including the predictor if one exists.
* A JSON file containing the serialized Loss, Metrics, and Optimizer.
* The Optimizer weighs as a npy file.
We also provide a custom load_model() function that restores the ContrastiveModel and supports continued training.
End of explanation
TEST_EPOCHS = 10
TEST_STEPS_PER_EPOCH = int(len(x_train) // BATCH_SIZE)
@tf.function
def eval_augmenter(img):
    # random resize and crop. Increase the size before we crop.
    img = tfsim.augmenters.augmentation_utils.cropping.crop_and_resize(
        img, CIFAR_IMG_SIZE, CIFAR_IMG_SIZE, area_range=(0.2, 1.0)
    )
    # random horizontal flip
    img = tf.image.random_flip_left_right(img)
    img = tf.clip_by_value(img, 0., 255.)
    return img
eval_train_ds = tf.data.Dataset.from_tensor_slices((x_train, tf.keras.utils.to_categorical(y_train, 10)))
eval_train_ds = eval_train_ds.repeat()
eval_train_ds = eval_train_ds.shuffle(1024)
eval_train_ds = eval_train_ds.map(lambda x, y: (eval_augmenter(x), y), tf.data.AUTOTUNE)
eval_train_ds = eval_train_ds.map(lambda x, y: (img_scaling(x), y), tf.data.AUTOTUNE)
eval_train_ds = eval_train_ds.batch(BATCH_SIZE)
eval_train_ds = eval_train_ds.prefetch(tf.data.AUTOTUNE)
eval_val_ds = tf.data.Dataset.from_tensor_slices((x_val, tf.keras.utils.to_categorical(y_val, 10)))
eval_val_ds = eval_val_ds.repeat()
eval_val_ds = eval_val_ds.shuffle(1024)
eval_val_ds = eval_val_ds.map(lambda x, y: (img_scaling(tf.cast(x, dtype=tf.float32)), y), tf.data.AUTOTUNE)
eval_val_ds = eval_val_ds.batch(BATCH_SIZE)
eval_val_ds = eval_val_ds.prefetch(tf.data.AUTOTUNE)
eval_test_ds = tf.data.Dataset.from_tensor_slices((x_test, tf.keras.utils.to_categorical(y_test, 10)))
eval_test_ds = eval_test_ds.map(lambda x, y: (img_scaling(tf.cast(x, dtype=tf.float32)), y), tf.data.AUTOTUNE)
eval_test_ds = eval_test_ds.batch(BATCH_SIZE)
eval_test_ds = eval_test_ds.prefetch(tf.data.AUTOTUNE)
def get_eval_model(img_size, backbone, total_steps, trainable=True, lr=1.8):
    backbone.trainable = trainable
    inputs = tf.keras.layers.Input((img_size, img_size, 3), name="eval_input")
    x = backbone(inputs, training=trainable)
    o = tf.keras.layers.Dense(10, activation="softmax")(x)
    model = tf.keras.Model(inputs, o)
    cosine_decayed_lr = tf.keras.experimental.CosineDecay(initial_learning_rate=lr, decay_steps=total_steps)
    opt = tf.keras.optimizers.SGD(cosine_decayed_lr, momentum=0.9)
    model.compile(optimizer=opt, loss="categorical_crossentropy", metrics=["acc"])
    return model
Explanation: Evaluation
This final section trains two different classifiers.
No Pre-training: Uses a ResNet18 model and a simple linear layer.
Pre-trained Uses the frozen pre-trained backbone from the ContrastiveModel and only trains the weights in the linear layer.
The original train data is partitioned into eval_train and eval_val splits and a simplified augmentation is applied to the training data. 
The models are then trained for 10 epochs and the classification accuracy is evaluated on the held out test split.
End of explanation
no_pt_eval_model = get_eval_model(
    img_size=CIFAR_IMG_SIZE,
    backbone=get_backbone(CIFAR_IMG_SIZE, DIM),
    total_steps=TEST_EPOCHS * TEST_STEPS_PER_EPOCH,
    trainable=True,
    lr=1e-3,
)
no_pt_history = no_pt_eval_model.fit(
    eval_train_ds,
    batch_size=BATCH_SIZE,
    epochs=TEST_EPOCHS,
    steps_per_epoch=TEST_STEPS_PER_EPOCH,
    validation_data=eval_val_ds,
    validation_steps=VAL_STEPS_PER_EPOCH,
)
Explanation: No Pretrain
End of explanation
pt_eval_model = get_eval_model(
    img_size=CIFAR_IMG_SIZE,
    backbone=contrastive_model.backbone,
    total_steps=TEST_EPOCHS * TEST_STEPS_PER_EPOCH,
    trainable=False,
    lr=30.0,
)
pt_eval_model.summary()
pt_history = pt_eval_model.fit(
    eval_train_ds,
    batch_size=BATCH_SIZE,
    epochs=TEST_EPOCHS,
    steps_per_epoch=TEST_STEPS_PER_EPOCH,
    validation_data=eval_val_ds,
    validation_steps=VAL_STEPS_PER_EPOCH,
)
Explanation: Pretrained
End of explanation
print("no pretrain", no_pt_eval_model.evaluate(eval_test_ds))
print("pretrained", pt_eval_model.evaluate(eval_test_ds))
Explanation: Comparison
End of explanation
<END_TASK>