{ "cells": [ { "cell_type": "markdown", "id": "b0f0f4f3", "metadata": {}, "source": [ "# Import packages & functions" ] }, { "cell_type": "code", "execution_count": 1, "id": "5bad764b-45c1-45ce-a716-8d055e09821a", "metadata": { "tags": [] }, "outputs": [], "source": [ "import os\n", "import sys\n", "import json\n", "import argparse\n", "import numpy as np\n", "import math\n", "from einops import rearrange\n", "import time\n", "import random\n", "import string\n", "import h5py\n", "from tqdm import tqdm\n", "import webdataset as wds\n", "\n", "import matplotlib.pyplot as plt\n", "import torch\n", "import torch.nn as nn\n", "from torchvision import transforms\n", "from accelerate import Accelerator\n", "\n", "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n", "sys.path.append('generative_models/')\n", "import sgm\n", "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\n", "\n", "# tf32 data type is faster than standard float32\n", "torch.backends.cuda.matmul.allow_tf32 = True\n", "\n", "# custom functions #\n", "import utils" ] }, { "cell_type": "code", "execution_count": 31, "id": "f9cdfdb2-a81e-495f-a777-31acd22d9746", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Top-1 Precision: 0.00%\n" ] } ], "source": [ "import torch\n", "import torch.nn.functional as F\n", "\n", "def classPrecision(logits, y_true, top=1):\n", " \"\"\"\n", " Calculate the precision of the top-n predictions.\n", " \n", " Parameters:\n", " logits (torch.Tensor): The output logits from the model (shape: [batch_size, num_classes]).\n", " y_true (torch.Tensor): The ground truth labels (shape: [batch_size]).\n", " top (int): The number of top predictions to consider.\n", " \n", " Returns:\n", " float: The precision percentage of the top-n predictions.\n", " \"\"\"\n", " # Apply softmax to get probabilities\n", " probs = F.softmax(logits, dim=1).detach().cpu()\n", " \n", " # Get the top-n predictions\n", " top_n_preds = torch.topk(probs, top, dim=1).indices.detach().cpu()\n", "\n", " # Move y_true to CPU and detach\n", " y_true = y_true.detach().cpu()\n", "\n", " # Check if y_true is in top-n predictions\n", " correct = top_n_preds.eq(y_true.view(-1, 1).expand_as(top_n_preds))\n", "\n", " # Calculate precision\n", " precision = correct.sum().item() / y_true.size(0)\n", " \n", " return precision * 100\n", "\n", "# Example usage:\n", "logits = torch.randn(8, 41) # Example logits tensor\n", "y_true = torch.randint(0, 41, (8,)) # Example ground truth labels\n", "\n", "top_n_precision = classPrecision(logits, y_true, top=1)\n", "print(f\"Top-1 Precision: {top_n_precision:.2f}%\")\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "cc5d2e32-6027-4a19-bef4-5ca068db35bb", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "LOCAL RANK 0\n" ] } ], "source": [ "### Multi-GPU config ###\n", "local_rank = os.getenv('RANK')\n", "if local_rank is None: \n", " local_rank = 0\n", "else:\n", " local_rank = int(local_rank)\n", "print(\"LOCAL RANK \", local_rank) \n", "\n", "data_type = torch.float16 # change depending on your mixed_precision\n", "num_devices = torch.cuda.device_count()\n", "if num_devices==0: num_devices = 1\n", "\n", "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n", "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n", "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n", " global_batch_size = batch_size = 16\n", "else:\n", " global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n", " batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices" ] }, { "cell_type": "code", "execution_count": 4, "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "PID of this process = 5887\n", "device: cuda\n", "Distributed environment: DistributedType.NO\n", "Num processes: 1\n", "Process index: 0\n", "Local process index: 0\n", "Device: cuda\n", "\n", "Mixed precision type: fp16\n", "\n", "distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n" ] } ], "source": [ "print(\"PID of this process =\",os.getpid())\n", "device = accelerator.device\n", "print(\"device:\",device)\n", "world_size = accelerator.state.num_processes\n", "distributed = not accelerator.state.distributed_type == 'NO'\n", "num_devices = torch.cuda.device_count()\n", "if num_devices==0 or not distributed: num_devices = 1\n", "num_workers = num_devices\n", "print(accelerator.state)\n", "\n", "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\n", "print = accelerator.print # only print if local_rank=0" ] }, { "cell_type": "markdown", "id": "9018b82b-c054-4463-9527-4b0c2a75bda6", "metadata": { "tags": [] }, "source": [ "# Configurations" ] }, { "cell_type": "code", "execution_count": 5, "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "model_name: testing-tt3\n", "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=testing-tt3 --no-multi_subject --subj=1 --batch_size=16 --num_sessions=40 --hidden_dim=1024 --clip_scale=1. --no-blurry_recon --blur_scale=.5 --use_prior --prior_scale=30 --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug --ckpt_interval=999 --no-ckpt_saving --wandb_log\n" ] } ], "source": [ "# if running this interactively, can specify jupyter_args here for argparser to use\n", "if utils.is_interactive():\n", " model_name = \"testing-tt3\"\n", " print(\"model_name:\", model_name)\n", " \n", " # global_batch_size and batch_size should already be defined in the 2nd cell block\n", " jupyter_args = f\"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \\\n", " --cache_dir=/weka/proj-medarc/shared/cache \\\n", " --model_name={model_name} \\\n", " --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \\\n", " --hidden_dim=1024 --clip_scale=1. \\\n", " --no-blurry_recon --blur_scale=.5 \\\n", " --use_prior --prior_scale=30 \\\n", " --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n", " --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n", " # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n", "\n", " print(jupyter_args)\n", " jupyter_args = jupyter_args.split()\n", " \n", " from IPython.display import clear_output # function to clear print outputs in cell\n", " %load_ext autoreload \n", " # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n", " %autoreload 2 " ] }, { "cell_type": "code", "execution_count": 6, "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "subj_list [1] num_sessions 40\n" ] } ], "source": [ "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n", "parser.add_argument(\n", " \"--model_name\", type=str, default=\"testing2\",\n", " help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n", ")\n", "parser.add_argument(\n", " \"--data_path\", type=str, default=os.getcwd(),\n", " help=\"Path to where NSD data is stored / where to download it to\",\n", ")\n", "parser.add_argument(\n", " \"--cache_dir\", type=str, default=os.getcwd(),\n", " help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n", ")\n", "parser.add_argument(\n", " \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n", " help=\"Validate on which subject?\",\n", ")\n", "parser.add_argument(\n", " \"--multisubject_ckpt\", type=str, default=None,\n", " help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n", ")\n", "parser.add_argument(\n", " \"--num_sessions\", type=int, default=1,\n", " help=\"Number of training sessions to include\",\n", ")\n", "parser.add_argument(\n", " \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n", " help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n", ")\n", "parser.add_argument(\n", " \"--batch_size\", type=int, default=16,\n", " help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\n", ")\n", "parser.add_argument(\n", " \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n", " help=\"whether to log to wandb\",\n", ")\n", "parser.add_argument(\n", " \"--wandb_project\",type=str,default=\"stability\",\n", " help=\"wandb project name\",\n", ")\n", "parser.add_argument(\n", " \"--mixup_pct\",type=float,default=.33,\n", " help=\"proportion of way through training when to switch from BiMixCo to SoftCLIP\",\n", ")\n", "parser.add_argument(\n", " \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n", " help=\"whether to output blurry reconstructions\",\n", ")\n", "parser.add_argument(\n", " \"--blur_scale\",type=float,default=.5,\n", " help=\"multiply loss from blurry recons by this number\",\n", ")\n", "parser.add_argument(\n", " \"--clip_scale\",type=float,default=1.,\n", " help=\"multiply contrastive loss by this number\",\n", ")\n", "parser.add_argument(\n", " \"--prior_scale\",type=float,default=30,\n", " help=\"multiply diffusion prior loss by this\",\n", ")\n", "parser.add_argument(\n", " \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n", " help=\"whether to use image augmentation\",\n", ")\n", "parser.add_argument(\n", " \"--num_epochs\",type=int,default=150,\n", " help=\"number of epochs of training\",\n", ")\n", "parser.add_argument(\n", " \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n", ")\n", "parser.add_argument(\n", " \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n", ")\n", "parser.add_argument(\n", " \"--n_blocks\",type=int,default=4,\n", ")\n", "parser.add_argument(\n", " \"--hidden_dim\",type=int,default=1024,\n", ")\n", "parser.add_argument(\n", " \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n", ")\n", "parser.add_argument(\n", " \"--ckpt_saving\",action=argparse.BooleanOptionalAction,default=True,\n", ")\n", "parser.add_argument(\n", " \"--ckpt_interval\",type=int,default=5,\n", " help=\"save backup ckpt and reconstruct every x epochs\",\n", ")\n", "parser.add_argument(\n", " \"--seed\",type=int,default=42,\n", ")\n", "parser.add_argument(\n", " \"--max_lr\",type=float,default=3e-5,\n", ")\n", "\n", "if utils.is_interactive():\n", " args = parser.parse_args(jupyter_args)\n", "else:\n", " args = parser.parse_args()\n", "\n", "# create global variables without the args prefix\n", "for attribute_name in vars(args).keys():\n", " globals()[attribute_name] = getattr(args, attribute_name)\n", " \n", "# seed all random functions\n", "utils.seed_everything(seed)\n", "\n", "outdir = os.path.abspath(f'../train_logs/{model_name}')\n", "if not os.path.exists(outdir) and ckpt_saving:\n", " os.makedirs(outdir,exist_ok=True)\n", " \n", "if use_image_aug or blurry_recon:\n", " import kornia\n", " from kornia.augmentation.container import AugmentationSequential\n", "if use_image_aug:\n", " img_augment = AugmentationSequential(\n", " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n", " same_on_batch=False,\n", " data_keys=[\"input\"],\n", " )\n", " \n", "if multi_subject:\n", " subj_list = np.arange(1,9)\n", " subj_list = subj_list[subj_list != subj]\n", "else:\n", " subj_list = [subj]\n", "\n", "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)" ] }, { "cell_type": "code", "execution_count": 7, "id": "2fb18bf3-27f5-470c-be05-002215d391b9", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "1e-05" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "max_lr" ] }, { "cell_type": "markdown", "id": "42d13c25-1369-4c49-81d4-83d713586096", "metadata": { "tags": [] }, "source": [ "# Prep data, models, and dataloaders" ] }, { "cell_type": "markdown", "id": "1c023f24-5233-4a15-a2f5-78487b3a8546", "metadata": {}, "source": [ "### Creating wds dataloader, preload betas and all 73k possible images" ] }, { "cell_type": "code", "execution_count": 8, "id": "aefe7c27-ab39-4b2c-90f4-480f4087b7ab", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "dividing batch size by subj_list, which will then be concatenated across subj during training...\n", "batch_size = 16 num_iterations_per_epoch = 1875 num_samples_per_epoch = 30000\n" ] } ], "source": [ "def my_split_by_node(urls): return urls\n", "num_voxels_list = []\n", "\n", "if multi_subject:\n", " nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n", " num_samples_per_epoch = (750*40) // num_devices \n", "else:\n", " num_samples_per_epoch = (750*num_sessions) // num_devices \n", "\n", "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n", "batch_size = batch_size // len(subj_list)\n", "\n", "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n", "\n", "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)" ] }, { "cell_type": "code", "execution_count": 9, "id": "81084834-035f-4465-ad59-59e6b806a2f5", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Training with 40 sessions\n", "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar\n", "num_voxels for subj01: 15724\n", "Loaded all subj train dls and betas!\n", "\n", "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar\n", "Loaded test dl for subj1!\n", "\n" ] } ], "source": [ "train_data = {}\n", "train_dl = {}\n", "num_voxels = {}\n", "voxels = {}\n", "for s in subj_list:\n", " print(f\"Training with {num_sessions} sessions\")\n", " if multi_subject:\n", " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n", " else:\n", " train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n", " print(train_url)\n", " \n", " train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,nodesplitter=my_split_by_node)\\\n", " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", " .decode(\"torch\")\\\n", " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", " train_dl[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n", "\n", " f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n", " betas = f['betas'][:]\n", " betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n", " num_voxels_list.append(betas[0].shape[-1])\n", " num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n", " voxels[f'subj0{s}'] = betas\n", " print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n", "\n", "print(\"Loaded all subj train dls and betas!\\n\")\n", "\n", "# Validate only on one subject\n", "if multi_subject: \n", " subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n", "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n", " if subj==3:\n", " num_test=2113\n", " elif subj==4:\n", " num_test=1985\n", " elif subj==6:\n", " num_test=2113\n", " elif subj==8:\n", " num_test=1985\n", " else:\n", " num_test=2770\n", " test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n", "elif new_test: # using larger test set from after full dataset released\n", " if subj==3:\n", " num_test=2371\n", " elif subj==4:\n", " num_test=2188\n", " elif subj==6:\n", " num_test=2371\n", " elif subj==8:\n", " num_test=2188\n", " else:\n", " num_test=3000\n", " test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n", "print(test_url)\n", "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n", " .shuffle(750, initial=1500, rng=random.Random(42))\\\n", " .decode(\"torch\")\\\n", " .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n", " .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n", "test_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)\n", "print(f\"Loaded test dl for subj{subj}!\\n\")" ] }, { "cell_type": "code", "execution_count": 10, "id": "c13b4b84-094c-4b5b-bace-26c155aa6181", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loaded all 73k possible NSD images to cpu! (73000, 3, 224, 224)\n" ] } ], "source": [ "# Load 73k NSD images\n", "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n", "images = f['images']\n", "print(\"Loaded all 73k possible NSD images to cpu!\", images.shape)" ] }, { "cell_type": "markdown", "id": "10ec4517-dbdf-4ece-98f6-4714d5de4e15", "metadata": {}, "source": [ "## Load models" ] }, { "cell_type": "markdown", "id": "48d6160e-1ee8-4da7-a755-9dbb452a6fa5", "metadata": {}, "source": [ "### CLIP image embeddings model" ] }, { "cell_type": "code", "execution_count": 11, "id": "b0420dc0-199e-4c1a-857d-b1747058b467", "metadata": { "tags": [] }, "outputs": [], "source": [ "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n", " arch=\"ViT-bigG-14\",\n", " version=\"laion2b_s39b_b160k\",\n", " output_tokens=True,\n", " only_tokens=True,\n", ")\n", "clip_img_embedder.to(device)\n", "\n", "clip_seq_dim = 256\n", "clip_emb_dim = 1664" ] }, { "cell_type": "markdown", "id": "5b79bd38-6990-4504-8d45-4a68d57d8885", "metadata": {}, "source": [ "### SD VAE" ] }, { "cell_type": "code", "execution_count": 12, "id": "01baff79-8114-482b-b115-6f05aa8ad691", "metadata": { "tags": [] }, "outputs": [], "source": [ "if blurry_recon:\n", " from diffusers import AutoencoderKL \n", " autoenc = AutoencoderKL(\n", " down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n", " up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n", " block_out_channels=[128, 256, 512, 512],\n", " layers_per_block=2,\n", " sample_size=256,\n", " )\n", " ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n", " autoenc.load_state_dict(ckpt)\n", " \n", " autoenc.eval()\n", " autoenc.requires_grad_(False)\n", " autoenc.to(device)\n", " utils.count_params(autoenc)\n", " \n", " from autoencoder.convnext import ConvnextXL\n", " cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n", " cnx.requires_grad_(False)\n", " cnx.eval()\n", " cnx.to(device)\n", " \n", " mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n", " std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n", " \n", " blur_augs = AugmentationSequential(\n", " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n", " kornia.augmentation.RandomGrayscale(p=0.1),\n", " kornia.augmentation.RandomSolarize(p=0.1),\n", " kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n", " data_keys=[\"input\"],\n", " )" ] }, { "cell_type": "markdown", "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0", "metadata": {}, "source": [ "### MindEye modules" ] }, { "cell_type": "code", "execution_count": 13, "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5", "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "MindEyeModule()" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class MindEyeModule(nn.Module):\n", " def __init__(self):\n", " super(MindEyeModule, self).__init__()\n", " def forward(self, x):\n", " return x\n", " \n", "model = MindEyeModule()\n", "model" ] }, { "cell_type": "code", "execution_count": 14, "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "param counts:\n", "16,102,400 total\n", "16,102,400 trainable\n", "param counts:\n", "16,102,400 total\n", "16,102,400 trainable\n", "torch.Size([2, 1, 15724]) torch.Size([2, 1, 1024])\n" ] } ], "source": [ "class RidgeRegression(torch.nn.Module):\n", " # make sure to add weight_decay when initializing optimizer to enable regularization\n", " def __init__(self, input_sizes, out_features): \n", " super(RidgeRegression, self).__init__()\n", " self.out_features = out_features\n", " self.linears = torch.nn.ModuleList([\n", " torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n", " ])\n", " def forward(self, x, subj_idx):\n", " out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n", " return out\n", " \n", "class IndividRidgeRegression(torch.nn.Module):\n", " def __init__(self, input_size, out_features):\n", " super(IndividRidgeRegression, self).__init__()\n", " self.out_features = out_features\n", " self.linear = torch.nn.Linear(input_size, out_features)\n", " def forward(self, x):\n", " out = self.linear(x)\n", " return out\n", " \n", "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n", "utils.count_params(model.ridge)\n", "utils.count_params(model)\n", "\n", "# test on subject 1 with fake data\n", "b = torch.randn((2,1,num_voxels_list[0]))\n", "print(b.shape, model.ridge(b,0).shape)" ] }, { "cell_type": "code", "execution_count": 15, "id": "7b8de65a-6d3b-4248-bea9-9b6f4d562321", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "param counts:\n", "453,360,280 total\n", "453,360,280 trainable\n", "param counts:\n", "469,462,680 total\n", "469,462,680 trainable\n", "b.shape torch.Size([2, 1, 1024])\n", "torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1])\n" ] } ], "source": [ "from models import BrainNetwork\n", "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n", " clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n", " blurry_recon=blurry_recon, clip_scale=clip_scale)\n", "utils.count_params(model.backbone)\n", "utils.count_params(model)\n", "\n", "# test that the model works on some fake data\n", "b = torch.randn((2,1,hidden_dim))\n", "print(\"b.shape\",b.shape)\n", "\n", "backbone_, clip_, blur_ = model.backbone(b)\n", "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)" ] }, { "cell_type": "markdown", "id": "e63b26a9-3a36-4638-be3c-dbf28705dd76", "metadata": {}, "source": [ "### Load semantic clusters" ] }, { "cell_type": "code", "execution_count": 16, "id": "2950add9-3828-4d55-bbf7-0d9f3519de5c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "semantic_cluster_onehot.shape torch.Size([73000, 41])\n", "num_seman_clusters 41\n" ] } ], "source": [ "path_semantic_names = \"/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy\"\n", "path_semantic_cluster = \"/weka/proj-fmri/ckadirt/MindEyeV2/src/COCO_73k_semantic_cluster.npy\"\n", "semantic_cluster_names = np.load(path_semantic_names)\n", "semantic_cluster = np.load(path_semantic_cluster)\n", "possible_semantic_clusters = np.unique(semantic_cluster)\n", "\n", "# one-hot encode semantic clusters\n", "# move possible_semantic_clusters to numbers and create a dictionary\n", "semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}\n", "semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))\n", "for i, cluster in enumerate(semantic_cluster):\n", " semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1\n", "\n", "\n", "print(\"semantic_cluster_onehot.shape\", semantic_cluster_onehot.shape)\n", "\n", "num_seman_clusters = len(np.unique(semantic_cluster))\n", "print(\"num_seman_clusters\", num_seman_clusters)" ] }, { "cell_type": "markdown", "id": "b397c0d7-52a3-4153-823b-c27d2eb3eeba", "metadata": {}, "source": [ "### Adding the ridge regression to the class" ] }, { "cell_type": "code", "execution_count": 17, "id": "69965344-9346-4592-9cc5-e537e31d5fce", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "param counts:\n", "17,465,385 total\n", "17,465,385 trainable\n", "param counts:\n", "486,928,065 total\n", "486,928,065 trainable\n" ] }, { "data": { "text/plain": [ "486928065" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# if use_prior:\n", "# from models import *\n", "\n", "# # setup diffusion prior network\n", "# out_dim = clip_emb_dim\n", "# depth = 6\n", "# dim_head = 52\n", "# heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n", "# timesteps = 100\n", "\n", "# prior_network = PriorNetwork(\n", "# dim=out_dim,\n", "# depth=depth,\n", "# dim_head=dim_head,\n", "# heads=heads,\n", "# causal=False,\n", "# num_tokens = clip_seq_dim,\n", "# learned_query_mode=\"pos_emb\"\n", "# )\n", "\n", "# model.diffusion_prior = BrainDiffusionPrior(\n", "# net=prior_network,\n", "# image_embed_dim=out_dim,\n", "# condition_on_text_encodings=False,\n", "# timesteps=timesteps,\n", "# cond_drop_prob=0.2,\n", "# image_embed_scale=None,\n", "# )\n", " \n", "# utils.count_params(model.diffusion_prior)\n", "# utils.count_params(model)\n", "\n", "model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)\n", "utils.count_params(model.RRClassifier)\n", "utils.count_params(model)" ] }, { "cell_type": "markdown", "id": "ec25271a-2209-400c-8026-df3b8ddc1eef", "metadata": {}, "source": [ "### Setup optimizer / lr / ckpt saving" ] }, { "cell_type": "code", "execution_count": 18, "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "total_steps 281250\n", "\n", "Done with model preparations!\n", "param counts:\n", "486,928,065 total\n", "486,928,065 trainable\n" ] } ], "source": [ "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n", "\n", "opt_grouped_parameters = [\n", " {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},\n", " {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", " {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n", " {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},\n", "]\n", "# if use_prior:\n", "# opt_grouped_parameters.extend([\n", "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n", "# {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n", "# ])\n", "# opt_grouped_parameters.extend([\n", "# \n", "# ])\n", "\n", "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n", "\n", "if lr_scheduler_type == 'linear':\n", " lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n", " optimizer,\n", " total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n", " last_epoch=-1\n", " )\n", "elif lr_scheduler_type == 'cycle':\n", " total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n", " print(\"total_steps\", total_steps)\n", " lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n", " optimizer, \n", " max_lr=max_lr,\n", " total_steps=total_steps,\n", " final_div_factor=1000,\n", " last_epoch=-1, pct_start=2/num_epochs\n", " )\n", " \n", "def save_ckpt(tag):\n", " ckpt_path = outdir+f'/{tag}.pth'\n", " if accelerator.is_main_process:\n", " unwrapped_model = accelerator.unwrap_model(model)\n", " torch.save({\n", " 'epoch': epoch,\n", " 'model_state_dict': unwrapped_model.state_dict(),\n", " 'optimizer_state_dict': optimizer.state_dict(),\n", " 'lr_scheduler': lr_scheduler.state_dict(),\n", " 'train_losses': losses,\n", " 'test_losses': test_losses,\n", " 'lrs': lrs,\n", " }, ckpt_path)\n", " print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n", "\n", "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n", " print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n", " checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n", " state_dict = checkpoint['model_state_dict']\n", " if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n", " state_dict.pop('ridge.linears.0.weight',None)\n", " model.load_state_dict(state_dict, strict=strict)\n", " if load_epoch:\n", " globals()[\"epoch\"] = checkpoint['epoch']\n", " print(\"Epoch\",epoch)\n", " if load_optimizer:\n", " optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n", " if load_lr:\n", " lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n", " del checkpoint\n", "\n", "print(\"\\nDone with model preparations!\")\n", "num_params = utils.count_params(model)" ] }, { "cell_type": "markdown", "id": "983f458b-35b8-49f2-b6db-80296cece730", "metadata": {}, "source": [ "# Weights and Biases" ] }, { "cell_type": "code", "execution_count": 19, "id": "0a25a662-daa8-4de9-9233-8364800fcb6b", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "wandb mindeye_semantic_cluster run testing-tt3\n", "wandb_config:\n", " {'model_name': 'testing-tt3', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 486928065, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}\n", "wandb_id: testing-tt3\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mckadirt\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n" ] }, { "data": { "text/html": [ "wandb version 0.17.4 is available! To upgrade, please run:\n", " $ pip install wandb --upgrade" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Tracking run with wandb version 0.17.1" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Run data is saved locally in /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_011732-testing-tt3" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Syncing run testing-tt3 to Weights & Biases (docs)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ " View project at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ " View run at https://stability.wandb.io/ckadirt/mindeye_semantic_cluster/runs/testing-tt3" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", " import wandb\n", " wandb_project = 'mindeye_semantic_cluster'\n", " print(f\"wandb {wandb_project} run {model_name}\")\n", " # need to configure wandb beforehand in terminal with \"wandb init\"!\n", " wandb_config = {\n", " \"model_name\": model_name,\n", " \"global_batch_size\": global_batch_size,\n", " \"batch_size\": batch_size,\n", " \"num_epochs\": num_epochs,\n", " \"num_sessions\": num_sessions,\n", " \"num_params\": num_params,\n", " \"clip_scale\": clip_scale,\n", " \"prior_scale\": prior_scale,\n", " \"blur_scale\": blur_scale,\n", " \"use_image_aug\": use_image_aug,\n", " \"max_lr\": max_lr,\n", " \"mixup_pct\": mixup_pct,\n", " \"num_samples_per_epoch\": num_samples_per_epoch,\n", " \"num_test\": num_test,\n", " \"ckpt_interval\": ckpt_interval,\n", " \"ckpt_saving\": ckpt_saving,\n", " \"seed\": seed,\n", " \"distributed\": distributed,\n", " \"num_devices\": num_devices,\n", " \"world_size\": world_size,\n", " \"train_url\": train_url,\n", " \"test_url\": test_url,\n", " }\n", " print(\"wandb_config:\\n\",wandb_config)\n", " print(\"wandb_id:\",model_name)\n", " wandb.login(host='https://stability.wandb.io')\n", " wandb.init(\n", " id=model_name,\n", " project=wandb_project,\n", " name=model_name,\n", " config=wandb_config,\n", " resume=\"allow\",\n", " )\n", "else:\n", " wandb_log = False" ] }, { "cell_type": "markdown", "id": "d5690151-2131-4918-b750-e869cbd1a8a8", "metadata": {}, "source": [ "# Main" ] }, { "cell_type": "code", "execution_count": 20, "id": "12de6387-6e18-4e4b-b5ce-a847d625330a", "metadata": { "tags": [] }, "outputs": [], "source": [ "epoch = 0\n", "losses, test_losses, lrs = [], [], []\n", "best_test_loss = 1e9\n", "torch.cuda.empty_cache()" ] }, { "cell_type": "code", "execution_count": 21, "id": "607a7c7b-fe5e-41a4-80bf-d2814b3a57cc", "metadata": { "tags": [] }, "outputs": [], "source": [ "# load multisubject stage1 ckpt if set\n", "if multisubject_ckpt is not None:\n", " load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)" ] }, { "cell_type": "code", "execution_count": 22, "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4", "metadata": { "tags": [] }, "outputs": [], "source": [ "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n", "\n", "model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)\n", "# leaving out test_dl since we will only have local_rank 0 device do evals" ] }, { "cell_type": "code", "execution_count": 23, "id": "e16fa583-064c-4cfa-9f47-08d08f53f504", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1875\n" ] } ], "source": [ "print(num_iterations_per_epoch)" ] }, { "cell_type": "code", "execution_count": 28, "id": "60be0d5f-3e94-4612-9373-61b53d836393", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "testing-tt3 starting with epoch 4 / 150\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/146 [00:00 90\u001b[0m image \u001b[38;5;241m=\u001b[39m \u001b[43mimage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 92\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m use_image_aug: \n\u001b[1;32m 93\u001b[0m image \u001b[38;5;241m=\u001b[39m img_augment(image)\n", "\u001b[0;31mKeyboardInterrupt\u001b[0m: " ] } ], "source": [ "print(f\"{model_name} starting with epoch {epoch} / {num_epochs}\")\n", "progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))\n", "test_image, test_voxel = None, None\n", "mse = nn.MSELoss()\n", "l1 = nn.L1Loss()\n", "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n", "\n", "for epoch in progress_bar:\n", " model.train()\n", "\n", " fwd_percent_correct = 0.\n", " bwd_percent_correct = 0.\n", " test_fwd_percent_correct = 0.\n", " test_bwd_percent_correct = 0.\n", " \n", " recon_cossim = 0.\n", " test_recon_cossim = 0.\n", " recon_mse = 0.\n", " test_recon_mse = 0.\n", "\n", " loss_clip_total = 0.\n", " loss_blurry_total = 0.\n", " loss_blurry_cont_total = 0.\n", " test_loss_clip_total = 0.\n", " \n", " loss_prior_total = 0.\n", " test_loss_prior_total = 0.\n", " \n", " loss_RR_total = 0.\n", " test_loss_RR_total = 0.\n", "\n", " blurry_pixcorr = 0.\n", " test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1\n", "\n", " class_precisions_1 = 0\n", " test_class_precisions_1 = 0\n", "\n", " class_precisions_5 = 0\n", " test_class_precisions_5 = 0\n", "\n", " class_precisions_10 = 0\n", " test_class_precisions_10 = 0\n", "\n", " # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)\n", " voxel_iters = {} # empty dict because diff subjects have differing # of voxels\n", " image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()\n", " annot_iters = {}\n", " perm_iters, betas_iters, select_iters = {}, {}, {}\n", " images_indexes = {}\n", " for s, train_dl in enumerate(train_dls):\n", " with torch.cuda.amp.autocast(dtype=data_type):\n", " iter = -1\n", " for behav0, past_behav0, future_behav0, old_behav0 in train_dl: \n", " # Load images to cpu from hdf5 (requires sorted indexing)\n", " image_idx = behav0[:,0,0].cpu().long().numpy()\n", "\n", " image0, image_sorted_idx = np.unique(image_idx, return_index=True) \n", " if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing\n", " continue\n", " iter += 1\n", " image0 = torch.tensor(images[image0], dtype=data_type)\n", " image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n", " images_indexes[f\"subj0{s}_iter{iter}\"] = image_sorted_idx\n", " \n", " # Load voxels for current batch, matching above indexing\n", " voxel_idx = behav0[:,0,5].cpu().long().numpy()\n", " voxel_sorted_idx = voxel_idx[image_sorted_idx]\n", " voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]\n", " voxel0 = torch.Tensor(voxel0).unsqueeze(1)\n", "\n", " if epoch < int(mixup_pct * num_epochs):\n", " voxel0, perm, betas, select = utils.mixco(voxel0)\n", " perm_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = perm\n", " betas_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = betas\n", " select_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = select\n", "\n", " voxel_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = voxel0\n", "\n", " if iter >= num_iterations_per_epoch-1:\n", " break\n", "\n", " # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each\n", " for train_i in range(num_iterations_per_epoch):\n", " with torch.cuda.amp.autocast(dtype=data_type):\n", " optimizer.zero_grad()\n", " loss=0.\n", "\n", " voxel_list = [voxel_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n", " image = image_iters[train_i].detach()\n", " image = image.to(device)\n", "\n", " if use_image_aug: \n", " image = img_augment(image)\n", "\n", " clip_target = clip_img_embedder(image)\n", " assert not torch.any(torch.isnan(clip_target))\n", "\n", " if epoch < int(mixup_pct * num_epochs):\n", " perm_list = [perm_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n", " perm = torch.cat(perm_list, dim=0)\n", " betas_list = [betas_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n", " betas = torch.cat(betas_list, dim=0)\n", " select_list = [select_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n", " select = torch.cat(select_list, dim=0)\n", "\n", " voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]\n", " voxel_ridge = torch.cat(voxel_ridge_list, dim=0)\n", "\n", " backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n", "\n", " if clip_scale>0:\n", " clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n", " clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n", "\n", " # if use_prior:\n", " # loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)\n", " # loss_prior_total += loss_prior.item()\n", " # loss_prior *= prior_scale\n", " # loss += loss_prior\n", "\n", " # recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()\n", " # recon_mse += mse(prior_out, clip_target).item()\n", " print(backbone.shape)\n", " logits = model.RRClassifier(backbone.flatten(1))\n", " print(logits.shape)\n", " print(semantic_cluster[images_indexes[f\"subj0{s}_iter{train_i}\"]])\n", " print(torch.Tensor([semantic_cluster_dict[i] for i in semantic_cluster[images_indexes[f\"subj0{s}_iter{train_i}\"]]]))\n", " #print(logits.shape, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).shape)\n", " #print(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1))\n", " loss_RR = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device))\n", " #print(\"backbone.shape\",backbone.shape, \"clip_voxels.shape\",clip_voxels.shape, \"blurry_image_enc_[0].shape\",blurry_image_enc_[0].shape, \"blurry_image_enc_[1].shape\",blurry_image_enc_[1].shape)\n", " #something \n", "\n", " loss_RR_total += loss_RR.item()\n", " loss += loss_RR\n", "\n", " if (torch.rand(1) < 0.03).item():\n", " print(\"loss_RR\", loss_RR.item())\n", "\n", " if clip_scale>0:\n", " if epoch < int(mixup_pct * num_epochs): \n", " loss_clip = utils.mixco_nce(\n", " clip_voxels_norm,\n", " clip_target_norm,\n", " temp=.006,\n", " perm=perm, betas=betas, select=select)\n", " else:\n", " epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]\n", " loss_clip = utils.soft_clip_loss(\n", " clip_voxels_norm,\n", " clip_target_norm,\n", " temp=epoch_temp)\n", "\n", " loss_clip_total += loss_clip.item()\n", " loss_clip *= clip_scale\n", " # loss += loss_clip\n", "\n", " if blurry_recon: \n", " image_enc_pred, transformer_feats = blurry_image_enc_\n", "\n", " image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215\n", " loss_blurry = l1(image_enc_pred, image_enc)\n", " loss_blurry_total += loss_blurry.item()\n", "\n", " if epoch < int(mixup_pct * num_epochs):\n", " image_enc_shuf = image_enc[perm]\n", " betas_shape = [-1] + [1]*(len(image_enc.shape)-1)\n", " image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \\\n", " image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)\n", "\n", " image_norm = (image - mean)/std\n", " image_aug = (blur_augs(image) - mean)/std\n", " _, cnx_embeds = cnx(image_norm)\n", " _, cnx_aug_embeds = cnx(image_aug)\n", "\n", " cont_loss = utils.soft_cont_loss(\n", " nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),\n", " nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n", " nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n", " temp=0.2)\n", " loss_blurry_cont_total += cont_loss.item()\n", "\n", " # loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215\n", "\n", " if clip_scale>0:\n", " # forward and backward top 1 accuracy \n", " labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) \n", " fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()\n", " bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n", "\n", " if blurry_recon:\n", " with torch.no_grad():\n", " # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()\n", " random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n", " blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)\n", " pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n", " blurry_pixcorr += pixcorr.item()\n", "\n", " class_precisions_1 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device))\n", " class_precisions_5 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device), 5)\n", " class_precisions_10 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device), 10)\n", "\n", " utils.check_loss(loss)\n", " accelerator.backward(loss)\n", " optimizer.step()\n", "\n", " losses.append(loss.item())\n", " lrs.append(optimizer.param_groups[0]['lr'])\n", "\n", " if lr_scheduler_type is not None:\n", " lr_scheduler.step()\n", "\n", " model.eval()\n", " if local_rank==0:\n", " with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): \n", " for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl): \n", " # all test samples should be loaded per batch such that test_i should never exceed 0\n", " assert len(behav) == num_test\n", "\n", " ## Average same-image repeats ##\n", " if test_image is None:\n", " voxel = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)\n", " \n", " image = behav[:,0,0].cpu().long()\n", "\n", " unique_image, sort_indices = torch.unique(image, return_inverse=True)\n", " for im in unique_image:\n", " locs = torch.where(im == image)[0]\n", " if len(locs)==1:\n", " locs = locs.repeat(3)\n", " elif len(locs)==2:\n", " locs = locs.repeat(2)[:3]\n", " assert len(locs)==3\n", " if test_image is None:\n", " test_image = torch.Tensor(images[im][None])\n", " test_voxel = voxel[locs][None]\n", " else:\n", " test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))\n", " test_voxel = torch.vstack((test_voxel, voxel[locs][None]))\n", "\n", " loss=0.\n", " \n", " test_indices = torch.arange(len(test_voxel))[:300]\n", " voxel = test_voxel[test_indices].to(device)\n", " image = test_image[test_indices].to(device)\n", " assert len(image) == 300\n", "\n", " clip_target = clip_img_embedder(image.float())\n", "\n", " for rep in range(3):\n", " voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list\n", " backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)\n", "\n", " logits0 = model.RRClassifier(backbone0.flatten(1))\n", "\n", " if rep==0:\n", " clip_voxels = clip_voxels0\n", " backbone = backbone0\n", " logits = logits0\n", " else:\n", " clip_voxels += clip_voxels0\n", " backbone += backbone0\n", " logits += logits0\n", " clip_voxels /= 3\n", " backbone /= 3\n", " logits /= 3\n", "\n", " print(logits.shape, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).shape)\n", " RR_loss = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))\n", " test_loss_RR_total += RR_loss.item()\n", " loss += RR_loss\n", "\n", " if clip_scale>0:\n", " clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n", " clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n", " \n", " # for some evals, only doing a subset of the samples per batch because of computational cost\n", " random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n", " \n", " # if use_prior:\n", " # loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])\n", " # test_loss_prior_total += loss_prior.item()\n", " # loss_prior *= prior_scale\n", " # loss += loss_prior\n", " \n", " if clip_scale>0:\n", " loss_clip = utils.soft_clip_loss(\n", " clip_voxels_norm,\n", " clip_target_norm,\n", " temp=.006)\n", "\n", " test_loss_clip_total += loss_clip.item()\n", " loss_clip = loss_clip * clip_scale\n", " loss += loss_clip\n", "\n", " if blurry_recon:\n", " image_enc_pred, _ = blurry_image_enc_\n", " blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)\n", " pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n", " test_blurry_pixcorr += pixcorr.item()\n", "\n", " if clip_scale>0:\n", " # forward and backward top 1 accuracy \n", " labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) \n", " test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()\n", " test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n", "\n", " test_class_precisions_1 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))\n", " test_class_precisions_5 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device), 5)\n", " test_class_precisions_10 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device), 10)\n", "\n", " \n", " utils.check_loss(loss) \n", " test_losses.append(loss.item())\n", "\n", " assert (test_i+1) == 1\n", " logs = {\"train/loss\": np.mean(losses[-(train_i+1):]),\n", " \"test/loss\": np.mean(test_losses[-(test_i+1):]),\n", " \"train/lr\": lrs[-1],\n", " \"train/num_steps\": len(losses),\n", " \"test/num_steps\": len(test_losses),\n", " \"train/fwd_pct_correct\": fwd_percent_correct / (train_i + 1),\n", " \"train/bwd_pct_correct\": bwd_percent_correct / (train_i + 1),\n", " \"test/test_fwd_pct_correct\": test_fwd_percent_correct / (test_i + 1),\n", " \"test/test_bwd_pct_correct\": test_bwd_percent_correct / (test_i + 1),\n", " \"train/loss_clip_total\": loss_clip_total / (train_i + 1),\n", " \"train/loss_blurry_total\": loss_blurry_total / (train_i + 1),\n", " \"train/loss_blurry_cont_total\": loss_blurry_cont_total / (train_i + 1),\n", " \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n", " \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n", " \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n", " \"train/recon_cossim\": recon_cossim / (train_i + 1),\n", " \"test/recon_cossim\": test_recon_cossim / (test_i + 1),\n", " \"train/recon_mse\": recon_mse / (train_i + 1),\n", " \"test/recon_mse\": test_recon_mse / (test_i + 1),\n", " \"train/loss_prior\": loss_prior_total / (train_i + 1),\n", " \"test/loss_prior\": test_loss_prior_total / (test_i + 1),\n", " \"train/loss_RR\": loss_RR_total / (train_i + 1),\n", " \"test/loss_RR\": test_loss_RR_total / (test_i + 1),\n", " \"train/class_precisions_1\": class_precisions_1 / (train_i + 1),\n", " \"test/class_precisions_1\": test_class_precisions_1 / (test_i + 1),\n", " \"train/class_precisions_5\": class_precisions_5 / (train_i + 1),\n", " \"test/class_precisions_5\": test_class_precisions_5 / (test_i + 1),\n", " \"train/class_precisions_10\": class_precisions_10 / (train_i + 1),\n", " \"test/class_precisions_10\": test_class_precisions_10 / (test_i + 1),\n", " }\n", "\n", " # if finished training, save jpg recons if they exist\n", " if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):\n", " if blurry_recon: \n", " image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215\n", " # transform blurry recon latents to images and plot it\n", " fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n", " jj=-1\n", " for j in [0,1,2,3]:\n", " jj+=1\n", " axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", " axes[jj].axis('off')\n", " jj+=1\n", " axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", " axes[jj].axis('off')\n", "\n", " if wandb_log:\n", " logs[f\"test/blur_recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\")\n", " plt.close()\n", " else:\n", " plt.show()\n", "\n", " progress_bar.set_postfix(**logs)\n", "\n", " if wandb_log: wandb.log(logs)\n", " \n", " # Save model checkpoint and reconstruct\n", " if (ckpt_saving) and (epoch % ckpt_interval == 0):\n", " save_ckpt(f'last')\n", "\n", " # wait for other GPUs to catch up if needed\n", " accelerator.wait_for_everyone()\n", " torch.cuda.empty_cache()\n", "\n", "print(\"\\n===Finished!===\\n\")\n", "if ckpt_saving:\n", " save_ckpt(f'last')" ] }, { "cell_type": "code", "execution_count": 27, "id": "4211dd56-6a12-493e-8ccf-9b5229195cbd", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'photo of a person': 0,\n", " 'photo of airplane': 1,\n", " 'photo of baseball': 2,\n", " 'photo of bathroom': 3,\n", " 'photo of bear': 4,\n", " 'photo of bedroom': 5,\n", " 'photo of bike': 6,\n", " 'photo of bird': 7,\n", " 'photo of boat': 8,\n", " 'photo of bus': 9,\n", " 'photo of cat': 10,\n", " 'photo of clocktower': 11,\n", " 'photo of computer': 12,\n", " 'photo of cow': 13,\n", " 'photo of dog': 14,\n", " 'photo of elephant': 15,\n", " 'photo of flower': 16,\n", " 'photo of food': 17,\n", " 'photo of fruits': 18,\n", " 'photo of giraffe': 19,\n", " 'photo of group of people': 20,\n", " 'photo of horse': 21,\n", " 'photo of hydrant': 22,\n", " 'photo of living room': 23,\n", " 'photo of person eating': 24,\n", " 'photo of pizza': 25,\n", " 'photo of sheep': 26,\n", " 'photo of skate': 27,\n", " 'photo of ski': 28,\n", " 'photo of sky': 29,\n", " 'photo of soccer': 30,\n", " 'photo of sports': 31,\n", " 'photo of stop sign': 32,\n", " 'photo of surfer': 33,\n", " 'photo of sweets': 34,\n", " 'photo of tennis': 35,\n", " 'photo of toy': 36,\n", " 'photo of train': 37,\n", " 'photo of umbrella': 38,\n", " 'photo of vehicle': 39,\n", " 'photo of zebra': 40}" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "semantic_cluster_dict" ] }, { "cell_type": "code", "execution_count": null, "id": "a7e81ae3-171f-40ad-a3e8-24bee4472325", "metadata": { "tags": [] }, "outputs": [], "source": [ "plt.plot(losses)\n", "plt.show()\n", "plt.plot(test_losses)\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "id": "45d5bc17-0914-43a5-bb4e-d98f0ba238f0", "metadata": {}, "outputs": [], "source": [ "import wandb\n", "wandb.login()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "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.11.9" }, "toc": { "base_numbering": 1, "nav_menu": {}, "number_sections": true, "sideBar": true, "skip_h1_title": false, "title_cell": "Table of Contents", "title_sidebar": "Contents", "toc_cell": false, "toc_position": { "height": "calc(100% - 180px)", "left": "10px", "top": "150px", "width": "165px" }, "toc_section_display": true, "toc_window_display": true }, "toc-autonumbering": true, "vscode": { "interpreter": { "hash": "62aae01ef0cf7b6af841ab1c8ce59175c4332e693ab3d00bc32ceffb78a35376" } } }, "nbformat": 4, "nbformat_minor": 5 }