CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/vision/masked_image_modeling/" />  <meta property="og:title" content="Keras documentation: Masked image modeling with Autoencoders"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Masked image modeling with Autoencoders"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Masked image modeling with Autoencoders</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink active" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink2" href="/examples/vision/image_classification_from_scratch/">Image classification from scratch</a> <a class="nav-sublink2" href="/examples/vision/mnist_convnet/">Simple MNIST convnet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_efficientnet_fine_tuning/">Image classification via fine-tuning with EfficientNet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_with_vision_transformer/">Image classification with Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/attention_mil_classification/">Classification using Attention-based Deep Multiple Instance Learning</a> <a class="nav-sublink2" href="/examples/vision/mlp_image_classification/">Image classification with modern MLP models</a> <a class="nav-sublink2" href="/examples/vision/mobilevit/">A mobile-friendly Transformer-based model for image classification</a> <a class="nav-sublink2" href="/examples/vision/xray_classification_with_tpus/">Pneumonia Classification on TPU</a> <a class="nav-sublink2" href="/examples/vision/cct/">Compact Convolutional Transformers</a> <a class="nav-sublink2" href="/examples/vision/convmixer/">Image classification with ConvMixer</a> <a class="nav-sublink2" href="/examples/vision/eanet/">Image classification with EANet (External Attention Transformer)</a> <a class="nav-sublink2" href="/examples/vision/involution/">Involutional neural networks</a> <a class="nav-sublink2" href="/examples/vision/perceiver_image_classification/">Image classification with Perceiver</a> <a class="nav-sublink2" href="/examples/vision/reptile/">Few-Shot learning with Reptile</a> <a class="nav-sublink2" href="/examples/vision/semisupervised_simclr/">Semi-supervised image classification using contrastive pretraining with SimCLR</a> <a class="nav-sublink2" href="/examples/vision/swin_transformers/">Image classification with Swin Transformers</a> <a class="nav-sublink2" href="/examples/vision/vit_small_ds/">Train a Vision Transformer on small datasets</a> <a class="nav-sublink2" href="/examples/vision/shiftvit/">A Vision Transformer without Attention</a> <a class="nav-sublink2" href="/examples/vision/image_classification_using_global_context_vision_transformer/">Image Classification using Global Context Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/temporal_latent_bottleneck/">When Recurrence meets Transformers</a> <a class="nav-sublink2" href="/examples/vision/oxford_pets_image_segmentation/">Image segmentation with a U-Net-like architecture</a> <a class="nav-sublink2" href="/examples/vision/deeplabv3_plus/">Multiclass semantic segmentation using DeepLabV3+</a> <a class="nav-sublink2" href="/examples/vision/basnet_segmentation/">Highly accurate boundaries segmentation using BASNet</a> <a class="nav-sublink2" href="/examples/vision/fully_convolutional_network/">Image Segmentation using Composable Fully-Convolutional Networks</a> <a class="nav-sublink2" href="/examples/vision/retinanet/">Object Detection with RetinaNet</a> <a class="nav-sublink2" href="/examples/vision/keypoint_detection/">Keypoint Detection with Transfer Learning</a> <a class="nav-sublink2" href="/examples/vision/object_detection_using_vision_transformer/">Object detection with Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/3D_image_classification/">3D image classification from CT scans</a> <a class="nav-sublink2" href="/examples/vision/depth_estimation/">Monocular depth estimation</a> <a class="nav-sublink2" href="/examples/vision/nerf/">3D volumetric rendering with NeRF</a> <a class="nav-sublink2" href="/examples/vision/pointnet_segmentation/">Point cloud segmentation with PointNet</a> <a class="nav-sublink2" href="/examples/vision/pointnet/">Point cloud classification</a> <a class="nav-sublink2" href="/examples/vision/captcha_ocr/">OCR model for reading Captchas</a> <a class="nav-sublink2" href="/examples/vision/handwriting_recognition/">Handwriting recognition</a> <a class="nav-sublink2" href="/examples/vision/autoencoder/">Convolutional autoencoder for image denoising</a> <a class="nav-sublink2" href="/examples/vision/mirnet/">Low-light image enhancement using MIRNet</a> <a class="nav-sublink2" href="/examples/vision/super_resolution_sub_pixel/">Image Super-Resolution using an Efficient Sub-Pixel CNN</a> <a class="nav-sublink2" href="/examples/vision/edsr/">Enhanced Deep Residual Networks for single-image super-resolution</a> <a class="nav-sublink2" href="/examples/vision/zero_dce/">Zero-DCE for low-light image enhancement</a> <a class="nav-sublink2" href="/examples/vision/cutmix/">CutMix data augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/mixup/">MixUp augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/randaugment/">RandAugment for Image Classification for Improved Robustness</a> <a class="nav-sublink2" href="/examples/vision/image_captioning/">Image captioning</a> <a class="nav-sublink2" href="/examples/vision/nl_image_search/">Natural language image search with a Dual Encoder</a> <a class="nav-sublink2" href="/examples/vision/visualizing_what_convnets_learn/">Visualizing what convnets learn</a> <a class="nav-sublink2" href="/examples/vision/integrated_gradients/">Model interpretability with Integrated Gradients</a> <a class="nav-sublink2" href="/examples/vision/probing_vits/">Investigating Vision Transformer representations</a> <a class="nav-sublink2" href="/examples/vision/grad_cam/">Grad-CAM class activation visualization</a> <a class="nav-sublink2" href="/examples/vision/near_dup_search/">Near-duplicate image search</a> <a class="nav-sublink2" href="/examples/vision/semantic_image_clustering/">Semantic Image Clustering</a> <a class="nav-sublink2" href="/examples/vision/siamese_contrastive/">Image similarity estimation using a Siamese Network with a contrastive loss</a> <a class="nav-sublink2" href="/examples/vision/siamese_network/">Image similarity estimation using a Siamese Network with a triplet loss</a> <a class="nav-sublink2" href="/examples/vision/metric_learning/">Metric learning for image similarity search</a> <a class="nav-sublink2" href="/examples/vision/metric_learning_tf_similarity/">Metric learning for image similarity search using TensorFlow Similarity</a> <a class="nav-sublink2" href="/examples/vision/nnclr/">Self-supervised contrastive learning with NNCLR</a> <a class="nav-sublink2" href="/examples/vision/video_classification/">Video Classification with a CNN-RNN Architecture</a> <a class="nav-sublink2" href="/examples/vision/conv_lstm/">Next-Frame Video Prediction with Convolutional LSTMs</a> <a class="nav-sublink2" href="/examples/vision/video_transformers/">Video Classification with Transformers</a> <a class="nav-sublink2" href="/examples/vision/vivit/">Video Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/bit/">Image Classification using BigTransfer (BiT)</a> <a class="nav-sublink2" href="/examples/vision/gradient_centralization/">Gradient Centralization for Better Training Performance</a> <a class="nav-sublink2" href="/examples/vision/token_learner/">Learning to tokenize in Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/knowledge_distillation/">Knowledge Distillation</a> <a class="nav-sublink2" href="/examples/vision/fixres/">FixRes: Fixing train-test resolution discrepancy</a> <a class="nav-sublink2" href="/examples/vision/cait/">Class Attention Image Transformers with LayerScale</a> <a class="nav-sublink2" href="/examples/vision/patch_convnet/">Augmenting convnets with aggregated attention</a> <a class="nav-sublink2" href="/examples/vision/learnable_resizer/">Learning to Resize</a> <a class="nav-sublink2" href="/examples/vision/adamatch/">Semi-supervision and domain adaptation with AdaMatch</a> <a class="nav-sublink2" href="/examples/vision/barlow_twins/">Barlow Twins for Contrastive SSL</a> <a class="nav-sublink2" href="/examples/vision/consistency_training/">Consistency training with supervision</a> <a class="nav-sublink2" href="/examples/vision/deit/">Distilling Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/focal_modulation_network/">Focal Modulation: A replacement for Self-Attention</a> <a class="nav-sublink2" href="/examples/vision/forwardforward/">Using the Forward-Forward Algorithm for Image Classification</a> <a class="nav-sublink2 active" href="/examples/vision/masked_image_modeling/">Masked image modeling with Autoencoders</a> <a class="nav-sublink2" href="/examples/vision/sam/">Segment Anything Model with 🤗Transformers</a> <a class="nav-sublink2" href="/examples/vision/segformer/">Semantic segmentation with SegFormer and Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/vision/simsiam/">Self-supervised contrastive learning with SimSiam</a> <a class="nav-sublink2" href="/examples/vision/supervised-contrastive-learning/">Supervised Contrastive Learning</a> <a class="nav-sublink2" href="/examples/vision/yolov8/">Efficient Object Detection with YOLOV8 and KerasCV</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/vision/'>Computer Vision</a> / Masked image modeling with Autoencoders </div> <div class='k-content'> <h1 id="masked-image-modeling-with-autoencoders">Masked image modeling with Autoencoders</h1> Author: <a href="https://twitter.com/arig23498">Aritra Roy Gosthipaty</a>, <a href="https://twitter.com/RisingSayak">Sayak Paul</a> Date created: 2021/12/20 Last modified: 2021/12/21 Description: Implementing Masked Autoencoders for self-supervised pretraining. <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/vision/ipynb/masked_image_modeling.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/vision/masked_image_modeling.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> In deep learning, models with growing capacity and capability can easily overfit on large datasets (ImageNet-1K). In the field of natural language processing, the appetite for data has been successfully addressed by self-supervised pretraining. In the academic paper <a href="https://arxiv.org/abs/2111.06377">Masked Autoencoders Are Scalable Vision Learners</a> by He et. al. the authors propose a simple yet effective method to pretrain large vision models (here <a href="https://arxiv.org/abs/2010.11929">ViT Huge</a>). Inspired from the pretraining algorithm of BERT (<a href="https://arxiv.org/abs/1810.04805">Devlin et al.</a>), they mask patches of an image and, through an autoencoder predict the masked patches. In the spirit of "masked language modeling", this pretraining task could be referred to as "masked image modeling". In this example, we implement <a href="https://arxiv.org/abs/2111.06377">Masked Autoencoders Are Scalable Vision Learners</a> with the <a href="https://www.cs.toronto.edu/~kriz/cifar.html">CIFAR-10</a> dataset. After pretraining a scaled down version of ViT, we also implement the linear evaluation pipeline on CIFAR-10. This implementation covers (MAE refers to Masked Autoencoder): <ul> <li>The masking algorithm</li> <li>MAE encoder</li> <li>MAE decoder</li> <li>Evaluation with linear probing</li> </ul> As a reference, we reuse some of the code presented in <a href="https://keras.io/examples/vision/image_classification_with_vision_transformer/">this example</a>. <hr /> <h2 id="imports">Imports</h2> <div class="codehilite"><pre><code>import os os.environ["KERAS_BACKEND"] = "tensorflow" import tensorflow as tf import keras from keras import layers import matplotlib.pyplot as plt import numpy as np import random # Setting seeds for reproducibility. SEED = 42 keras.utils.set_random_seed(SEED) </code></pre></div> <hr /> <h2 id="hyperparameters-for-pretraining">Hyperparameters for pretraining</h2> Please feel free to change the hyperparameters and check your results. The best way to get an intuition about the architecture is to experiment with it. Our hyperparameters are heavily inspired by the design guidelines laid out by the authors in <a href="https://arxiv.org/abs/2111.06377">the original paper</a>. <div class="codehilite"><pre><code># DATA BUFFER_SIZE = 1024 BATCH_SIZE = 256 AUTO = tf.data.AUTOTUNE INPUT_SHAPE = (32, 32, 3) NUM_CLASSES = 10 # OPTIMIZER LEARNING_RATE = 5e-3 WEIGHT_DECAY = 1e-4 # PRETRAINING EPOCHS = 100 # AUGMENTATION IMAGE_SIZE = 48 # We will resize input images to this size. PATCH_SIZE = 6 # Size of the patches to be extracted from the input images. NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2 MASK_PROPORTION = 0.75 # We have found 75% masking to give us the best results. # ENCODER and DECODER LAYER_NORM_EPS = 1e-6 ENC_PROJECTION_DIM = 128 DEC_PROJECTION_DIM = 64 ENC_NUM_HEADS = 4 ENC_LAYERS = 6 DEC_NUM_HEADS = 4 DEC_LAYERS = ( 2 # The decoder is lightweight but should be reasonably deep for reconstruction. ) ENC_TRANSFORMER_UNITS = [ ENC_PROJECTION_DIM * 2, ENC_PROJECTION_DIM, ] # Size of the transformer layers. DEC_TRANSFORMER_UNITS = [ DEC_PROJECTION_DIM * 2, DEC_PROJECTION_DIM, ] </code></pre></div> <hr /> <h2 id="load-and-prepare-the-cifar10-dataset">Load and prepare the CIFAR-10 dataset</h2> <div class="codehilite"><pre><code>(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data() (x_train, y_train), (x_val, y_val) = ( (x_train[:40000], y_train[:40000]), (x_train[40000:], y_train[40000:]), ) print(f"Training samples: {len(x_train)}") print(f"Validation samples: {len(x_val)}") print(f"Testing samples: {len(x_test)}") train_ds = tf.data.Dataset.from_tensor_slices(x_train) train_ds = train_ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE).prefetch(AUTO) val_ds = tf.data.Dataset.from_tensor_slices(x_val) val_ds = val_ds.batch(BATCH_SIZE).prefetch(AUTO) test_ds = tf.data.Dataset.from_tensor_slices(x_test) test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Training samples: 40000 Validation samples: 10000 Testing samples: 10000 </code></pre></div> </div> <hr /> <h2 id="data-augmentation">Data augmentation</h2> In previous self-supervised pretraining methodologies (<a href="https://arxiv.org/abs/2002.05709">SimCLR</a> alike), we have noticed that the data augmentation pipeline plays an important role. On the other hand the authors of this paper point out that Masked Autoencoders do not rely on augmentations. They propose a simple augmentation pipeline of: <ul> <li>Resizing</li> <li>Random cropping (fixed-sized or random sized)</li> <li>Random horizontal flipping</li> </ul> <div class="codehilite"><pre><code>def get_train_augmentation_model(): model = keras.Sequential( [ layers.Rescaling(1 / 255.0), layers.Resizing(INPUT_SHAPE[0] + 20, INPUT_SHAPE[0] + 20), layers.RandomCrop(IMAGE_SIZE, IMAGE_SIZE), layers.RandomFlip("horizontal"), ], name="train_data_augmentation", ) return model def get_test_augmentation_model(): model = keras.Sequential( [ layers.Rescaling(1 / 255.0), layers.Resizing(IMAGE_SIZE, IMAGE_SIZE), ], name="test_data_augmentation", ) return model </code></pre></div> <hr /> <h2 id="a-layer-for-extracting-patches-from-images">A layer for extracting patches from images</h2> This layer takes images as input and divides them into patches. The layer also includes two utility method: <ul> <li><code>show_patched_image</code> – Takes a batch of images and its corresponding patches to plot a random pair of image and patches.</li> <li><code>reconstruct_from_patch</code> – Takes a single instance of patches and stitches them together into the original image.</li> </ul> <div class="codehilite"><pre><code>class Patches(layers.Layer): def __init__(self, patch_size=PATCH_SIZE, **kwargs): super().__init__(**kwargs) self.patch_size = patch_size # Assuming the image has three channels each patch would be # of size (patch_size, patch_size, 3). self.resize = layers.Reshape((-1, patch_size * patch_size * 3)) def call(self, images): # Create patches from the input images patches = tf.image.extract_patches( images=images, sizes=[1, self.patch_size, self.patch_size, 1], strides=[1, self.patch_size, self.patch_size, 1], rates=[1, 1, 1, 1], padding="VALID", ) # Reshape the patches to (batch, num_patches, patch_area) and return it. patches = self.resize(patches) return patches def show_patched_image(self, images, patches): # This is a utility function which accepts a batch of images and its # corresponding patches and help visualize one image and its patches # side by side. idx = np.random.choice(patches.shape[0]) print(f"Index selected: {idx}.") plt.figure(figsize=(4, 4)) plt.imshow(keras.utils.array_to_img(images[idx])) plt.axis("off") plt.show() n = int(np.sqrt(patches.shape[1])) plt.figure(figsize=(4, 4)) for i, patch in enumerate(patches[idx]): ax = plt.subplot(n, n, i + 1) patch_img = tf.reshape(patch, (self.patch_size, self.patch_size, 3)) plt.imshow(keras.utils.img_to_array(patch_img)) plt.axis("off") plt.show() # Return the index chosen to validate it outside the method. return idx # taken from https://stackoverflow.com/a/58082878/10319735 def reconstruct_from_patch(self, patch): # This utility function takes patches from a *single* image and # reconstructs it back into the image. This is useful for the train # monitor callback. num_patches = patch.shape[0] n = int(np.sqrt(num_patches)) patch = tf.reshape(patch, (num_patches, self.patch_size, self.patch_size, 3)) rows = tf.split(patch, n, axis=0) rows = [tf.concat(tf.unstack(x), axis=1) for x in rows] reconstructed = tf.concat(rows, axis=0) return reconstructed </code></pre></div> Let's visualize the image patches. <div class="codehilite"><pre><code># Get a batch of images. image_batch = next(iter(train_ds)) # Augment the images. augmentation_model = get_train_augmentation_model() augmented_images = augmentation_model(image_batch) # Define the patch layer. patch_layer = Patches() # Get the patches from the batched images. patches = patch_layer(images=augmented_images) # Now pass the images and the corresponding patches # to the `show_patched_image` method. random_index = patch_layer.show_patched_image(images=augmented_images, patches=patches) # Chose the same chose image and try reconstructing the patches # into the original image. image = patch_layer.reconstruct_from_patch(patches[random_index]) plt.imshow(image) plt.axis("off") plt.show() </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Index selected: 102. </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_13_1.png" /> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_13_2.png" /> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_13_3.png" /> <hr /> <h2 id="patch-encoding-with-masking">Patch encoding with masking</h2> Quoting the paper <blockquote> Following ViT, we divide an image into regular non-overlapping patches. Then we sample a subset of patches and mask (i.e., remove) the remaining ones. Our sampling strategy is straightforward: we sample random patches without replacement, following a uniform distribution. We simply refer to this as “random sampling”. </blockquote> This layer includes masking and encoding the patches. The utility methods of the layer are: <ul> <li><code>get_random_indices</code> – Provides the mask and unmask indices.</li> <li><code>generate_masked_image</code> – Takes patches and unmask indices, results in a random masked image. This is an essential utility method for our training monitor callback (defined later).</li> </ul> <div class="codehilite"><pre><code>class PatchEncoder(layers.Layer): def __init__( self, patch_size=PATCH_SIZE, projection_dim=ENC_PROJECTION_DIM, mask_proportion=MASK_PROPORTION, downstream=False, **kwargs, ): super().__init__(**kwargs) self.patch_size = patch_size self.projection_dim = projection_dim self.mask_proportion = mask_proportion self.downstream = downstream # This is a trainable mask token initialized randomly from a normal # distribution. self.mask_token = tf.Variable( tf.random.normal([1, patch_size * patch_size * 3]), trainable=True ) def build(self, input_shape): (_, self.num_patches, self.patch_area) = input_shape # Create the projection layer for the patches. self.projection = layers.Dense(units=self.projection_dim) # Create the positional embedding layer. self.position_embedding = layers.Embedding( input_dim=self.num_patches, output_dim=self.projection_dim ) # Number of patches that will be masked. self.num_mask = int(self.mask_proportion * self.num_patches) def call(self, patches): # Get the positional embeddings. batch_size = tf.shape(patches)[0] positions = tf.range(start=0, limit=self.num_patches, delta=1) pos_embeddings = self.position_embedding(positions[tf.newaxis, ...]) pos_embeddings = tf.tile( pos_embeddings, [batch_size, 1, 1] ) # (B, num_patches, projection_dim) # Embed the patches. patch_embeddings = ( self.projection(patches) + pos_embeddings ) # (B, num_patches, projection_dim) if self.downstream: return patch_embeddings else: mask_indices, unmask_indices = self.get_random_indices(batch_size) # The encoder input is the unmasked patch embeddings. Here we gather # all the patches that should be unmasked. unmasked_embeddings = tf.gather( patch_embeddings, unmask_indices, axis=1, batch_dims=1 ) # (B, unmask_numbers, projection_dim) # Get the unmasked and masked position embeddings. We will need them # for the decoder. unmasked_positions = tf.gather( pos_embeddings, unmask_indices, axis=1, batch_dims=1 ) # (B, unmask_numbers, projection_dim) masked_positions = tf.gather( pos_embeddings, mask_indices, axis=1, batch_dims=1 ) # (B, mask_numbers, projection_dim) # Repeat the mask token number of mask times. # Mask tokens replace the masks of the image. mask_tokens = tf.repeat(self.mask_token, repeats=self.num_mask, axis=0) mask_tokens = tf.repeat( mask_tokens[tf.newaxis, ...], repeats=batch_size, axis=0 ) # Get the masked embeddings for the tokens. masked_embeddings = self.projection(mask_tokens) + masked_positions return ( unmasked_embeddings, # Input to the encoder. masked_embeddings, # First part of input to the decoder. unmasked_positions, # Added to the encoder outputs. mask_indices, # The indices that were masked. unmask_indices, # The indices that were unmaksed. ) def get_random_indices(self, batch_size): # Create random indices from a uniform distribution and then split # it into mask and unmask indices. rand_indices = tf.argsort( tf.random.uniform(shape=(batch_size, self.num_patches)), axis=-1 ) mask_indices = rand_indices[:, : self.num_mask] unmask_indices = rand_indices[:, self.num_mask :] return mask_indices, unmask_indices def generate_masked_image(self, patches, unmask_indices): # Choose a random patch and it corresponding unmask index. idx = np.random.choice(patches.shape[0]) patch = patches[idx] unmask_index = unmask_indices[idx] # Build a numpy array of same shape as patch. new_patch = np.zeros_like(patch) # Iterate of the new_patch and plug the unmasked patches. count = 0 for i in range(unmask_index.shape[0]): new_patch[unmask_index[i]] = patch[unmask_index[i]] return new_patch, idx </code></pre></div> Let's see the masking process in action on a sample image. <div class="codehilite"><pre><code># Create the patch encoder layer. patch_encoder = PatchEncoder() # Get the embeddings and positions. ( unmasked_embeddings, masked_embeddings, unmasked_positions, mask_indices, unmask_indices, ) = patch_encoder(patches=patches) # Show a maksed patch image. new_patch, random_index = patch_encoder.generate_masked_image(patches, unmask_indices) plt.figure(figsize=(10, 10)) plt.subplot(1, 2, 1) img = patch_layer.reconstruct_from_patch(new_patch) plt.imshow(keras.utils.array_to_img(img)) plt.axis("off") plt.title("Masked") plt.subplot(1, 2, 2) img = augmented_images[random_index] plt.imshow(keras.utils.array_to_img(img)) plt.axis("off") plt.title("Original") plt.show() </code></pre></div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_17_0.png" /> <hr /> <h2 id="mlp">MLP</h2> This serves as the fully connected feed forward network of the transformer architecture. <div class="codehilite"><pre><code>def mlp(x, dropout_rate, hidden_units): for units in hidden_units: x = layers.Dense(units, activation=tf.nn.gelu)(x) x = layers.Dropout(dropout_rate)(x) return x </code></pre></div> <hr /> <h2 id="mae-encoder">MAE encoder</h2> The MAE encoder is ViT. The only point to note here is that the encoder outputs a layer normalized output. <div class="codehilite"><pre><code>def create_encoder(num_heads=ENC_NUM_HEADS, num_layers=ENC_LAYERS): inputs = layers.Input((None, ENC_PROJECTION_DIM)) x = inputs for _ in range(num_layers): # Layer normalization 1. x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x) # Create a multi-head attention layer. attention_output = layers.MultiHeadAttention( num_heads=num_heads, key_dim=ENC_PROJECTION_DIM, dropout=0.1 )(x1, x1) # Skip connection 1. x2 = layers.Add()([attention_output, x]) # Layer normalization 2. x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2) # MLP. x3 = mlp(x3, hidden_units=ENC_TRANSFORMER_UNITS, dropout_rate=0.1) # Skip connection 2. x = layers.Add()([x3, x2]) outputs = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x) return keras.Model(inputs, outputs, name="mae_encoder") </code></pre></div> <hr /> <h2 id="mae-decoder">MAE decoder</h2> The authors point out that they use an asymmetric autoencoder model. They use a lightweight decoder that takes "<10% computation per token vs. the encoder". We are not specific with the "<10% computation" in our implementation but have used a smaller decoder (both in terms of depth and projection dimensions). <div class="codehilite"><pre><code>def create_decoder( num_layers=DEC_LAYERS, num_heads=DEC_NUM_HEADS, image_size=IMAGE_SIZE ): inputs = layers.Input((NUM_PATCHES, ENC_PROJECTION_DIM)) x = layers.Dense(DEC_PROJECTION_DIM)(inputs) for _ in range(num_layers): # Layer normalization 1. x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x) # Create a multi-head attention layer. attention_output = layers.MultiHeadAttention( num_heads=num_heads, key_dim=DEC_PROJECTION_DIM, dropout=0.1 )(x1, x1) # Skip connection 1. x2 = layers.Add()([attention_output, x]) # Layer normalization 2. x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2) # MLP. x3 = mlp(x3, hidden_units=DEC_TRANSFORMER_UNITS, dropout_rate=0.1) # Skip connection 2. x = layers.Add()([x3, x2]) x = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x) x = layers.Flatten()(x) pre_final = layers.Dense(units=image_size * image_size * 3, activation="sigmoid")(x) outputs = layers.Reshape((image_size, image_size, 3))(pre_final) return keras.Model(inputs, outputs, name="mae_decoder") </code></pre></div> <hr /> <h2 id="mae-trainer">MAE trainer</h2> This is the trainer module. We wrap the encoder and decoder inside of a <a href="https://www.tensorflow.org/api_docs/python/tf/keras/Model"><code>tf.keras.Model</code></a> subclass. This allows us to customize what happens in the <code>model.fit()</code> loop. <div class="codehilite"><pre><code>class MaskedAutoencoder(keras.Model): def __init__( self, train_augmentation_model, test_augmentation_model, patch_layer, patch_encoder, encoder, decoder, **kwargs, ): super().__init__(**kwargs) self.train_augmentation_model = train_augmentation_model self.test_augmentation_model = test_augmentation_model self.patch_layer = patch_layer self.patch_encoder = patch_encoder self.encoder = encoder self.decoder = decoder def calculate_loss(self, images, test=False): # Augment the input images. if test: augmented_images = self.test_augmentation_model(images) else: augmented_images = self.train_augmentation_model(images) # Patch the augmented images. patches = self.patch_layer(augmented_images) # Encode the patches. ( unmasked_embeddings, masked_embeddings, unmasked_positions, mask_indices, unmask_indices, ) = self.patch_encoder(patches) # Pass the unmaksed patche to the encoder. encoder_outputs = self.encoder(unmasked_embeddings) # Create the decoder inputs. encoder_outputs = encoder_outputs + unmasked_positions decoder_inputs = tf.concat([encoder_outputs, masked_embeddings], axis=1) # Decode the inputs. decoder_outputs = self.decoder(decoder_inputs) decoder_patches = self.patch_layer(decoder_outputs) loss_patch = tf.gather(patches, mask_indices, axis=1, batch_dims=1) loss_output = tf.gather(decoder_patches, mask_indices, axis=1, batch_dims=1) # Compute the total loss. total_loss = self.compute_loss(y=loss_patch, y_pred=loss_output) return total_loss, loss_patch, loss_output def train_step(self, images): with tf.GradientTape() as tape: total_loss, loss_patch, loss_output = self.calculate_loss(images) # Apply gradients. train_vars = [ self.train_augmentation_model.trainable_variables, self.patch_layer.trainable_variables, self.patch_encoder.trainable_variables, self.encoder.trainable_variables, self.decoder.trainable_variables, ] grads = tape.gradient(total_loss, train_vars) tv_list = [] for grad, var in zip(grads, train_vars): for g, v in zip(grad, var): tv_list.append((g, v)) self.optimizer.apply_gradients(tv_list) # Report progress. results = {} for metric in self.metrics: metric.update_state(loss_patch, loss_output) results[metric.name] = metric.result() return results def test_step(self, images): total_loss, loss_patch, loss_output = self.calculate_loss(images, test=True) # Update the trackers. results = {} for metric in self.metrics: metric.update_state(loss_patch, loss_output) results[metric.name] = metric.result() return results </code></pre></div> <hr /> <h2 id="model-initialization">Model initialization</h2> <div class="codehilite"><pre><code>train_augmentation_model = get_train_augmentation_model() test_augmentation_model = get_test_augmentation_model() patch_layer = Patches() patch_encoder = PatchEncoder() encoder = create_encoder() decoder = create_decoder() mae_model = MaskedAutoencoder( train_augmentation_model=train_augmentation_model, test_augmentation_model=test_augmentation_model, patch_layer=patch_layer, patch_encoder=patch_encoder, encoder=encoder, decoder=decoder, ) </code></pre></div> <hr /> <h2 id="training-callbacks">Training callbacks</h2> <h3 id="visualization-callback">Visualization callback</h3> <div class="codehilite"><pre><code># Taking a batch of test inputs to measure model's progress. test_images = next(iter(test_ds)) class TrainMonitor(keras.callbacks.Callback): def __init__(self, epoch_interval=None): self.epoch_interval = epoch_interval def on_epoch_end(self, epoch, logs=None): if self.epoch_interval and epoch % self.epoch_interval == 0: test_augmented_images = self.model.test_augmentation_model(test_images) test_patches = self.model.patch_layer(test_augmented_images) ( test_unmasked_embeddings, test_masked_embeddings, test_unmasked_positions, test_mask_indices, test_unmask_indices, ) = self.model.patch_encoder(test_patches) test_encoder_outputs = self.model.encoder(test_unmasked_embeddings) test_encoder_outputs = test_encoder_outputs + test_unmasked_positions test_decoder_inputs = tf.concat( [test_encoder_outputs, test_masked_embeddings], axis=1 ) test_decoder_outputs = self.model.decoder(test_decoder_inputs) # Show a maksed patch image. test_masked_patch, idx = self.model.patch_encoder.generate_masked_image( test_patches, test_unmask_indices ) print(f"\nIdx chosen: {idx}") original_image = test_augmented_images[idx] masked_image = self.model.patch_layer.reconstruct_from_patch( test_masked_patch ) reconstructed_image = test_decoder_outputs[idx] fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5)) ax[0].imshow(original_image) ax[0].set_title(f"Original: {epoch:03d}") ax[1].imshow(masked_image) ax[1].set_title(f"Masked: {epoch:03d}") ax[2].imshow(reconstructed_image) ax[2].set_title(f"Resonstructed: {epoch:03d}") plt.show() plt.close() </code></pre></div> <h3 id="learning-rate-scheduler">Learning rate scheduler</h3> <div class="codehilite"><pre><code># Some code is taken from: # https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2. class WarmUpCosine(keras.optimizers.schedules.LearningRateSchedule): def __init__( self, learning_rate_base, total_steps, warmup_learning_rate, warmup_steps ): super().__init__() self.learning_rate_base = learning_rate_base self.total_steps = total_steps self.warmup_learning_rate = warmup_learning_rate self.warmup_steps = warmup_steps self.pi = tf.constant(np.pi) def __call__(self, step): if self.total_steps < self.warmup_steps: raise ValueError("Total_steps must be larger or equal to warmup_steps.") cos_annealed_lr = tf.cos( self.pi * (tf.cast(step, tf.float32) - self.warmup_steps) / float(self.total_steps - self.warmup_steps) ) learning_rate = 0.5 * self.learning_rate_base * (1 + cos_annealed_lr) if self.warmup_steps > 0: if self.learning_rate_base < self.warmup_learning_rate: raise ValueError( "Learning_rate_base must be larger or equal to " "warmup_learning_rate." ) slope = ( self.learning_rate_base - self.warmup_learning_rate ) / self.warmup_steps warmup_rate = slope * tf.cast(step, tf.float32) + self.warmup_learning_rate learning_rate = tf.where( step < self.warmup_steps, warmup_rate, learning_rate ) return tf.where( step > self.total_steps, 0.0, learning_rate, name="learning_rate" ) total_steps = int((len(x_train) / BATCH_SIZE) * EPOCHS) warmup_epoch_percentage = 0.15 warmup_steps = int(total_steps * warmup_epoch_percentage) scheduled_lrs = WarmUpCosine( learning_rate_base=LEARNING_RATE, total_steps=total_steps, warmup_learning_rate=0.0, warmup_steps=warmup_steps, ) lrs = [scheduled_lrs(step) for step in range(total_steps)] plt.plot(lrs) plt.xlabel("Step", fontsize=14) plt.ylabel("LR", fontsize=14) plt.show() # Assemble the callbacks. train_callbacks = [TrainMonitor(epoch_interval=5)] </code></pre></div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_32_0.png" /> <hr /> <h2 id="model-compilation-and-training">Model compilation and training</h2> <div class="codehilite"><pre><code>optimizer = keras.optimizers.AdamW( learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY ) # Compile and pretrain the model. mae_model.compile( optimizer=optimizer, loss=keras.losses.MeanSquaredError(), metrics=["mae"] ) history = mae_model.fit( train_ds, epochs=EPOCHS, validation_data=val_ds, callbacks=train_callbacks, ) # Measure its performance. loss, mae = mae_model.evaluate(test_ds) print(f"Loss: {loss:.2f}") print(f"MAE: {mae:.2f}") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 80ms/step - mae: 0.2035 - loss: 0.4828 Idx chosen: 92 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_1.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 47s 95ms/step - mae: 0.2033 - loss: 0.4828 - val_loss: 0.5225 - val_mae: 0.1600 Epoch 2/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1592 - loss: 0.5128 - val_loss: 0.5290 - val_mae: 0.1511 Epoch 3/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1530 - loss: 0.5193 - val_loss: 0.5336 - val_mae: 0.1478 Epoch 4/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1502 - loss: 0.5220 - val_loss: 0.5298 - val_mae: 0.1436 Epoch 5/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1458 - loss: 0.5245 - val_loss: 0.5296 - val_mae: 0.1405 Epoch 6/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 81ms/step - mae: 0.1414 - loss: 0.5265 Idx chosen: 14 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_3.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 14s 88ms/step - mae: 0.1414 - loss: 0.5265 - val_loss: 0.5328 - val_mae: 0.1402 Epoch 7/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1399 - loss: 0.5278 - val_loss: 0.5361 - val_mae: 0.1360 Epoch 8/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1389 - loss: 0.5285 - val_loss: 0.5365 - val_mae: 0.1424 Epoch 9/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1379 - loss: 0.5295 - val_loss: 0.5312 - val_mae: 0.1345 Epoch 10/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1352 - loss: 0.5308 - val_loss: 0.5374 - val_mae: 0.1321 Epoch 11/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 81ms/step - mae: 0.1339 - loss: 0.5317 Idx chosen: 106 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_5.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 14s 87ms/step - mae: 0.1339 - loss: 0.5317 - val_loss: 0.5392 - val_mae: 0.1330 Epoch 12/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1321 - loss: 0.5331 - val_loss: 0.5383 - val_mae: 0.1301 Epoch 13/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1317 - loss: 0.5343 - val_loss: 0.5405 - val_mae: 0.1322 Epoch 14/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1326 - loss: 0.5338 - val_loss: 0.5404 - val_mae: 0.1280 Epoch 15/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 84ms/step - mae: 0.1297 - loss: 0.5343 - val_loss: 0.5444 - val_mae: 0.1261 Epoch 16/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 82ms/step - mae: 0.1276 - loss: 0.5361 Idx chosen: 71 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_7.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 14s 91ms/step - mae: 0.1276 - loss: 0.5362 - val_loss: 0.5456 - val_mae: 0.1243 Epoch 17/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1262 - loss: 0.5382 - val_loss: 0.5427 - val_mae: 0.1233 Epoch 18/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1221 - loss: 0.5407 - val_loss: 0.5473 - val_mae: 0.1196 Epoch 19/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1209 - loss: 0.5412 - val_loss: 0.5511 - val_mae: 0.1176 Epoch 20/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1202 - loss: 0.5422 - val_loss: 0.5515 - val_mae: 0.1167 Epoch 21/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1186 - loss: 0.5430 Idx chosen: 188 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_9.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 85ms/step - mae: 0.1186 - loss: 0.5430 - val_loss: 0.5546 - val_mae: 0.1168 Epoch 22/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1171 - loss: 0.5446 - val_loss: 0.5500 - val_mae: 0.1155 Epoch 23/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1161 - loss: 0.5457 - val_loss: 0.5559 - val_mae: 0.1135 Epoch 24/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1135 - loss: 0.5479 - val_loss: 0.5521 - val_mae: 0.1112 Epoch 25/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1128 - loss: 0.5480 - val_loss: 0.5505 - val_mae: 0.1122 Epoch 26/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1123 - loss: 0.5470 Idx chosen: 20 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_11.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.1123 - loss: 0.5470 - val_loss: 0.5572 - val_mae: 0.1127 Epoch 27/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1114 - loss: 0.5487 - val_loss: 0.5555 - val_mae: 0.1092 Epoch 28/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1108 - loss: 0.5492 - val_loss: 0.5569 - val_mae: 0.1110 Epoch 29/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1104 - loss: 0.5491 - val_loss: 0.5517 - val_mae: 0.1110 Epoch 30/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1099 - loss: 0.5490 - val_loss: 0.5543 - val_mae: 0.1104 Epoch 31/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1095 - loss: 0.5501 Idx chosen: 102 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_13.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.1095 - loss: 0.5501 - val_loss: 0.5578 - val_mae: 0.1108 Epoch 32/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1089 - loss: 0.5503 - val_loss: 0.5620 - val_mae: 0.1081 Epoch 33/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1079 - loss: 0.5509 - val_loss: 0.5618 - val_mae: 0.1067 Epoch 34/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1067 - loss: 0.5524 - val_loss: 0.5627 - val_mae: 0.1059 Epoch 35/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1068 - loss: 0.5515 - val_loss: 0.5576 - val_mae: 0.1050 Epoch 36/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1057 - loss: 0.5526 Idx chosen: 121 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_15.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.1057 - loss: 0.5526 - val_loss: 0.5627 - val_mae: 0.1050 Epoch 37/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1065 - loss: 0.5534 - val_loss: 0.5638 - val_mae: 0.1050 Epoch 38/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1055 - loss: 0.5528 - val_loss: 0.5527 - val_mae: 0.1083 Epoch 39/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 20s 82ms/step - mae: 0.1056 - loss: 0.5516 - val_loss: 0.5562 - val_mae: 0.1044 Epoch 40/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1053 - loss: 0.5528 - val_loss: 0.5567 - val_mae: 0.1051 Epoch 41/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 78ms/step - mae: 0.1049 - loss: 0.5533 Idx chosen: 210 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_17.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 85ms/step - mae: 0.1049 - loss: 0.5533 - val_loss: 0.5620 - val_mae: 0.1030 Epoch 42/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1041 - loss: 0.5534 - val_loss: 0.5650 - val_mae: 0.1052 Epoch 43/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1048 - loss: 0.5526 - val_loss: 0.5619 - val_mae: 0.1027 Epoch 44/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1037 - loss: 0.5543 - val_loss: 0.5615 - val_mae: 0.1031 Epoch 45/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1036 - loss: 0.5535 - val_loss: 0.5575 - val_mae: 0.1026 Epoch 46/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 78ms/step - mae: 0.1032 - loss: 0.5537 Idx chosen: 214 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_19.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 85ms/step - mae: 0.1032 - loss: 0.5537 - val_loss: 0.5549 - val_mae: 0.1037 Epoch 47/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 84ms/step - mae: 0.1035 - loss: 0.5539 - val_loss: 0.5597 - val_mae: 0.1031 Epoch 48/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1033 - loss: 0.5533 - val_loss: 0.5650 - val_mae: 0.1013 Epoch 49/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.1027 - loss: 0.5543 - val_loss: 0.5571 - val_mae: 0.1028 Epoch 50/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1024 - loss: 0.5548 - val_loss: 0.5592 - val_mae: 0.1018 Epoch 51/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 78ms/step - mae: 0.1025 - loss: 0.5543 Idx chosen: 74 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_21.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 85ms/step - mae: 0.1025 - loss: 0.5543 - val_loss: 0.5645 - val_mae: 0.1007 Epoch 52/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.1025 - loss: 0.5544 - val_loss: 0.5616 - val_mae: 0.1004 Epoch 53/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1014 - loss: 0.5547 - val_loss: 0.5594 - val_mae: 0.1007 Epoch 54/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1014 - loss: 0.5550 - val_loss: 0.5687 - val_mae: 0.1012 Epoch 55/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1022 - loss: 0.5551 - val_loss: 0.5572 - val_mae: 0.1018 Epoch 56/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1015 - loss: 0.5558 Idx chosen: 202 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_23.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.1015 - loss: 0.5558 - val_loss: 0.5619 - val_mae: 0.0996 Epoch 57/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1008 - loss: 0.5550 - val_loss: 0.5614 - val_mae: 0.0996 Epoch 58/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1004 - loss: 0.5557 - val_loss: 0.5620 - val_mae: 0.0995 Epoch 59/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.1002 - loss: 0.5558 - val_loss: 0.5612 - val_mae: 0.0997 Epoch 60/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.1005 - loss: 0.5563 - val_loss: 0.5598 - val_mae: 0.1000 Epoch 61/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.1001 - loss: 0.5564 Idx chosen: 87 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_25.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.1001 - loss: 0.5564 - val_loss: 0.5606 - val_mae: 0.0998 Epoch 62/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.0998 - loss: 0.5562 - val_loss: 0.5643 - val_mae: 0.0988 Epoch 63/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.1001 - loss: 0.5556 - val_loss: 0.5657 - val_mae: 0.0985 Epoch 64/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0998 - loss: 0.5566 - val_loss: 0.5624 - val_mae: 0.0989 Epoch 65/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0994 - loss: 0.5564 - val_loss: 0.5576 - val_mae: 0.0999 Epoch 66/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - mae: 0.0993 - loss: 0.5567 Idx chosen: 116 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_27.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 86ms/step - mae: 0.0993 - loss: 0.5567 - val_loss: 0.5572 - val_mae: 0.1000 Epoch 67/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0990 - loss: 0.5570 - val_loss: 0.5619 - val_mae: 0.0981 Epoch 68/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0987 - loss: 0.5578 - val_loss: 0.5644 - val_mae: 0.0973 Epoch 69/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0981 - loss: 0.5577 - val_loss: 0.5639 - val_mae: 0.0976 Epoch 70/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0986 - loss: 0.5563 - val_loss: 0.5601 - val_mae: 0.0989 Epoch 71/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0982 - loss: 0.5578 Idx chosen: 99 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_29.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 84ms/step - mae: 0.0982 - loss: 0.5577 - val_loss: 0.5628 - val_mae: 0.0970 Epoch 72/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0979 - loss: 0.5569 - val_loss: 0.5637 - val_mae: 0.0968 Epoch 73/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0979 - loss: 0.5575 - val_loss: 0.5606 - val_mae: 0.0975 Epoch 74/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0977 - loss: 0.5572 - val_loss: 0.5628 - val_mae: 0.0967 Epoch 75/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0975 - loss: 0.5572 - val_loss: 0.5631 - val_mae: 0.0964 Epoch 76/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0973 - loss: 0.5580 Idx chosen: 103 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_31.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.0973 - loss: 0.5579 - val_loss: 0.5628 - val_mae: 0.0967 Epoch 77/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0974 - loss: 0.5579 - val_loss: 0.5638 - val_mae: 0.0963 Epoch 78/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0968 - loss: 0.5585 - val_loss: 0.5615 - val_mae: 0.0967 Epoch 79/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0969 - loss: 0.5578 - val_loss: 0.5641 - val_mae: 0.0959 Epoch 80/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0967 - loss: 0.5584 - val_loss: 0.5619 - val_mae: 0.0962 Epoch 81/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0965 - loss: 0.5578 Idx chosen: 151 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_33.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.0965 - loss: 0.5578 - val_loss: 0.5651 - val_mae: 0.0957 Epoch 82/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0965 - loss: 0.5583 - val_loss: 0.5644 - val_mae: 0.0957 Epoch 83/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0962 - loss: 0.5584 - val_loss: 0.5649 - val_mae: 0.0954 Epoch 84/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0962 - loss: 0.5586 - val_loss: 0.5611 - val_mae: 0.0962 Epoch 85/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0961 - loss: 0.5582 - val_loss: 0.5638 - val_mae: 0.0956 Epoch 86/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0961 - loss: 0.5584 Idx chosen: 130 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_35.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 83ms/step - mae: 0.0961 - loss: 0.5584 - val_loss: 0.5641 - val_mae: 0.0954 Epoch 87/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0959 - loss: 0.5580 - val_loss: 0.5641 - val_mae: 0.0953 Epoch 88/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0960 - loss: 0.5583 - val_loss: 0.5642 - val_mae: 0.0953 Epoch 89/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0958 - loss: 0.5591 - val_loss: 0.5635 - val_mae: 0.0953 Epoch 90/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0957 - loss: 0.5587 - val_loss: 0.5648 - val_mae: 0.0948 Epoch 91/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0957 - loss: 0.5585 Idx chosen: 149 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_37.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 84ms/step - mae: 0.0957 - loss: 0.5585 - val_loss: 0.5636 - val_mae: 0.0952 Epoch 92/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0957 - loss: 0.5593 - val_loss: 0.5642 - val_mae: 0.0950 Epoch 93/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0957 - loss: 0.5598 - val_loss: 0.5635 - val_mae: 0.0950 Epoch 94/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0956 - loss: 0.5587 - val_loss: 0.5641 - val_mae: 0.0950 Epoch 95/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0955 - loss: 0.5587 - val_loss: 0.5637 - val_mae: 0.0950 Epoch 96/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 0s 77ms/step - mae: 0.0956 - loss: 0.5585 Idx chosen: 52 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/masked_image_modeling/masked_image_modeling_34_39.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 157/157 ━━━━━━━━━━━━━━━━━━━━ 14s 87ms/step - mae: 0.0956 - loss: 0.5585 - val_loss: 0.5643 - val_mae: 0.0950 Epoch 97/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 81ms/step - mae: 0.0956 - loss: 0.5587 - val_loss: 0.5642 - val_mae: 0.0950 Epoch 98/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 82ms/step - mae: 0.0954 - loss: 0.5586 - val_loss: 0.5639 - val_mae: 0.0950 Epoch 99/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0954 - loss: 0.5580 - val_loss: 0.5641 - val_mae: 0.0950 Epoch 100/100 157/157 ━━━━━━━━━━━━━━━━━━━━ 13s 80ms/step - mae: 0.0955 - loss: 0.5587 - val_loss: 0.5639 - val_mae: 0.0951 40/40 ━━━━━━━━━━━━━━━━━━━━ 1s 13ms/step - mae: 0.0955 - loss: 0.5684 Loss: 0.57 MAE: 0.10 </code></pre></div> </div> <hr /> <h2 id="evaluation-with-linear-probing">Evaluation with linear probing</h2> <h3 id="extract-the-encoder-model-along-with-other-layers">Extract the encoder model along with other layers</h3> <div class="codehilite"><pre><code># Extract the augmentation layers. train_augmentation_model = mae_model.train_augmentation_model test_augmentation_model = mae_model.test_augmentation_model # Extract the patchers. patch_layer = mae_model.patch_layer patch_encoder = mae_model.patch_encoder patch_encoder.downstream = True # Swtich the downstream flag to True. # Extract the encoder. encoder = mae_model.encoder # Pack as a model. downstream_model = keras.Sequential( [ layers.Input((IMAGE_SIZE, IMAGE_SIZE, 3)), patch_layer, patch_encoder, encoder, layers.BatchNormalization(), # Refer to A.1 (Linear probing). layers.GlobalAveragePooling1D(), layers.Dense(NUM_CLASSES, activation="softmax"), ], name="linear_probe_model", ) # Only the final classification layer of the `downstream_model` should be trainable. for layer in downstream_model.layers[:-1]: layer.trainable = False downstream_model.summary() </code></pre></div> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">Model: "linear_probe_model" </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩ │ patches_1 (Patches) │ (None, 64, 108) │ 0 │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ patch_encoder_1 (PatchEncoder) │ (None, 64, 128) │ 22,144 │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ mae_encoder (Functional) │ (None, 64, 128) │ 1,981,696 │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ batch_normalization │ (None, 64, 128) │ 512 │ │ (BatchNormalization) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ global_average_pooling1d │ (None, 128) │ 0 │ │ (GlobalAveragePooling1D) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ dense_20 (Dense) │ (None, 10) │ 1,290 │ └─────────────────────────────────┴───────────────────────────┴────────────┘ </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Total params: 2,005,642 (7.65 MB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Trainable params: 1,290 (5.04 KB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Non-trainable params: 2,004,352 (7.65 MB) </pre> We are using average pooling to extract learned representations from the MAE encoder. Another approach would be to use a learnable dummy token inside the encoder during pretraining (resembling the [CLS] token). Then we can extract representations from that token during the downstream tasks. <h3 id="prepare-datasets-for-linear-probing">Prepare datasets for linear probing</h3> <div class="codehilite"><pre><code>def prepare_data(images, labels, is_train=True): if is_train: augmentation_model = train_augmentation_model else: augmentation_model = test_augmentation_model dataset = tf.data.Dataset.from_tensor_slices((images, labels)) if is_train: dataset = dataset.shuffle(BUFFER_SIZE) dataset = dataset.batch(BATCH_SIZE).map( lambda x, y: (augmentation_model(x), y), num_parallel_calls=AUTO ) return dataset.prefetch(AUTO) train_ds = prepare_data(x_train, y_train) val_ds = prepare_data(x_train, y_train, is_train=False) test_ds = prepare_data(x_test, y_test, is_train=False) </code></pre></div> <h3 id="perform-linear-probing">Perform linear probing</h3> <div class="codehilite"><pre><code>linear_probe_epochs = 50 linear_prob_lr = 0.1 warm_epoch_percentage = 0.1 steps = int((len(x_train) // BATCH_SIZE) * linear_probe_epochs) warmup_steps = int(steps * warm_epoch_percentage) scheduled_lrs = WarmUpCosine( learning_rate_base=linear_prob_lr, total_steps=steps, warmup_learning_rate=0.0, warmup_steps=warmup_steps, ) optimizer = keras.optimizers.SGD(learning_rate=scheduled_lrs, momentum=0.9) downstream_model.compile( optimizer=optimizer, loss="sparse_categorical_crossentropy", metrics=["accuracy"] ) downstream_model.fit(train_ds, validation_data=val_ds, epochs=linear_probe_epochs) loss, accuracy = downstream_model.evaluate(test_ds) accuracy = round(accuracy * 100, 2) print(f"Accuracy on the test set: {accuracy}%.") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/50 7/157 [37m━━━━━━━━━━━━━━━━━━━━ 3s 21ms/step - accuracy: 0.1183 - loss: 3.3939 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1700264823.481598 64012 device_compiler.h:187] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 157/157 ━━━━━━━━━━━━━━━━━━━━ 70s 242ms/step - accuracy: 0.1967 - loss: 2.6073 - val_accuracy: 0.3631 - val_loss: 1.7846 Epoch 2/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.3521 - loss: 1.8063 - val_accuracy: 0.3677 - val_loss: 1.7301 Epoch 3/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3580 - loss: 1.7580 - val_accuracy: 0.3649 - val_loss: 1.7326 Epoch 4/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3617 - loss: 1.7471 - val_accuracy: 0.3810 - val_loss: 1.7353 Epoch 5/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 35ms/step - accuracy: 0.3547 - loss: 1.7728 - val_accuracy: 0.3526 - val_loss: 1.8496 Epoch 6/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 35ms/step - accuracy: 0.3546 - loss: 1.7866 - val_accuracy: 0.3896 - val_loss: 1.7583 Epoch 7/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 37ms/step - accuracy: 0.3587 - loss: 1.7924 - val_accuracy: 0.3674 - val_loss: 1.7729 Epoch 8/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 38ms/step - accuracy: 0.3616 - loss: 1.7912 - val_accuracy: 0.3685 - val_loss: 1.7928 Epoch 9/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 36ms/step - accuracy: 0.3707 - loss: 1.7543 - val_accuracy: 0.3568 - val_loss: 1.7943 Epoch 10/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3719 - loss: 1.7451 - val_accuracy: 0.3859 - val_loss: 1.7230 Epoch 11/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3781 - loss: 1.7384 - val_accuracy: 0.3711 - val_loss: 1.7608 Epoch 12/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 35ms/step - accuracy: 0.3791 - loss: 1.7249 - val_accuracy: 0.4004 - val_loss: 1.6961 Epoch 13/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3818 - loss: 1.7303 - val_accuracy: 0.3501 - val_loss: 1.8506 Epoch 14/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3841 - loss: 1.7179 - val_accuracy: 0.3810 - val_loss: 1.8033 Epoch 15/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3818 - loss: 1.7172 - val_accuracy: 0.4168 - val_loss: 1.6507 Epoch 16/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 36ms/step - accuracy: 0.3851 - loss: 1.7059 - val_accuracy: 0.3806 - val_loss: 1.7581 Epoch 17/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3747 - loss: 1.7356 - val_accuracy: 0.4094 - val_loss: 1.6466 Epoch 18/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.3828 - loss: 1.7221 - val_accuracy: 0.4015 - val_loss: 1.6757 Epoch 19/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3889 - loss: 1.6939 - val_accuracy: 0.4102 - val_loss: 1.6392 Epoch 20/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3943 - loss: 1.6857 - val_accuracy: 0.4028 - val_loss: 1.6518 Epoch 21/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3870 - loss: 1.6970 - val_accuracy: 0.3949 - val_loss: 1.7283 Epoch 22/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3893 - loss: 1.6838 - val_accuracy: 0.4207 - val_loss: 1.6292 Epoch 23/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.4005 - loss: 1.6606 - val_accuracy: 0.4152 - val_loss: 1.6320 Epoch 24/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3978 - loss: 1.6556 - val_accuracy: 0.4042 - val_loss: 1.6657 Epoch 25/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4029 - loss: 1.6464 - val_accuracy: 0.4198 - val_loss: 1.6033 Epoch 26/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.3974 - loss: 1.6638 - val_accuracy: 0.4278 - val_loss: 1.5731 Epoch 27/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 37ms/step - accuracy: 0.4035 - loss: 1.6370 - val_accuracy: 0.4302 - val_loss: 1.5663 Epoch 28/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4027 - loss: 1.6349 - val_accuracy: 0.4458 - val_loss: 1.5349 Epoch 29/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4054 - loss: 1.6196 - val_accuracy: 0.4349 - val_loss: 1.5709 Epoch 30/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.4070 - loss: 1.6061 - val_accuracy: 0.4297 - val_loss: 1.5578 Epoch 31/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4105 - loss: 1.6172 - val_accuracy: 0.4250 - val_loss: 1.5735 Epoch 32/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4197 - loss: 1.5960 - val_accuracy: 0.4259 - val_loss: 1.5677 Epoch 33/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4156 - loss: 1.5989 - val_accuracy: 0.4400 - val_loss: 1.5395 Epoch 34/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.4214 - loss: 1.5862 - val_accuracy: 0.4486 - val_loss: 1.5237 Epoch 35/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4208 - loss: 1.5763 - val_accuracy: 0.4188 - val_loss: 1.5925 Epoch 36/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4227 - loss: 1.5803 - val_accuracy: 0.4525 - val_loss: 1.5174 Epoch 37/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4267 - loss: 1.5700 - val_accuracy: 0.4463 - val_loss: 1.5330 Epoch 38/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 37ms/step - accuracy: 0.4283 - loss: 1.5649 - val_accuracy: 0.4348 - val_loss: 1.5482 Epoch 39/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4332 - loss: 1.5581 - val_accuracy: 0.4486 - val_loss: 1.5251 Epoch 40/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4290 - loss: 1.5596 - val_accuracy: 0.4489 - val_loss: 1.5221 Epoch 41/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4318 - loss: 1.5589 - val_accuracy: 0.4494 - val_loss: 1.5202 Epoch 42/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4317 - loss: 1.5514 - val_accuracy: 0.4505 - val_loss: 1.5184 Epoch 43/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4353 - loss: 1.5504 - val_accuracy: 0.4561 - val_loss: 1.5081 Epoch 44/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4369 - loss: 1.5510 - val_accuracy: 0.4581 - val_loss: 1.5092 Epoch 45/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 35ms/step - accuracy: 0.4379 - loss: 1.5428 - val_accuracy: 0.4555 - val_loss: 1.5099 Epoch 46/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4421 - loss: 1.5475 - val_accuracy: 0.4579 - val_loss: 1.5073 Epoch 47/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4434 - loss: 1.5390 - val_accuracy: 0.4593 - val_loss: 1.5052 Epoch 48/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 34ms/step - accuracy: 0.4418 - loss: 1.5373 - val_accuracy: 0.4600 - val_loss: 1.5038 Epoch 49/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 6s 38ms/step - accuracy: 0.4400 - loss: 1.5367 - val_accuracy: 0.4596 - val_loss: 1.5045 Epoch 50/50 157/157 ━━━━━━━━━━━━━━━━━━━━ 5s 35ms/step - accuracy: 0.4448 - loss: 1.5321 - val_accuracy: 0.4595 - val_loss: 1.5048 40/40 ━━━━━━━━━━━━━━━━━━━━ 3s 71ms/step - accuracy: 0.4496 - loss: 1.5088 Accuracy on the test set: 44.66%. </code></pre></div> </div> We believe that with a more sophisticated hyperparameter tuning process and a longer pretraining it is possible to improve this performance further. For comparison, we took the encoder architecture and <a href="https://github.com/ariG23498/mae-scalable-vision-learners/blob/master/regular-classification.ipynb">trained it from scratch</a> in a fully supervised manner. This gave us ~76% test top-1 accuracy. The authors of MAE demonstrates strong performance on the ImageNet-1k dataset as well as other downstream tasks like object detection and semantic segmentation. <hr /> <h2 id="final-notes">Final notes</h2> We refer the interested readers to other examples on self-supervised learning present on keras.io: <ul> <li><a href="https://keras.io/examples/vision/semisupervised_simclr/">SimCLR</a></li> <li><a href="https://keras.io/examples/vision/nnclr">NNCLR</a></li> <li><a href="https://keras.io/examples/vision/simsiam">SimSiam</a></li> </ul> This idea of using BERT flavored pretraining in computer vision was also explored in <a href="https://arxiv.org/abs/1906.02940">Selfie</a>, but it could not demonstrate strong results. Another concurrent work that explores the idea of masked image modeling is <a href="https://arxiv.org/abs/2111.09886">SimMIM</a>. Finally, as a fun fact, we, the authors of this example also explored the idea of <a href="https://i.ibb.co/k5CpwDX/image.png">"reconstruction as a pretext task"</a> in 2020 but we could not prevent the network from representation collapse, and hence we did not get strong downstream performance. We would like to thank <a href="http://xinleic.xyz/">Xinlei Chen</a> (one of the authors of MAE) for helpful discussions. We are grateful to <a href="https://jarvislabs.ai/">JarvisLabs</a> and <a href="https://developers.google.com/programs/experts/">Google Developers Experts</a> program for helping with GPU credits. </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#masked-image-modeling-with-autoencoders'>Masked image modeling with Autoencoders</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#imports'>Imports</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#hyperparameters-for-pretraining'>Hyperparameters for pretraining</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#load-and-prepare-the-cifar10-dataset'>Load and prepare the CIFAR-10 dataset</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#data-augmentation'>Data augmentation</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#a-layer-for-extracting-patches-from-images'>A layer for extracting patches from images</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#patch-encoding-with-masking'>Patch encoding with masking</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#mlp'>MLP</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#mae-encoder'>MAE encoder</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#mae-decoder'>MAE decoder</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#mae-trainer'>MAE trainer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#model-initialization'>Model initialization</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#training-callbacks'>Training callbacks</a> </div> <div class='k-outline-depth-3'> <a href='#visualization-callback'>Visualization callback</a> </div> <div class='k-outline-depth-3'> <a href='#learning-rate-scheduler'>Learning rate scheduler</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#model-compilation-and-training'>Model compilation and training</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#evaluation-with-linear-probing'>Evaluation with linear probing</a> </div> <div class='k-outline-depth-3'> <a href='#extract-the-encoder-model-along-with-other-layers'>Extract the encoder model along with other layers</a> </div> <div class='k-outline-depth-3'> <a href='#prepare-datasets-for-linear-probing'>Prepare datasets for linear probing</a> </div> <div class='k-outline-depth-3'> <a href='#perform-linear-probing'>Perform linear probing</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#final-notes'>Final notes</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>