Plot for Relative log amplitudes of Fourier transformed feature maps about how-do-vits-work HOT 5 CLOSED

import math
import matplotlib.pyplot as plt


def fourier(x):  # 2D Fourier transform
    f = torch.fft.fft2(x)
    f = f.abs() + 1e-6
    f = f.log()
    return f

def shift(x):
    b, c, h, w = x.shape
    return torch.roll(x, shifts=(int(h/2), int(w/2)), dims=(2, 3))

fig, ax = plt.subplots(1, 1, figsize=(3.3, 4))
for latent in latents:  # `latents` is a list of hidden feature maps in latent spaces
    if len(latent.shape) == 3:  # For ViT
        b, n, c = latent.shape
        h, w = int(math.sqrt(n)), int(math.sqrt(n))
        latent = latent.permute(0, 2, 1).reshape(b, c, h, w)
    elif len(latent.shape) == 4:  # For CNN
        b, c, h, w = latent.shape
    else:
        raise Exception("shape: %s" % str(latent.shape))
    latent = fourier(latent)
    latent = shift(latent).mean(dim=(0, 1))
    latent = latent.diag()[int(h/2):]  # Only use the half-diagonal components
    latent = latent - latent[0]  # Visualize 'relative' log amplitudes

    # Plot Fourier transformed relative log amplitudes
    freq = np.linspace(0, 1, len(latent))
    ax.plot(freq, latent)

import copy
import timm
import torch
import torch.nn as nn


# Divide the pretrained timm model into blocks.
name = 'vit_tiny_patch16_224'
model = timm.create_model(name, pretrained=True)


class PatchEmbed(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = copy.deepcopy(model)
        
    def forward(self, x, **kwargs):
        x = self.model.patch_embed(x)
        cls_token = self.model.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_token, x), dim=1)
        x = self.model.pos_drop(x + self.model.pos_embed)
        return x
    

class Residual(nn.Module):
    def __init__(self, *fn):
        super().__init__()
        self.fn = nn.Sequential(*fn)
        
    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x
    
    
class Lambda(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
        
    def forward(self, x):
        return self.fn(x)


def flatten(xs_list):
    return [x for xs in xs_list for x in xs]


# `blocks` is a sequence of blocks
blocks = [
    PatchEmbed(model),
    *flatten([[Residual(b.norm1, b.attn), Residual(b.norm2, b.mlp)] 
              for b in model.blocks]),
    nn.Sequential(Lambda(lambda x: x[:, 0]), model.norm, model.head),
]

# This snippet below build off https://github.com/facebookresearch/mae
import requests
import torch
import numpy as np

from PIL import Image
from einops import rearrange, reduce, repeat


imagenet_mean = np.array([0.485, 0.456, 0.406])
imagenet_std = np.array([0.229, 0.224, 0.225])

# Load an image
img_url = 'https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'
xs = Image.open(requests.get(img_url, stream=True).raw)
xs = xs.resize((224, 224))
xs = np.array(xs) / 255.

assert xs.shape == (224, 224, 3)

# Normalize by ImageNet mean and std
xs = xs - imagenet_mean
xs = xs / imagenet_std
xs = rearrange(torch.tensor(xs, dtype=torch.float32), 'h w c -> 1 c h w')

# Accumulate `latents` by collecting hidden states of a model
latents = []
with torch.no_grad():
    for block in blocks:
        xs = block(xs)
        latents.append(xs)
        
latents = [latent[:,1:] for latent in latents]  # Drop CLS token
latents = latents[:-1]  # Drop logit