{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ec6d7909",
   "metadata": {},
   "source": [
    "# Get activations from a foveated model\n",
    "\n",
    "Here we will demonstrate two methods for getting activitations. The first uses the model class directly. \n",
    "\n",
    "Let's load a pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5424f0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model with base_fn fovi-dinov3-splus_a-2.78_res-64_in1k not found in ../models\n",
      "Attempting to download fovi-dinov3-splus_a-2.78_res-64_in1k from HuggingFace Hub...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4568a95faa74c9abce173223af7a53a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model with base_fn fovi-dinov3-splus_a-2.78_res-64_in1k downloaded from HuggingFace Hub to /home/nblauch/.cache/fovi/fovi-dinov3-splus_a-2.78_res-64_in1k\n",
      "adjusting FOV for fixation: 16.0 (full: 16.0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nblauch/git/fovi/fovi/arch/knn.py:65: 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).\n",
      "  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))\n",
      "/home/nblauch/miniconda3/envs/fovi/lib/python3.9/site-packages/torch/functional.py:554: 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:4322.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n",
      "/home/nblauch/git/fovi/fovi/arch/knn.py:170: 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).\n",
      "  num_neighbors = torch.minimum(torch.tensor(self.k*m), torch.tensor(self.in_coords.shape[0]))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "minimum k to use all inputs: 103\n",
      "Note: horizontal flip always done in the loader, to avoid differences across fixations\n",
      "Number of coords per layer: [3976, 64]\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from fovi import get_model_from_base_fn\n",
    "from fovi.fovinet import FoviNet\n",
    "\n",
    "device = 'cuda'\n",
    "\n",
    "# base_fn = 'fovi-alexnet_a-1_res-64_rfmult-2_in1k'\n",
    "base_fn = 'fovi-dinov3-splus_a-2.78_res-64_in1k'\n",
    "model = get_model_from_base_fn(base_fn, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66eca93",
   "metadata": {},
   "source": [
    "### Now we can create some fake data and get activations.\n",
    "First, let's see which layers are available to hook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9f6fbca3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['',\n",
       " 'backbone',\n",
       " 'backbone.embeddings',\n",
       " 'backbone.embeddings.patch_embeddings',\n",
       " 'backbone.embeddings.patch_embeddings.parametrizations',\n",
       " 'backbone.embeddings.patch_embeddings.parametrizations.weight',\n",
       " 'backbone.embeddings.patch_embeddings.parametrizations.weight.0',\n",
       " 'backbone.rope_embeddings',\n",
       " 'backbone.layer',\n",
       " 'backbone.layer.0',\n",
       " 'backbone.layer.0.norm1',\n",
       " 'backbone.layer.0.attention',\n",
       " 'backbone.layer.0.attention.k_proj',\n",
       " 'backbone.layer.0.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.0.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.attention.v_proj',\n",
       " 'backbone.layer.0.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.0.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.attention.q_proj',\n",
       " 'backbone.layer.0.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.0.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.attention.o_proj',\n",
       " 'backbone.layer.0.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.0.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.layer_scale1',\n",
       " 'backbone.layer.0.drop_path',\n",
       " 'backbone.layer.0.norm2',\n",
       " 'backbone.layer.0.mlp',\n",
       " 'backbone.layer.0.mlp.gate_proj',\n",
       " 'backbone.layer.0.mlp.up_proj',\n",
       " 'backbone.layer.0.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.0.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.mlp.down_proj',\n",
       " 'backbone.layer.0.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.0.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.0.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.0.mlp.act_fn',\n",
       " 'backbone.layer.0.layer_scale2',\n",
       " 'backbone.layer.1',\n",
       " 'backbone.layer.1.norm1',\n",
       " 'backbone.layer.1.attention',\n",
       " 'backbone.layer.1.attention.k_proj',\n",
       " 'backbone.layer.1.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.1.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.attention.v_proj',\n",
       " 'backbone.layer.1.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.1.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.attention.q_proj',\n",
       " 'backbone.layer.1.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.1.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.attention.o_proj',\n",
       " 'backbone.layer.1.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.1.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.layer_scale1',\n",
       " 'backbone.layer.1.drop_path',\n",
       " 'backbone.layer.1.norm2',\n",
       " 'backbone.layer.1.mlp',\n",
       " 'backbone.layer.1.mlp.gate_proj',\n",
       " 'backbone.layer.1.mlp.up_proj',\n",
       " 'backbone.layer.1.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.1.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.mlp.down_proj',\n",
       " 'backbone.layer.1.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.1.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.1.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.1.mlp.act_fn',\n",
       " 'backbone.layer.1.layer_scale2',\n",
       " 'backbone.layer.2',\n",
       " 'backbone.layer.2.norm1',\n",
       " 'backbone.layer.2.attention',\n",
       " 'backbone.layer.2.attention.k_proj',\n",
       " 'backbone.layer.2.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.2.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.attention.v_proj',\n",
       " 'backbone.layer.2.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.2.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.attention.q_proj',\n",
       " 'backbone.layer.2.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.2.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.attention.o_proj',\n",
       " 'backbone.layer.2.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.2.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.layer_scale1',\n",
       " 'backbone.layer.2.drop_path',\n",
       " 'backbone.layer.2.norm2',\n",
       " 'backbone.layer.2.mlp',\n",
       " 'backbone.layer.2.mlp.gate_proj',\n",
       " 'backbone.layer.2.mlp.up_proj',\n",
       " 'backbone.layer.2.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.2.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.mlp.down_proj',\n",
       " 'backbone.layer.2.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.2.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.2.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.2.mlp.act_fn',\n",
       " 'backbone.layer.2.layer_scale2',\n",
       " 'backbone.layer.3',\n",
       " 'backbone.layer.3.norm1',\n",
       " 'backbone.layer.3.attention',\n",
       " 'backbone.layer.3.attention.k_proj',\n",
       " 'backbone.layer.3.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.3.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.attention.v_proj',\n",
       " 'backbone.layer.3.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.3.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.attention.q_proj',\n",
       " 'backbone.layer.3.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.3.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.attention.o_proj',\n",
       " 'backbone.layer.3.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.3.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.layer_scale1',\n",
       " 'backbone.layer.3.drop_path',\n",
       " 'backbone.layer.3.norm2',\n",
       " 'backbone.layer.3.mlp',\n",
       " 'backbone.layer.3.mlp.gate_proj',\n",
       " 'backbone.layer.3.mlp.up_proj',\n",
       " 'backbone.layer.3.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.3.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.mlp.down_proj',\n",
       " 'backbone.layer.3.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.3.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.3.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.3.mlp.act_fn',\n",
       " 'backbone.layer.3.layer_scale2',\n",
       " 'backbone.layer.4',\n",
       " 'backbone.layer.4.norm1',\n",
       " 'backbone.layer.4.attention',\n",
       " 'backbone.layer.4.attention.k_proj',\n",
       " 'backbone.layer.4.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.4.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.attention.v_proj',\n",
       " 'backbone.layer.4.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.4.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.attention.q_proj',\n",
       " 'backbone.layer.4.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.4.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.attention.o_proj',\n",
       " 'backbone.layer.4.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.4.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.layer_scale1',\n",
       " 'backbone.layer.4.drop_path',\n",
       " 'backbone.layer.4.norm2',\n",
       " 'backbone.layer.4.mlp',\n",
       " 'backbone.layer.4.mlp.gate_proj',\n",
       " 'backbone.layer.4.mlp.up_proj',\n",
       " 'backbone.layer.4.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.4.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.mlp.down_proj',\n",
       " 'backbone.layer.4.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.4.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.4.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.4.mlp.act_fn',\n",
       " 'backbone.layer.4.layer_scale2',\n",
       " 'backbone.layer.5',\n",
       " 'backbone.layer.5.norm1',\n",
       " 'backbone.layer.5.attention',\n",
       " 'backbone.layer.5.attention.k_proj',\n",
       " 'backbone.layer.5.attention.k_proj.parametrizations',\n",
       " 'backbone.layer.5.attention.k_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.attention.k_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.attention.v_proj',\n",
       " 'backbone.layer.5.attention.v_proj.parametrizations',\n",
       " 'backbone.layer.5.attention.v_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.attention.v_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.attention.q_proj',\n",
       " 'backbone.layer.5.attention.q_proj.parametrizations',\n",
       " 'backbone.layer.5.attention.q_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.attention.q_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.attention.o_proj',\n",
       " 'backbone.layer.5.attention.o_proj.parametrizations',\n",
       " 'backbone.layer.5.attention.o_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.attention.o_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.layer_scale1',\n",
       " 'backbone.layer.5.drop_path',\n",
       " 'backbone.layer.5.norm2',\n",
       " 'backbone.layer.5.mlp',\n",
       " 'backbone.layer.5.mlp.gate_proj',\n",
       " 'backbone.layer.5.mlp.up_proj',\n",
       " 'backbone.layer.5.mlp.up_proj.parametrizations',\n",
       " 'backbone.layer.5.mlp.up_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.mlp.up_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.mlp.down_proj',\n",
       " 'backbone.layer.5.mlp.down_proj.parametrizations',\n",
       " 'backbone.layer.5.mlp.down_proj.parametrizations.weight',\n",
       " 'backbone.layer.5.mlp.down_proj.parametrizations.weight.0',\n",
       " 'backbone.layer.5.mlp.act_fn',\n",
       " 'backbone.layer.5.layer_scale2',\n",
       " 'backbone.layer.6',\n",
       " 'backbone.layer.6.norm1',\n",
       " 'backbone.layer.6.attention',\n",
       " 'backbone.layer.6.attention.k_proj',\n",
       " 'backbone.layer.6.attention.v_proj',\n",
       " 'backbone.layer.6.attention.q_proj',\n",
       " 'backbone.layer.6.attention.o_proj',\n",
       " 'backbone.layer.6.layer_scale1',\n",
       " 'backbone.layer.6.drop_path',\n",
       " 'backbone.layer.6.norm2',\n",
       " 'backbone.layer.6.mlp',\n",
       " 'backbone.layer.6.mlp.gate_proj',\n",
       " 'backbone.layer.6.mlp.up_proj',\n",
       " 'backbone.layer.6.mlp.down_proj',\n",
       " 'backbone.layer.6.mlp.act_fn',\n",
       " 'backbone.layer.6.layer_scale2',\n",
       " 'backbone.layer.7',\n",
       " 'backbone.layer.7.norm1',\n",
       " 'backbone.layer.7.attention',\n",
       " 'backbone.layer.7.attention.k_proj',\n",
       " 'backbone.layer.7.attention.v_proj',\n",
       " 'backbone.layer.7.attention.q_proj',\n",
       " 'backbone.layer.7.attention.o_proj',\n",
       " 'backbone.layer.7.layer_scale1',\n",
       " 'backbone.layer.7.drop_path',\n",
       " 'backbone.layer.7.norm2',\n",
       " 'backbone.layer.7.mlp',\n",
       " 'backbone.layer.7.mlp.gate_proj',\n",
       " 'backbone.layer.7.mlp.up_proj',\n",
       " 'backbone.layer.7.mlp.down_proj',\n",
       " 'backbone.layer.7.mlp.act_fn',\n",
       " 'backbone.layer.7.layer_scale2',\n",
       " 'backbone.layer.8',\n",
       " 'backbone.layer.8.norm1',\n",
       " 'backbone.layer.8.attention',\n",
       " 'backbone.layer.8.attention.k_proj',\n",
       " 'backbone.layer.8.attention.v_proj',\n",
       " 'backbone.layer.8.attention.q_proj',\n",
       " 'backbone.layer.8.attention.o_proj',\n",
       " 'backbone.layer.8.layer_scale1',\n",
       " 'backbone.layer.8.drop_path',\n",
       " 'backbone.layer.8.norm2',\n",
       " 'backbone.layer.8.mlp',\n",
       " 'backbone.layer.8.mlp.gate_proj',\n",
       " 'backbone.layer.8.mlp.up_proj',\n",
       " 'backbone.layer.8.mlp.down_proj',\n",
       " 'backbone.layer.8.mlp.act_fn',\n",
       " 'backbone.layer.8.layer_scale2',\n",
       " 'backbone.layer.9',\n",
       " 'backbone.layer.9.norm1',\n",
       " 'backbone.layer.9.attention',\n",
       " 'backbone.layer.9.attention.k_proj',\n",
       " 'backbone.layer.9.attention.v_proj',\n",
       " 'backbone.layer.9.attention.q_proj',\n",
       " 'backbone.layer.9.attention.o_proj',\n",
       " 'backbone.layer.9.layer_scale1',\n",
       " 'backbone.layer.9.drop_path',\n",
       " 'backbone.layer.9.norm2',\n",
       " 'backbone.layer.9.mlp',\n",
       " 'backbone.layer.9.mlp.gate_proj',\n",
       " 'backbone.layer.9.mlp.up_proj',\n",
       " 'backbone.layer.9.mlp.down_proj',\n",
       " 'backbone.layer.9.mlp.act_fn',\n",
       " 'backbone.layer.9.layer_scale2',\n",
       " 'backbone.layer.10',\n",
       " 'backbone.layer.10.norm1',\n",
       " 'backbone.layer.10.attention',\n",
       " 'backbone.layer.10.attention.k_proj',\n",
       " 'backbone.layer.10.attention.v_proj',\n",
       " 'backbone.layer.10.attention.q_proj',\n",
       " 'backbone.layer.10.attention.o_proj',\n",
       " 'backbone.layer.10.layer_scale1',\n",
       " 'backbone.layer.10.drop_path',\n",
       " 'backbone.layer.10.norm2',\n",
       " 'backbone.layer.10.mlp',\n",
       " 'backbone.layer.10.mlp.gate_proj',\n",
       " 'backbone.layer.10.mlp.up_proj',\n",
       " 'backbone.layer.10.mlp.down_proj',\n",
       " 'backbone.layer.10.mlp.act_fn',\n",
       " 'backbone.layer.10.layer_scale2',\n",
       " 'backbone.layer.11',\n",
       " 'backbone.layer.11.norm1',\n",
       " 'backbone.layer.11.attention',\n",
       " 'backbone.layer.11.attention.k_proj',\n",
       " 'backbone.layer.11.attention.v_proj',\n",
       " 'backbone.layer.11.attention.q_proj',\n",
       " 'backbone.layer.11.attention.o_proj',\n",
       " 'backbone.layer.11.layer_scale1',\n",
       " 'backbone.layer.11.drop_path',\n",
       " 'backbone.layer.11.norm2',\n",
       " 'backbone.layer.11.mlp',\n",
       " 'backbone.layer.11.mlp.gate_proj',\n",
       " 'backbone.layer.11.mlp.up_proj',\n",
       " 'backbone.layer.11.mlp.down_proj',\n",
       " 'backbone.layer.11.mlp.act_fn',\n",
       " 'backbone.layer.11.layer_scale2',\n",
       " 'backbone.norm',\n",
       " 'projector',\n",
       " 'projector.layers',\n",
       " 'projector.layers.fc_block_6',\n",
       " 'projector.layers.fc_block_6.0',\n",
       " 'projector.layers.fc_block_6.1',\n",
       " 'projector.layers.fc_block_6.2',\n",
       " '',\n",
       " 'fix_projector',\n",
       " 'fix_projector.dropout',\n",
       " 'fix_projector.probe']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.list_available_layers()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c859038f",
   "metadata": {},
   "source": [
    "Let's hook the the fourth backbone block (layers.3), the full backbone (conv layers), and the projector (MLP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d070ab0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "inputs = torch.rand((10, 3, 256, 256)).to(device)\n",
    "outputs, acts = model.get_activations(inputs, layer_names=['backbone.layers.3', 'backbone', 'projector'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18daf1a9",
   "metadata": {},
   "source": [
    "Note that the intermediate backbone block retains a spatial dimension ($n=60$), whereas the full backbone has been globally pooled and has no spatial dimension, similarly to the projector.\n",
    "\n",
    "Note also that each activation tensor contains a fixation dimension as the second dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da0a2d3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'backbone.layers.3': torch.Size([10, 4, 1, 384]),\n",
       " 'backbone': torch.Size([10, 4, 1, 384]),\n",
       " 'projector': torch.Size([10, 4, 1024])}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{k: v.shape for k, v in acts.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508746f5",
   "metadata": {},
   "source": [
    "# Using the trainer class\n",
    "\n",
    "An even more stream-lined way of getting activations is to use the Trainer class. \n",
    "\n",
    "For this to work, you will need to define paths to existing dataset files. For now, these must be FFCV files. Soon, we will allow for standard image datasets. \n",
    "\n",
    "When loading a trainer from pre-trained, it is generally easiest to use the utility `get_trainer_from_base_fn`, which does a few basic things under the hood so we don't need to manually edit the config to turn off distributed training, etc. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a4f1204",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model with base_fn fovi-dinov3-splus_a-2.78_res-64_in1k not found in ../models\n",
      "Attempting to download fovi-dinov3-splus_a-2.78_res-64_in1k from HuggingFace Hub...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db63b0da35654973b00001956458d641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model with base_fn fovi-dinov3-splus_a-2.78_res-64_in1k downloaded from HuggingFace Hub to /home/nblauch/.cache/fovi/fovi-dinov3-splus_a-2.78_res-64_in1k\n",
      "adjusting FOV for fixation: 16.0 (full: 16.0)\n",
      "minimum k to use all inputs: 103\n",
      "Note: horizontal flip always done in the loader, to avoid differences across fixations\n",
      "Number of coords per layer: [3976, 64]\n",
      "FoviNet(\n",
      "  (network): BackboneProjectorWrapper(\n",
      "    (backbone): DINOv3ViTModel(\n",
      "      (embeddings): DINOv3ViTEmbeddings(\n",
      "        (patch_embeddings): ParametrizedKNNPartitioningPatchEmbedding(\n",
      "        \tin_channels=3\n",
      "        \tout_channels=384\n",
      "        \tk=103\n",
      "        \tn_ref=256\n",
      "        \tin_coords=SamplingCoords(length=3976, fov=16.0, cmf_a=2.785765, resolution=44, style=isotropic)\n",
      "        \tout_coords=SamplingCoords(length=64, fov=16.0, cmf_a=2.785765, resolution=6, style=isotropic)\n",
      "        \tsample_cortex=geodesic\n",
      "        )\n",
      "      )\n",
      "      (rope_embeddings): FoviDinoV3RoPE()\n",
      "      (layer): ModuleList(\n",
      "        (0-5): 6 x DINOv3ViTLayer(\n",
      "          (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "          (attention): DINOv3ViTAttention(\n",
      "            (k_proj): ParametrizedLinear(\n",
      "              in_features=384, out_features=384, bias=False\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "            (v_proj): ParametrizedLinear(\n",
      "              in_features=384, out_features=384, bias=True\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "            (q_proj): ParametrizedLinear(\n",
      "              in_features=384, out_features=384, bias=True\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "            (o_proj): ParametrizedLinear(\n",
      "              in_features=384, out_features=384, bias=True\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "          )\n",
      "          (layer_scale1): DINOv3ViTLayerScale()\n",
      "          (drop_path): Identity()\n",
      "          (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): DINOv3ViTGatedMLP(\n",
      "            (gate_proj): Linear(in_features=384, out_features=1536, bias=True)\n",
      "            (up_proj): ParametrizedLinear(\n",
      "              in_features=384, out_features=1536, bias=True\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "            (down_proj): ParametrizedLinear(\n",
      "              in_features=1536, out_features=384, bias=True\n",
      "              (parametrizations): ModuleDict(\n",
      "                (weight): ParametrizationList(\n",
      "                  (0): LoRAParam()\n",
      "                )\n",
      "              )\n",
      "            )\n",
      "            (act_fn): SiLUActivation()\n",
      "          )\n",
      "          (layer_scale2): DINOv3ViTLayerScale()\n",
      "        )\n",
      "        (6-11): 6 x DINOv3ViTLayer(\n",
      "          (norm1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "          (attention): DINOv3ViTAttention(\n",
      "            (k_proj): Linear(in_features=384, out_features=384, bias=False)\n",
      "            (v_proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "            (q_proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "            (o_proj): Linear(in_features=384, out_features=384, bias=True)\n",
      "          )\n",
      "          (layer_scale1): DINOv3ViTLayerScale()\n",
      "          (drop_path): Identity()\n",
      "          (norm2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): DINOv3ViTGatedMLP(\n",
      "            (gate_proj): Linear(in_features=384, out_features=1536, bias=True)\n",
      "            (up_proj): Linear(in_features=384, out_features=1536, bias=True)\n",
      "            (down_proj): Linear(in_features=1536, out_features=384, bias=True)\n",
      "            (act_fn): SiLUActivation()\n",
      "          )\n",
      "          (layer_scale2): DINOv3ViTLayerScale()\n",
      "        )\n",
      "      )\n",
      "      (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (projector): MLPWrapper(\n",
      "      (layers): Sequential(\n",
      "        (fc_block_6): LayerBlock(\n",
      "          (0): Dropout(p=0.5, inplace=False)\n",
      "          (1): Linear(in_features=384, out_features=1024, bias=False)\n",
      "          (2): ReLU(inplace=True)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (retinal_transform): RetinalTransform(\n",
      "    (foveal_color): GaussianColorDecay(sigma=None)\n",
      "    (sampler): GridSampler(fov=16.0, cmf_a=2.785765, style=isotropic, resolution=44, mode=nearest, n=3976)\n",
      "  )\n",
      "  (ssl_fixator): NoSaccadePolicy(\n",
      "    retinal_transform=RetinalTransform(\n",
      "    (foveal_color): GaussianColorDecay(sigma=None)\n",
      "    (sampler): GridSampler(fov=16.0, cmf_a=2.785765, style=isotropic, resolution=44, mode=nearest, n=3976)\n",
      "  ),\n",
      "    n_fixations=1\n",
      "  )\n",
      "  (sup_fixator): MultiRandomSaccadePolicy(\n",
      "    retinal_transform=RetinalTransform(\n",
      "    (foveal_color): GaussianColorDecay(sigma=None)\n",
      "    (sampler): GridSampler(fov=16.0, cmf_a=2.785765, style=isotropic, resolution=44, mode=nearest, n=3976)\n",
      "  ),\n",
      "    n_fixations=4,\n",
      "    nonrandom_first=1,\n",
      "    nonrandom_val=False,\n",
      "    crop_area_range=[1, 1],\n",
      "    add_aspect_variation=None,\n",
      "    val_crop_size=1,\n",
      "    norm_dist_from_center=0.25\n",
      "  )\n",
      "  (head): FoviNetProbe(\n",
      "    (fix_projector): LinearProbe(\n",
      "      (dropout): Dropout(p=0.5, inplace=False)\n",
      "      (probe): Linear(in_features=1024, out_features=1000, bias=True)\n",
      "    )\n",
      "  )\n",
      ")\n",
      "NUM PROBE LAYERS: 2\n",
      "n_fixations_val: [1, 2, 3, 5, 10, 20]\n",
      "just resizing, no crops\n",
      "output_size: (256, 256), scale: (1.0, 1.0), ratio: (1, 1)\n",
      "train loader: FlashLoader(\n",
      "\tData Path: /home/nblauch/data/ffcv/imagenet/train_compressed.ffcv\n",
      "\tBatch Size: 256\n",
      "\tOrder: OrderOption.QUASI_RANDOM\n",
      "\tNumber of Workers: 22\n",
      "\tOS Cache: 1\n",
      "\tDistributed: 0\n",
      "\tDrop Last: True\n",
      "\tRecompile: False\n",
      "\tAfter Batch Pipelines:\n",
      " {'image': Compose(\n",
      "    ToTorchImage(device=cuda, dtype=torch.float32, from_numpy=True)\n",
      "    RandomHorizontalFlip(p=0.5, seed=None)\n",
      ")}\n",
      ")\n",
      "val loader crop ratio: 1.0\n",
      "val loader: FlashLoader(\n",
      "\tData Path: /home/nblauch/data/ffcv/imagenet/val_compressed.ffcv\n",
      "\tBatch Size: 128\n",
      "\tOrder: OrderOption.SEQUENTIAL\n",
      "\tNumber of Workers: 22\n",
      "\tOS Cache: True\n",
      "\tDistributed: 0\n",
      "\tDrop Last: False\n",
      "\tRecompile: False\n",
      "\tAfter Batch Pipelines:\n",
      " {'image': Compose(\n",
      "    ToTorchImage(device=cuda, dtype=torch.float32, from_numpy=True)\n",
      "    NormalizeGPU(mean=tensor([0.4850, 0.4560, 0.4060], device='cuda:0', dtype=torch.float64), std=tensor([0.2290, 0.2240, 0.2250], device='cuda:0', dtype=torch.float64), inplace=True)\n",
      ")}\n",
      ")\n",
      "NUM TRAINING EXAMPLES: 1281167\n",
      "=> Logging in /home/nblauch/data/fovi/logs/None\n",
      "HydraConfig was not set\n",
      "skipping hydra directory copying\n",
      "Training backbone: True\n"
     ]
    }
   ],
   "source": [
    "from fovi import get_trainer_from_base_fn\n",
    "from fovi.paths import DATASETS_DIR\n",
    "\n",
    "# base_fn = 'fovi-alexnet_a-1_res-64_rfmult-2_in1k'\n",
    "base_fn = 'fovi-dinov3-splus_a-2.78_res-64_in1k'\n",
    "# edit the paths to those storing your ImageNet-1K FFCV files\n",
    "# in general, any kwarg you pass in will be used to update the loaded config file\n",
    "kwargs = {\n",
    "    'data.train_dataset': f'{DATASETS_DIR}/ffcv/imagenet/train_compressed.ffcv',\n",
    "    'data.val_dataset': f'{DATASETS_DIR}/ffcv/imagenet/val_compressed.ffcv',\n",
    "          }\n",
    "trainer = get_trainer_from_base_fn(base_fn, load=True, model_dirs=['../models'], **kwargs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5d06ca4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|█▎                                                                                                                                                                                | 3/391 [00:10<22:36,  3.50s/it]\n"
     ]
    }
   ],
   "source": [
    "outputs, activations, targets = trainer.compute_activations(trainer.val_loader, layer_names=['backbone.layers.3', 'backbone', 'projector'], max_batches=4, do_postproc=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "db8c1b13",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'backbone.layers.3': (512, 20, 1, 384),\n",
       " 'backbone': (512, 20, 1, 384),\n",
       " 'projector': (512, 20, 1024)}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{k: v.shape for k, v in activations.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86367c8f",
   "metadata": {},
   "source": [
    "note that we also now have the network outputs, which have been aggregated over fixations (since we passed `do_postproc=True`, which applies the fixation aggregator head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d79b820f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(512, 1000)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7a916a2",
   "metadata": {},
   "source": [
    "we can quickly check our top-1 accuracy (note: this is an unstable estimate since we used a small number of batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ed4a521c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.7305, device='cuda:0')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.val_meters['top_1_val'](torch.tensor(outputs), torch.tensor(targets))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (fovi)",
   "language": "python",
   "name": "fovi"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.25"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}