<jupyter_start><jupyter_text>Visualization demo: Masked Autoencoders (MAE)In this notebook, we are going to visualize some predictions of a Vision Transformer (ViT) pre-trained on a very simple objective, namely masked patch prediction. The model needs to reconstruct pixel values for masked patches (as shown in the figure below).* Paper: https://arxiv.org/abs/2111.06377* Original repo (on which this notebook is based): https://github.com/facebookresearch/mae* Video explaining the paper: https://www.youtube.com/watch?v=Dp6iICL2dVI  Set-up environment<jupyter_code>!pip install -q git+https://github.com/huggingface/transformers.git<jupyter_output>Installing build dependencies ... [?25l[?25hdone
  Getting requirements to build wheel ... [?25l[?25hdone
    Preparing wheel metadata ... [?25l[?25hdone
[K     |████████████████████████████████| 67 kB 2.6 MB/s 
[K     |████████████████████████████████| 596 kB 11.0 MB/s 
[K     |████████████████████████████████| 6.8 MB 52.2 MB/s 
[K     |████████████████████████████████| 895 kB 73.6 MB/s 
[?25h  Building wheel for transformers (PEP 517) ... [?25l[?25hdone<jupyter_text>Prepare imageHere we apply some very basic image processing, namely resizing to 224x224 and normalizing the channels.<jupyter_code>from transformers import ViTFeatureExtractor
import requests
from PIL import Image

feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/vit-mae-base")
url = "https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image
pixel_values = feature_extractor(image, return_tensors="pt").pixel_values<jupyter_output><empty_output><jupyter_text>VisualizationNext, we forward the pixel values through the model. The encoder (which is a standard Vision Transformer) will first encode the visual patches. Next, a learnable mask token is added at the positions of the masked patches, and the decoder (which is also a Transformer) reconstructs the pixel values based on the encoded visual patches + mask tokens.The authors saw the best performance when masking out a large portion (75%) of the image patches.<jupyter_code>import torch
import numpy as np
import matplotlib.pyplot as plt

imagenet_mean = np.array(feature_extractor.image_mean)
imagenet_std = np.array(feature_extractor.image_std)

def show_image(image, title=''):
    # image is [H, W, 3]
    assert image.shape[2] == 3
    plt.imshow(torch.clip((image * imagenet_std + imagenet_mean) * 255, 0, 255).int())
    plt.title(title, fontsize=16)
    plt.axis('off')
    return

def visualize(pixel_values, model):
    # forward pass
    outputs = model(pixel_values)
    y = model.unpatchify(outputs.logits)
    y = torch.einsum('nchw->nhwc', y).detach().cpu()
    
    # visualize the mask
    mask = outputs.mask.detach()
    mask = mask.unsqueeze(-1).repeat(1, 1, model.config.patch_size**2 *3)  # (N, H*W, p*p*3)
    mask = model.unpatchify(mask)  # 1 is removing, 0 is keeping
    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()
    
    x = torch.einsum('nchw->nhwc', pixel_values)

    # masked image
    im_masked = x * (1 - mask)

    # MAE reconstruction pasted with visible patches
    im_paste = x * (1 - mask) + y * mask

    # make the plt figure larger
    plt.rcParams['figure.figsize'] = [24, 24]

    plt.subplot(1, 4, 1)
    show_image(x[0], "original")

    plt.subplot(1, 4, 2)
    show_image(im_masked[0], "masked")

    plt.subplot(1, 4, 3)
    show_image(y[0], "reconstruction")

    plt.subplot(1, 4, 4)
    show_image(im_paste[0], "reconstruction + visible")

    plt.show()
from transformers import ViTMAEForPreTraining

# make random mask reproducible (comment out to make it change)
torch.manual_seed(2)

model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")

visualize(pixel_values, model)<jupyter_output><empty_output><jupyter_text>After pre-trainingNote that, after pre-training, you can directly load the weights of the encoder into a `ViTForImageClassification` and start fine-tuning/doing linear probing. Note that a warning will be printed, indicating that the weights of the decoder aren't used.<jupyter_code>from transformers import ViTForImageClassification

model = ViTForImageClassification.from_pretrained("facebook/vit-mae-base")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13679137020885931, 'fraction_numerical': 0.03259123249942621, 'mean_word_length': 4.779840848806366, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 19, 'https://': 5, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '118877', 'license': 'MIT', 'n_tokens_mistral': 1648, 'n_tokens_neox': 1449, 'n_words': 397, 'path': 'ViTMAE/ViT_MAE_visualization_demo.ipynb', 'repo_name': 'FrancescoSaverioZuppichini/Transformers-Tutorials', 'size': '4357'}