Foveated ViTs and adapting pre-trained models (DINOv3)

Last, we will show how to build a foveated ViT with our code-base, walk through our approach to adapting pre-trained models, and load in our adapted foveated variant of DINOv3

Foveated ViTs are much simpler to build than foveated CNNs. In a foveated ViT, we simply modify the sensor and patch embedding, and then the architecture is a standard ViT!

Let’s look at the patch embedding

Setup and execution instructions are in the notebook README. Run from the notebooks/ directory.

%load_ext autoreload
%autoreload 2

import matplotlib.pyplot as plt
import numpy as np
import torch
import matplotlib.cm as cm
import os

def visualize_patch_embed(patch_embed):
    in_coords = patch_embed.in_coords
    out_coords = patch_embed.out_coords
    in_cart_coords = in_coords.cartesian.cpu().numpy()
    in_plotting_coords = in_coords.plotting.cpu().numpy()
    out_cart_coords = out_coords.cartesian.cpu().numpy()
    out_plotting_coords = out_coords.plotting.cpu().numpy()
    knns = patch_embed.knn_indices.cpu().numpy()

    cmap = cm.get_cmap('tab20b')
    alpha = 0.5
    rf_scalar = 0.01

    # plot output unit coordinates
    fig, axs = plt.subplots(1, 2, figsize=(6, 2), dpi=800, gridspec_kw={'width_ratios': [2, 1]})
    axs[0].scatter(out_plotting_coords[:,0], out_plotting_coords[:,1], s=0.8, alpha=alpha)
    axs[1].scatter(out_cart_coords[:,0], out_cart_coords[:,1], s=0.8, alpha=alpha)
    rf_size = in_coords.polar[:,0].cpu().numpy()
    rf_size = rf_size + np.sort(rf_size)[1] # add second smallest value
    rf_size = rf_scalar*(rf_size / rf_size.min())

    # plot input kNNs
    for jj in range(knns.shape[1]):
        knn = knns[:,jj]
        mask = torch.logical_not(patch_embed.knn_indices_pad_mask[:,jj]).cpu()
        knn = knn[mask] # remove padding tokens
        np.random.seed(jj)
        color = cmap(np.random.randint(knns.shape[1]) / max(1, knns.shape[1] - 1))  # Normalize to [0,1] range
        axs[0].scatter(in_plotting_coords[knn,0], in_plotting_coords[knn,1], s=2, alpha=alpha, color=color, marker='o')
        axs[1].scatter(in_cart_coords[knn,0], in_cart_coords[knn,1], s=rf_size[knn], alpha=alpha, color=color, marker='o')

    for ax in axs:
        ax.axis('off')
        ax.axis('equal')
    axs[1].set_xlim(-1.1, 1.1)
    axs[1].set_ylim(-1.1, 1.1)

    plt.tight_layout()
    plt.show()

We will demonstrate three choices for patch embeddings.

• KNNPatchEmbedding: this is the most direct translation of a KNNConv layer to the patch embedding setting.

• PartitioningPatchEmbedding: this embedding yields a non-overlapping partition of the input coordinate space into output tokens. This produces some odd shapes around the boundaries.

• KNNPartitioningPatchEmbedding: this uses the partitioning solution from a PartitioningPatchEmbedding to discover the minimum geodesic distance on the sensor manifold, and thereby \(k\) value, needed for all input units to be assigned into output kNNs, producing slight overlap due to the circular nature of kNNs. We prefer this embedding due to the consistent patch shapes, and ease of parameterization relative to KNNPatchEmbedding.

from fovi.models.knnvit import KNNPartitioningPatchEmbedding

device = "cpu"
resolution = 64
fov = 16
cmf_a = 2

patch_embed = KNNPartitioningPatchEmbedding(
    3,
    384,
    resolution,
    fov,
    cmf_a,
    auto_match_cart_resources=1,
    in_cart_res=64,
    cart_patch_size=8,
    device=device,
    force_patches_less_than_matched=True,
    max_coord_val='auto',
    sample_cortex='geodesic',
)

visualize_patch_embed(patch_embed)

/home/nblauch/git/dex_unified_ws/fovi-isaaceye/fovi/arch/knn.py:139: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))
/home/nblauch/.venvs/dex-dev-dex_unified_ws/lib/python3.12/site-packages/torch/functional.py:505: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /pytorch/aten/src/ATen/native/TensorShape.cpp:4381.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
/home/nblauch/git/dex_unified_ws/fovi-isaaceye/fovi/arch/knn.py:252: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))
/tmp/ipykernel_677168/3433196042.py:10: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.
  cmap = cm.get_cmap('tab20b')
Ignoring fixed y limits to fulfill fixed data aspect with adjustable data limits.
Ignoring fixed x limits to fulfill fixed data aspect with adjustable data limits.

minimum k to use all inputs: 118

from fovi.models.knnvit import PartitioningPatchEmbedding
alt_patch_embed = PartitioningPatchEmbedding(
    3,
    384,
    resolution,
    fov,
    cmf_a,
    auto_match_cart_resources=1,
    in_cart_res=64,
    cart_patch_size=8,
    device=device,
    force_patches_less_than_matched=True,
    max_coord_val='auto',
    sample_cortex='geodesic',
)

visualize_patch_embed(alt_patch_embed)

/home/nblauch/git/dex_unified_ws/fovi-isaaceye/fovi/models/knnvit.py:208: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  self.k = torch.tensor(max_len)
/tmp/ipykernel_677168/3433196042.py:10: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.
  cmap = cm.get_cmap('tab20b')
Ignoring fixed y limits to fulfill fixed data aspect with adjustable data limits.
Ignoring fixed x limits to fulfill fixed data aspect with adjustable data limits.

minimum k to use all inputs: 81

load a pre-trained Foveated DINOv3

from fovi.models.loading import get_model_from_base_fn

base_fn = 'fovi-dinov3-splus_a-2.78_res-64_in1k'
model = get_model_from_base_fn(base_fn, device=device).eval()

adjusting FOV for fixation: 16.0 (full: 16.0)

/home/nblauch/git/dex_unified_ws/fovi-isaaceye/fovi/arch/knn.py:139: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))
/home/nblauch/git/dex_unified_ws/fovi-isaaceye/fovi/arch/knn.py:252: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))

minimum k to use all inputs: 103
Note: horizontal flip always done in the loader, to avoid differences across fixations
Number of coords per layer: [3976, 64]