diff --git "a/src/Train_MLPMixer.ipynb" "b/src/Train_MLPMixer.ipynb" new file mode 100644--- /dev/null +++ "b/src/Train_MLPMixer.ipynb" @@ -0,0 +1,2476 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "f4d95fac-ac1d-473c-ab96-650f76e6aaf5", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# # Code to convert this notebook to .py if you want to run it via command line or with Slurm\n", + "# from subprocess import call\n", + "# command = \"jupyter nbconvert Train_with_autoencoder_MLPMixer.ipynb --to python\"\n", + "# call(command,shell=True)" + ] + }, + { + "cell_type": "markdown", + "id": "b0f0f4f3", + "metadata": {}, + "source": [ + "# Import packages & functions" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5bad764b-45c1-45ce-a716-8d055e09821a", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[2023-10-28 20:46:20,021] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "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", + "\n", + "import webdataset as wds\n", + "import gc\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import torch\n", + "import torch.nn as nn\n", + "from torchvision import transforms\n", + "\n", + "from accelerate import Accelerator, DeepSpeedPlugin\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": 3, + "id": "c0267850-3785-4be6-b134-b2a52bf55113", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "utils.mixco" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "cc5d2e32-6027-4a19-bef4-5ca068db35bb", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LOCAL RANK 0\n", + "Setting batch_size to 128\n", + "[2023-10-28 20:46:28,070] [WARNING] [comm.py:152:init_deepspeed_backend] NCCL backend in DeepSpeed not yet implemented\n", + "[2023-10-28 20:46:28,071] [INFO] [comm.py:594:init_distributed] cdb=None\n", + "[2023-10-28 20:46:28,071] [INFO] [comm.py:625:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl\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", + "num_devices = torch.cuda.device_count()\n", + "if num_devices==0: num_devices = 1\n", + "\n", + "# ## UNCOMMENT BELOW SECTION AND COMMENT OUT DEEPSPEED SECTION TO AVOID USING DEEPSPEED ###\n", + "# accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n", + "# global_batch_size = batch_size = 128\n", + "# data_type = torch.float16 # change depending on your mixed_precision\n", + "\n", + "### DEEPSPEED INITIALIZATION ###\n", + "if num_devices <= 1 and utils.is_interactive():\n", + " global_batch_size = batch_size = 128\n", + " print(f\"Setting batch_size to {batch_size}\")\n", + " # can emulate a distributed environment for deepspeed to work in jupyter notebook\n", + " os.environ[\"MASTER_ADDR\"] = \"localhost\"\n", + " os.environ[\"MASTER_PORT\"] = str(np.random.randint(10000)+9000)\n", + " os.environ[\"RANK\"] = \"0\"\n", + " os.environ[\"LOCAL_RANK\"] = \"0\"\n", + " os.environ[\"WORLD_SIZE\"] = \"1\"\n", + " os.environ[\"GLOBAL_BATCH_SIZE\"] = str(global_batch_size) # set this to your batch size!\n", + "else:\n", + " global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"] \n", + " batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices\n", + "\n", + "# alter the deepspeed config according to your global and local batch size\n", + "if local_rank == 0:\n", + " with open('deepspeed_config_stage2.json', 'r') as file:\n", + " config = json.load(file)\n", + " config['train_batch_size'] = int(os.environ[\"GLOBAL_BATCH_SIZE\"])\n", + " config['train_micro_batch_size_per_gpu'] = batch_size\n", + " config['bf16'] = {'enabled': False}\n", + " config['fp16'] = {'enabled': True}\n", + " with open('deepspeed_config_stage2.json', 'w') as file:\n", + " json.dump(config, file)\n", + "else:\n", + " # give some time for the local_rank=0 gpu to prep new deepspeed config file\n", + " time.sleep(10)\n", + "deepspeed_plugin = DeepSpeedPlugin(\"deepspeed_config_stage2.json\")\n", + "accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "PID of this process = 1896724\n", + "device: cuda:0\n", + "Distributed environment: DEEPSPEED Backend: nccl\n", + "Num processes: 1\n", + "Process index: 0\n", + "Local process index: 0\n", + "Device: cuda:0\n", + "\n", + "Mixed precision type: fp16\n", + "ds_config: {'bf16': {'enabled': False}, 'fp16': {'enabled': True}, 'zero_optimization': {'stage': 2, 'contiguous_gradients': True, 'stage3_gather_16bit_weights_on_model_save': True, 'stage3_max_live_parameters': 1000000000.0, 'stage3_max_reuse_distance': 1000000000.0, 'stage3_prefetch_bucket_size': 10000000.0, 'stage3_param_persistence_threshold': 100000.0, 'reduce_bucket_size': 10000000.0, 'sub_group_size': 1000000000.0, 'offload_optimizer': {'device': 'none', 'nvme_path': '/scratch', 'pin_memory': True}, 'offload_param': {'device': 'none', 'nvme_path': '/scratch', 'buffer_size': 4000000000.0, 'pin_memory': True}}, 'aio': {'block_size': 26214400, 'queue_depth': 64, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}, 'gradient_accumulation_steps': 1, 'gradient_clipping': 1.0, 'steps_per_print': inf, 'train_batch_size': 128, 'train_micro_batch_size_per_gpu': 128, 'wall_clock_breakdown': False, 'zero_allow_untested_optimizer': True}\n", + "\n", + "distributed = True 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", + "num_workers = num_devices\n", + "print(accelerator.state)\n", + "world_size = accelerator.state.num_processes\n", + "distributed = not accelerator.state.distributed_type == 'NO'\n", + "\n", + "# set data_type to match your mixed precision (automatically set based on deepspeed config)\n", + "if accelerator.mixed_precision == \"bf16\":\n", + " data_type = torch.bfloat16\n", + "elif accelerator.mixed_precision == \"fp16\":\n", + " data_type = torch.float16\n", + "else:\n", + " data_type = torch.float32\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": 6, + "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "model_name: 0qiAxQoaKN_interactive_bsl\n", + "['--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset', '--model_name=0qiAxQoaKN_interactive_bsl', '--subj=1', '--batch_size=128', '--no-blurry_recon', '--no-depth_recon', '--hidden_dim=4096', '--clip_scale=1.', '--blur_scale=100.', '--depth_scale=100.', '--max_lr=3e-4', '--mixup_pct=.66', '--num_epochs=12', '--ckpt_interval=999', '--no-use_image_aug', '--no-ckpt_saving']\n" + ] + } + ], + "source": [ + "# if running this interactively, can specify jupyter_args here for argparser to use\n", + "if utils.is_interactive():\n", + " # create random model_name\n", + " model_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))\n", + " model_name = model_name + \"_interactive_bsl\"\n", + " print(\"model_name:\", model_name)\n", + "\n", + " # global_batch_size and batch_size should already be defined in the above cells\n", + " # other variables can be specified in the following string:\n", + " jupyter_args = f\"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \\\n", + " --model_name={model_name} \\\n", + " --subj=1 --batch_size={batch_size} --no-blurry_recon --no-depth_recon --hidden_dim=4096 \\\n", + " --clip_scale=1. --blur_scale=100. --depth_scale=100. \\\n", + " --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug --no-ckpt_saving\"\n", + "\n", + " jupyter_args = jupyter_args.split()\n", + " print(jupyter_args)\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": 7, + "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n", + "parser.add_argument(\n", + " \"--model_name\", type=str, default=\"testing\",\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=\"/fsx/proj-fmri/shared/natural-scenes-dataset\",\n", + " help=\"Path to where NSD data is stored / where to download it to\",\n", + ")\n", + "parser.add_argument(\n", + " \"--subj\",type=int, default=1, choices=[1,2,5,7],\n", + ")\n", + "parser.add_argument(\n", + " \"--batch_size\", type=int, default=32,\n", + " help=\"Batch size can be increased by 10x if only training v2c and not diffusion diffuser\",\n", + ")\n", + "parser.add_argument(\n", + " \"--wandb_log\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to log to wandb\",\n", + ")\n", + "parser.add_argument(\n", + " \"--resume_from_ckpt\",action=argparse.BooleanOptionalAction,default=False,\n", + " help=\"if not using wandb and want to resume from a ckpt\",\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", + " \"--depth_recon\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to output depth reconstructions\",\n", + ")\n", + "parser.add_argument(\n", + " \"--blur_scale\",type=float,default=100.,\n", + " help=\"multiply loss from blurry recons by this number\",\n", + ")\n", + "parser.add_argument(\n", + " \"--depth_scale\",type=float,default=100.,\n", + " help=\"multiply loss from depth 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", + " \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=True,\n", + " help=\"whether to use image augmentation\",\n", + ")\n", + "parser.add_argument(\n", + " \"--num_epochs\",type=int,default=120,\n", + " help=\"number of epochs of training\",\n", + ")\n", + "parser.add_argument(\n", + " \"--hidden_dim\",type=int,default=4096,\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-4,\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)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "60cd7f2c-37fd-426b-a0c6-633e51bc4c4d", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "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", + "if use_image_aug:\n", + " import kornia\n", + " from kornia.augmentation.container import AugmentationSequential\n", + " img_augment = AugmentationSequential(\n", + " kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),\n", + " kornia.augmentation.Resize((224, 224)),\n", + " kornia.augmentation.RandomHorizontalFlip(p=0.3),\n", + " kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),\n", + " kornia.augmentation.RandomGrayscale(p=0.3),\n", + " same_on_batch=False,\n", + " data_keys=[\"input\"],\n", + " )" + ] + }, + { + "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": [ + "## Dataloader" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "81084834-035f-4465-ad59-59e6b806a2f5", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/{0..36}.tar\n", + "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/test/0.tar\n" + ] + } + ], + "source": [ + "if subj==1:\n", + " num_train = 24958\n", + " num_test = 2770\n", + "test_batch_size = num_test\n", + "\n", + "def my_split_by_node(urls): return urls\n", + " \n", + "train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..36}.tar\"\n", + "# train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..1}.tar\"\n", + "print(train_url)\n", + "\n", + "train_data = wds.WebDataset(train_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", + "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)\n", + "\n", + "test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n", + "print(test_url)\n", + "\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=test_batch_size, shuffle=False, drop_last=True, pin_memory=True)" + ] + }, + { + "cell_type": "markdown", + "id": "203b060a-2dd2-4c35-929b-c576be82eb52", + "metadata": {}, + "source": [ + "### check dataloaders are working" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "e7a9c68c-c3c9-4080-bd99-067c4486dc37", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0 2770 2770\n", + "---\n", + "\n", + "194 24960 24960\n" + ] + } + ], + "source": [ + "test_vox_indices = []\n", + "test_73k_images = []\n", + "for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):\n", + " test_vox_indices = np.append(test_vox_indices, behav[:,0,5].cpu().numpy())\n", + " test_73k_images = np.append(test_73k_images, behav[:,0,0].cpu().numpy())\n", + "test_vox_indices = test_vox_indices.astype(np.int16)\n", + "print(test_i, (test_i+1) * test_batch_size, len(test_vox_indices))\n", + "print(\"---\\n\")\n", + "\n", + "train_vox_indices = []\n", + "train_73k_images = []\n", + "for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n", + " train_vox_indices = np.append(train_vox_indices, behav[:,0,5].long().cpu().numpy())\n", + " train_73k_images = np.append(train_73k_images, behav[:,0,0].cpu().numpy())\n", + "train_vox_indices = train_vox_indices.astype(np.int16)\n", + "print(train_i, (train_i+1) * batch_size, len(train_vox_indices))" + ] + }, + { + "cell_type": "markdown", + "id": "45fad12c-f9fb-4408-8fd4-9bca324ad634", + "metadata": {}, + "source": [ + "## Load data and images" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "039dd330-7339-4f88-8f00-45f95e47baa0", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "subj01 betas loaded into memory\n", + "voxels torch.Size([27750, 15724])\n", + "images torch.Size([73000, 3, 224, 224])\n" + ] + } + ], + "source": [ + "# load betas\n", + "f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')\n", + "# f = h5py.File(f'{data_path}/betas_subj0{subj}_thresholded_wholebrain.hdf5', 'r')\n", + "\n", + "voxels = f['betas'][:]\n", + "print(f\"subj0{subj} betas loaded into memory\")\n", + "voxels = torch.Tensor(voxels).to(\"cpu\").to(data_type)\n", + "print(\"voxels\", voxels.shape)\n", + "num_voxels = voxels.shape[-1]\n", + "\n", + "# load orig images\n", + "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n", + "images = f['images'][:]\n", + "images = torch.Tensor(images).to(\"cpu\").to(data_type)\n", + "print(\"images\", 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": 12, + "id": "b0420dc0-199e-4c1a-857d-b1747058b467", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ViT-L/14 cuda:0\n" + ] + } + ], + "source": [ + "from models import Clipper\n", + "clip_model = Clipper(\"ViT-L/14\", device=torch.device(f\"cuda:{local_rank}\"), hidden_state=True, norm_embs=True)\n", + "clip_seq_dim = 257\n", + "clip_emb_dim = 768 #1024\n", + "# hidden_dim = 4096\n", + "seq_len = 1 #2 #32 " + ] + }, + { + "cell_type": "markdown", + "id": "5b79bd38-6990-4504-8d45-4a68d57d8885", + "metadata": {}, + "source": [ + "### SD VAE" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "01baff79-8114-482b-b115-6f05aa8ad691", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "# if blurry_recon:\n", + "# from diffusers import AutoencoderKL\n", + "# autoenc = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16, cache_dir=\"/fsx/proj-fmri/shared/cache\")\n", + "# # autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')[\"model_state_dict\"])\n", + "# autoenc.eval()\n", + "# autoenc.requires_grad_(False)\n", + "# autoenc.to(device)\n", + "# utils.count_params(autoenc)\n", + "\n", + "if blurry_recon:# or depth_recon:\n", + " from diffusers import VQModel\n", + " autoenc = VQModel.from_pretrained(\"/fsx/proj-fmri/shared/cache/models--microsoft--vq-diffusion-ithq/snapshots/3f796fb49ee559370dc638dea1d8116af131d993/vqvae\", torch_dtype=data_type)\n", + " autoenc.eval()\n", + " autoenc.requires_grad_(False)\n", + " autoenc.to(device)\n", + " utils.count_params(autoenc)" + ] + }, + { + "cell_type": "markdown", + "id": "120c8eee-9834-437d-bb60-b38faef50138", + "metadata": {}, + "source": [ + "#### downsampled images" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "6d1ba8dd-64c2-4ac9-947e-725b7f2e3e50", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "if blurry_recon:\n", + " if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))\n", + "\n", + " input_batch = images[[30]].to(device)\n", + " print(input_batch.shape)\n", + "\n", + " downsampled_image = nn.functional.interpolate(input_batch, size=(8, 8), mode='bilinear', align_corners=False)\n", + " re_upsampled_image = nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest')\n", + " re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215\n", + " print(re_upsampled_enc.shape)\n", + " \n", + " if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(re_upsampled_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))" + ] + }, + { + "cell_type": "markdown", + "id": "6390a3a8-2bef-4e81-9b82-e154d26a1e1d", + "metadata": {}, + "source": [ + "#### MiDaS depth" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "f35573e2-95bf-463d-8937-68ad4c2c3c20", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "if depth_recon:\n", + " from controlnet_aux.midas import MidasDetector\n", + " \n", + " midas_depth = MidasDetector.from_pretrained(\n", + " \"valhalla/t2iadapter-aux-models\", filename=\"dpt_large_384.pt\", model_type=\"dpt_large\", cache_dir=\"/fsx/proj-fmri/shared/cache\").to(device)\n", + " midas_depth.model.eval()\n", + " midas_depth.model.requires_grad_(False)\n", + " midas_depth.model.to(device)\n", + " pass" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "ba3f9207-b98e-45da-baa6-5cfcfb2ae958", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "if depth_recon:\n", + " if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))\n", + "\n", + " input_batch = images[[30,31]].float().to(device)\n", + " print(input_batch.shape)\n", + " \n", + " midas_emb = midas_depth.model(input_batch).unsqueeze(1)\n", + " print(midas_emb.shape)\n", + "\n", + " prediction = utils.resize(midas_emb, 32) #/30).clamp(0,1).half() # 30 is roughly prediction.max()\n", + " print(prediction.shape)\n", + " \n", + " prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()\n", + " midas_emb_size = prediction.flatten(1).shape[1]\n", + " print(\"midas_emb\", prediction.shape, prediction.min(), prediction.max())\n", + " print(\"midas_emb_size\", midas_emb_size)\n", + " \n", + " if utils.is_interactive(): display(utils.torch_to_Image(utils.resize(prediction, 224))) \n", + "\n", + " if blurry_recon:\n", + " prediction = utils.resize(midas_emb, 128).half().repeat(1,3,1,1)\n", + " prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()\n", + " prediction_enc = autoenc.encode(2*prediction-1).latents * 0.18215\n", + " print(\"vae midas_emb\", prediction_enc.shape, prediction_enc.min(), prediction_enc.max())\n", + " \n", + " if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(prediction_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))" + ] + }, + { + "cell_type": "markdown", + "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0", + "metadata": {}, + "source": [ + "### MindEye modules" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "MindEyeModule()" + ] + }, + "execution_count": 17, + "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": 18, + "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "param counts:\n", + "64,409,600 total\n", + "64,409,600 trainable\n", + "param counts:\n", + "64,409,600 total\n", + "64,409,600 trainable\n", + "torch.Size([2, 1, 15724]) torch.Size([2, 1, 4096])\n" + ] + } + ], + "source": [ + "class RidgeRegression(torch.nn.Module):\n", + " # make sure to add weight_decay when initializing optimizer\n", + " def __init__(self, input_size, out_features): \n", + " super(RidgeRegression, self).__init__()\n", + " self.out_features = out_features\n", + " self.linear = torch.nn.Linear(input_size, out_features)\n", + " def forward(self, x):\n", + " return self.linear(x)\n", + " \n", + "model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)\n", + "utils.count_params(model.ridge)\n", + "utils.count_params(model)\n", + "\n", + "b = torch.randn((2,1,voxels.shape[1]))\n", + "print(b.shape, model.ridge(b).shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "3602c333-d029-465c-8fb4-c3ccffdba6fd", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "param counts:\n", + "950,287,384 total\n", + "950,287,384 trainable\n", + "param counts:\n", + "1,014,696,984 total\n", + "1,014,696,984 trainable\n", + "b.shape torch.Size([2, 1, 4096])\n", + "torch.Size([2, 257, 768]) torch.Size([1]) torch.Size([1])\n" + ] + } + ], + "source": [ + "from functools import partial\n", + "from diffusers.models.vae import Decoder\n", + "class BrainNetwork(nn.Module):\n", + " def __init__(self, out_dim=768, in_dim=15724, seq_len=2, h=4096, n_blocks=4, drop=.15, clip_size=768):\n", + " super().__init__()\n", + " self.seq_len = seq_len\n", + " self.h = h\n", + " self.clip_size = clip_size\n", + " \n", + " # Initial linear layer to match the input dimensions to hidden dimensions\n", + " # self.lin0 = nn.Linear(in_dim, seq_len * h)\n", + " \n", + " # Mixer Blocks\n", + " self.mixer_blocks1 = nn.ModuleList([\n", + " self.mixer_block1(h, drop) for _ in range(n_blocks)\n", + " ])\n", + " self.mixer_blocks2 = nn.ModuleList([\n", + " self.mixer_block2(seq_len, drop) for _ in range(n_blocks)\n", + " ])\n", + " \n", + " # Output linear layer\n", + " self.clin1 = nn.Linear(h * seq_len, out_dim, bias=True)\n", + "\n", + " # low-rank matrices\n", + " # self.rank = 500\n", + " # self.U = nn.Parameter(torch.randn(self.rank, out_dim))\n", + " # self.V = nn.Parameter(torch.randn(h * seq_len, self.rank))\n", + " # self.S = nn.Parameter(torch.randn(out_dim))\n", + "\n", + " self.clip_proj = nn.Sequential(\n", + " nn.LayerNorm(clip_size),\n", + " nn.GELU(),\n", + " nn.Linear(clip_size, 2048),\n", + " nn.LayerNorm(2048),\n", + " nn.GELU(),\n", + " nn.Linear(2048, 2048),\n", + " nn.LayerNorm(2048),\n", + " nn.GELU(),\n", + " nn.Linear(2048, clip_size)\n", + " )\n", + "\n", + " if blurry_recon:\n", + " # self.blin1 = nn.Sequential(\n", + " # nn.Linear(out_dim, 4096, bias=True),\n", + " # nn.LayerNorm(4096),\n", + " # nn.GELU(),\n", + " # nn.Linear(4096, 4096))\n", + " self.blin1 = nn.Linear(h*seq_len, 4096)\n", + " self.bgroupnorm = nn.GroupNorm(1, 256)\n", + " self.bupsampler = Decoder(\n", + " in_channels=256,\n", + " out_channels=128,\n", + " up_block_types=[\"UpDecoderBlock2D\",\"UpDecoderBlock2D\",\"UpDecoderBlock2D\"],\n", + " block_out_channels=[32, 64, 128],\n", + " layers_per_block=1,\n", + " )\n", + "\n", + " if depth_recon:\n", + " # self.dlin1 = nn.Sequential(\n", + " # nn.Linear(h, midas_emb_size),\n", + " # nn.Sigmoid(),\n", + " # )\n", + " self.dlin1 = nn.Linear(h*seq_len, 4096)\n", + " self.dgroupnorm = nn.GroupNorm(1, 256)\n", + " self.dupsampler = Decoder(\n", + " in_channels=256,\n", + " out_channels=1,#128,\n", + " up_block_types=[\"UpDecoderBlock2D\",\"UpDecoderBlock2D\",\"UpDecoderBlock2D\",\"UpDecoderBlock2D\"],\n", + " block_out_channels=[32, 64, 128, 256],\n", + " layers_per_block=1,\n", + " )\n", + " \n", + " def mixer_block1(self, h, drop):\n", + " return nn.Sequential(\n", + " nn.LayerNorm(h),\n", + " self.mlp(h, h, drop), # Token mixing\n", + " )\n", + "\n", + " def mixer_block2(self, seq_len, drop):\n", + " return nn.Sequential(\n", + " nn.LayerNorm(seq_len),\n", + " self.mlp(seq_len, seq_len, drop) # Channel mixing\n", + " )\n", + " \n", + " def mlp(self, in_dim, out_dim, drop):\n", + " return nn.Sequential(\n", + " nn.Linear(in_dim, out_dim),\n", + " nn.GELU(),\n", + " nn.Dropout(drop),\n", + " nn.Linear(out_dim, out_dim),\n", + " )\n", + " \n", + " def forward(self, x):\n", + " # make empty tensors for blur and depth outputs\n", + " b,d = torch.Tensor([0.]), torch.Tensor([0.])\n", + " \n", + " # Initial linear layer\n", + " # x = self.lin0(x)\n", + " \n", + " # Reshape to seq_len by dim\n", + " # x = x.reshape(-1, self.seq_len, self.h)\n", + " \n", + " # Mixer blocks\n", + " residual1 = x\n", + " residual2 = x.permute(0,2,1)\n", + " for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):\n", + " x = block1(x) + residual1\n", + " residual1 = x\n", + " x = x.permute(0,2,1)\n", + " \n", + " x = block2(x) + residual2\n", + " residual2 = x\n", + " x = x.permute(0,2,1)\n", + " \n", + " # Flatten\n", + " x = x.reshape(x.size(0), -1)\n", + " \n", + " c = self.clin1(x)\n", + "\n", + " # low rank linear to out dim cuts # params by nearly half compared to full linear mapping\n", + " # c = (x @ (self.V/100) @ (self.U/100)) + self.S\n", + " \n", + " c = self.clip_proj(c.reshape(len(c), -1, self.clip_size))\n", + "\n", + " if blurry_recon:\n", + " b = self.blin1(x)\n", + " b = b.reshape(len(b), 256, 4, 4)\n", + " b = self.bgroupnorm(b)\n", + " b = self.bupsampler(b)\n", + " \n", + " if depth_recon:\n", + " d = self.dlin1(x)#.reshape(len(x), 1, 32, 32)\n", + " d = d.reshape(len(d), 256, 4, 4)\n", + " d = self.dgroupnorm(d)\n", + " d = self.dupsampler(d)\n", + " \n", + " return c, b, d\n", + "\n", + "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=seq_len, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) \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,seq_len,hidden_dim))\n", + "print(\"b.shape\",b.shape)\n", + "with torch.no_grad():\n", + " clip_, blur_, depth_ = model.backbone(b)\n", + "print(clip_.shape, blur_.shape, depth_.shape)" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "total_steps 2339\n", + "\n", + "Done with model preparations!\n", + "param counts:\n", + "1,014,696,984 total\n", + "1,014,696,984 trainable\n" + ] + } + ], + "source": [ + "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\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", + "]\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_train/num_devices/batch_size))),\n", + " last_epoch=-1\n", + " )\n", + "elif lr_scheduler_type == 'cycle':\n", + " total_steps=int(np.floor(num_epochs*(num_train/num_devices/batch_size)))\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", + " print(f'saving {ckpt_path}',flush=True)\n", + " unwrapped_model = accelerator.unwrap_model(model)\n", + " try:\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", + " except:\n", + " print(\"Couldn't save... moving on to prevent crashing.\")\n", + " del unwrapped_model\n", + " \n", + "print(\"\\nDone with model preparations!\")\n", + "utils.count_params(model)" + ] + }, + { + "cell_type": "markdown", + "id": "983f458b-35b8-49f2-b6db-80296cece730", + "metadata": {}, + "source": [ + "# Weights and Biases" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "0a25a662-daa8-4de9-9233-8364800fcb6b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "wandb mindeyev2 run 0qiAxQoaKN_interactive_bsl\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" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "wandb_config:\n", + " {'model_name': '0qiAxQoaKN_interactive_bsl', 'global_batch_size': 128, 'batch_size': 128, 'num_epochs': 12, 'clip_scale': 1.0, 'blur_scale': 100.0, 'use_image_aug': False, 'max_lr': 0.0003, 'mixup_pct': 0.66, 'num_train': 24958, 'num_test': 2770, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': True, 'num_devices': 1, 'world_size': 1, 'train_url': '/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/{0..36}.tar', 'test_url': '/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/test/0.tar'}\n", + "wandb_id: 0qiAxQoaKN_interactive_bsl\n" + ] + }, + { + "data": { + "text/html": [ + "wandb version 0.15.12 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.15.5" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Run data is saved locally in /fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20231028_204841-0qiAxQoaKN_interactive_bsl" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Syncing run 0qiAxQoaKN_interactive_bsl to Weights & Biases (docs)
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View project at https://stability.wandb.io/ckadirt/mindeyev2" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + " View run at https://stability.wandb.io/ckadirt/mindeyev2/runs/0qiAxQoaKN_interactive_bsl" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "wandb_log = True\n", + "if local_rank==0 and wandb_log: # only use main process for wandb logging\n", + " import wandb\n", + " wandb_project = 'mindeyev2'\n", + " wandb_run = model_name\n", + " wandb_notes = ''\n", + " \n", + " print(f\"wandb {wandb_project} run {wandb_run}\")\n", + " wandb.login(host='https://stability.wandb.io')#, relogin=True)\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", + " \"clip_scale\": clip_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_train\": num_train,\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", + " if True: # wandb_auto_resume\n", + " print(\"wandb_id:\",model_name)\n", + " wandb.init(\n", + " id = model_name,\n", + " project=wandb_project,\n", + " name=wandb_run,\n", + " config=wandb_config,\n", + " notes=wandb_notes,\n", + " resume=\"allow\",\n", + " )\n", + " else:\n", + " wandb.init(\n", + " project=wandb_project,\n", + " name=wandb_run,\n", + " config=wandb_config,\n", + " notes=wandb_notes,\n", + " )\n", + "else:\n", + " wandb_log = False" + ] + }, + { + "cell_type": "markdown", + "id": "d5690151-2131-4918-b750-e869cbd1a8a8", + "metadata": {}, + "source": [ + "# Main" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "12de6387-6e18-4e4b-b5ce-a847d625330a", + "metadata": {}, + "outputs": [], + "source": [ + "epoch = 0\n", + "losses, test_losses, lrs = [], [], []\n", + "best_test_loss = 1e9\n", + "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n", + "\n", + "# Optionally resume from checkpoint #\n", + "if resume_from_ckpt:\n", + " print(\"\\n---resuming from last.pth ckpt---\\n\")\n", + " try:\n", + " checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n", + " except:\n", + " print('last.pth failed... trying last_backup.pth')\n", + " checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n", + " epoch = checkpoint['epoch']\n", + " print(\"Epoch\",epoch)\n", + " optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n", + " lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n", + " model.load_state_dict(checkpoint['model_state_dict'])\n", + " del checkpoint\n", + "elif wandb_log:\n", + " if wandb.run.resumed:\n", + " print(\"\\n---resuming from last.pth ckpt---\\n\")\n", + " try:\n", + " checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n", + " except:\n", + " print('last.pth failed... trying last_backup.pth')\n", + " checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n", + " epoch = checkpoint['epoch']\n", + " print(\"Epoch\",epoch)\n", + " optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n", + " lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n", + " model.load_state_dict(checkpoint['model_state_dict'])\n", + " del checkpoint\n", + "torch.cuda.empty_cache()" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[2023-10-28 20:48:51,902] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.9.5, git-hash=unknown, git-branch=unknown\n", + "[2023-10-28 20:48:53,263] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False\n", + "[2023-10-28 20:48:53,265] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the client Optimizer\n", + "[2023-10-28 20:48:53,266] [INFO] [logging.py:96:log_dist] [Rank 0] Using client Optimizer as basic optimizer\n", + "[2023-10-28 20:48:53,267] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = AdamW\n", + "[2023-10-28 20:48:53,268] [INFO] [utils.py:54:is_zero_supported_optimizer] Checking ZeRO support for optimizer=AdamW type=\n", + "[2023-10-28 20:48:53,269] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.float16 ZeRO stage 2 optimizer\n", + "[2023-10-28 20:48:53,269] [INFO] [stage_1_and_2.py:133:__init__] Reduce bucket size 10000000\n", + "[2023-10-28 20:48:53,270] [INFO] [stage_1_and_2.py:134:__init__] Allgather bucket size 500,000,000\n", + "[2023-10-28 20:48:53,271] [INFO] [stage_1_and_2.py:135:__init__] CPU Offload: False\n", + "[2023-10-28 20:48:53,271] [INFO] [stage_1_and_2.py:136:__init__] Round robin gradient partitioning: False\n", + "Rank: 0 partition count [1, 1, 1] and sizes[(64409600, False), (950031116, False), (256268, False)] \n", + "[2023-10-28 20:48:55,761] [INFO] [utils.py:785:see_memory_usage] Before initializing optimizer states\n", + "[2023-10-28 20:48:55,763] [INFO] [utils.py:786:see_memory_usage] MA 7.68 GB Max_MA 7.68 GB CA 7.71 GB Max_CA 8 GB \n", + "[2023-10-28 20:48:55,764] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory: used = 72.0 GB, percent = 6.4%\n", + "[2023-10-28 20:48:55,940] [INFO] [utils.py:785:see_memory_usage] After initializing optimizer states\n", + "[2023-10-28 20:48:55,941] [INFO] [utils.py:786:see_memory_usage] MA 15.24 GB Max_MA 26.1 GB CA 26.62 GB Max_CA 27 GB \n", + "[2023-10-28 20:48:55,942] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory: used = 72.0 GB, percent = 6.4%\n", + "[2023-10-28 20:48:55,943] [INFO] [stage_1_and_2.py:488:__init__] optimizer state initialized\n", + "[2023-10-28 20:48:56,073] [INFO] [utils.py:785:see_memory_usage] After initializing ZeRO optimizer\n", + "[2023-10-28 20:48:56,074] [INFO] [utils.py:786:see_memory_usage] MA 15.24 GB Max_MA 15.24 GB CA 26.62 GB Max_CA 27 GB \n", + "[2023-10-28 20:48:56,075] [INFO] [utils.py:793:see_memory_usage] CPU Virtual Memory: used = 71.99 GB, percent = 6.4%\n", + "[2023-10-28 20:48:56,078] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = AdamW\n", + "[2023-10-28 20:48:56,078] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using client LR scheduler\n", + "[2023-10-28 20:48:56,079] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = None\n", + "[2023-10-28 20:48:56,080] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[1.200000000000002e-05, 1.200000000000002e-05, 1.200000000000002e-05], mom=[(0.95, 0.999), (0.95, 0.999), (0.95, 0.999)]\n", + "[2023-10-28 20:48:56,081] [INFO] [config.py:960:print] DeepSpeedEngine configuration:\n", + "[2023-10-28 20:48:56,082] [INFO] [config.py:964:print] activation_checkpointing_config {\n", + " \"partition_activations\": false, \n", + " \"contiguous_memory_optimization\": false, \n", + " \"cpu_checkpointing\": false, \n", + " \"number_checkpoints\": null, \n", + " \"synchronize_checkpoint_boundary\": false, \n", + " \"profile\": false\n", + "}\n", + "[2023-10-28 20:48:56,082] [INFO] [config.py:964:print] aio_config ................... {'block_size': 26214400, 'queue_depth': 64, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}\n", + "[2023-10-28 20:48:56,083] [INFO] [config.py:964:print] amp_enabled .................. False\n", + "[2023-10-28 20:48:56,084] [INFO] [config.py:964:print] amp_params ................... False\n", + "[2023-10-28 20:48:56,085] [INFO] [config.py:964:print] autotuning_config ............ {\n", + " \"enabled\": false, \n", + " \"start_step\": null, \n", + " \"end_step\": null, \n", + " \"metric_path\": null, \n", + " \"arg_mappings\": null, \n", + " \"metric\": \"throughput\", \n", + " \"model_info\": null, \n", + " \"results_dir\": \"autotuning_results\", \n", + " \"exps_dir\": \"autotuning_exps\", \n", + " \"overwrite\": true, \n", + " \"fast\": true, \n", + " \"start_profile_step\": 3, \n", + " \"end_profile_step\": 5, \n", + " \"tuner_type\": \"gridsearch\", \n", + " \"tuner_early_stopping\": 5, \n", + " \"tuner_num_trials\": 50, \n", + " \"model_info_path\": null, \n", + " \"mp_size\": 1, \n", + " \"max_train_batch_size\": null, \n", + " \"min_train_batch_size\": 1, \n", + " \"max_train_micro_batch_size_per_gpu\": 1.024000e+03, \n", + " \"min_train_micro_batch_size_per_gpu\": 1, \n", + " \"num_tuning_micro_batch_sizes\": 3\n", + "}\n", + "[2023-10-28 20:48:56,085] [INFO] [config.py:964:print] bfloat16_enabled ............. False\n", + "[2023-10-28 20:48:56,086] [INFO] [config.py:964:print] checkpoint_parallel_write_pipeline False\n", + "[2023-10-28 20:48:56,087] [INFO] [config.py:964:print] checkpoint_tag_validation_enabled True\n", + "[2023-10-28 20:48:56,087] [INFO] [config.py:964:print] checkpoint_tag_validation_fail False\n", + "[2023-10-28 20:48:56,088] [INFO] [config.py:964:print] comms_config ................. \n", + "[2023-10-28 20:48:56,088] [INFO] [config.py:964:print] communication_data_type ...... None\n", + "[2023-10-28 20:48:56,089] [INFO] [config.py:964:print] compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}\n", + "[2023-10-28 20:48:56,090] [INFO] [config.py:964:print] curriculum_enabled_legacy .... False\n", + "[2023-10-28 20:48:56,091] [INFO] [config.py:964:print] curriculum_params_legacy ..... False\n", + "[2023-10-28 20:48:56,091] [INFO] [config.py:964:print] data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}\n", + "[2023-10-28 20:48:56,092] [INFO] [config.py:964:print] data_efficiency_enabled ...... False\n", + "[2023-10-28 20:48:56,093] [INFO] [config.py:964:print] dataloader_drop_last ......... False\n", + "[2023-10-28 20:48:56,093] [INFO] [config.py:964:print] disable_allgather ............ False\n", + "[2023-10-28 20:48:56,094] [INFO] [config.py:964:print] dump_state ................... False\n", + "[2023-10-28 20:48:56,095] [INFO] [config.py:964:print] dynamic_loss_scale_args ...... None\n", + "[2023-10-28 20:48:56,095] [INFO] [config.py:964:print] eigenvalue_enabled ........... False\n", + "[2023-10-28 20:48:56,096] [INFO] [config.py:964:print] eigenvalue_gas_boundary_resolution 1\n", + "[2023-10-28 20:48:56,097] [INFO] [config.py:964:print] eigenvalue_layer_name ........ bert.encoder.layer\n", + "[2023-10-28 20:48:56,097] [INFO] [config.py:964:print] eigenvalue_layer_num ......... 0\n", + "[2023-10-28 20:48:56,098] [INFO] [config.py:964:print] eigenvalue_max_iter .......... 100\n", + "[2023-10-28 20:48:56,099] [INFO] [config.py:964:print] eigenvalue_stability ......... 1e-06\n", + "[2023-10-28 20:48:56,099] [INFO] [config.py:964:print] eigenvalue_tol ............... 0.01\n", + "[2023-10-28 20:48:56,100] [INFO] [config.py:964:print] eigenvalue_verbose ........... False\n", + "[2023-10-28 20:48:56,100] [INFO] [config.py:964:print] elasticity_enabled ........... False\n", + "[2023-10-28 20:48:56,101] [INFO] [config.py:964:print] flops_profiler_config ........ {\n", + " \"enabled\": false, \n", + " \"recompute_fwd_factor\": 0.0, \n", + " \"profile_step\": 1, \n", + " \"module_depth\": -1, \n", + " \"top_modules\": 1, \n", + " \"detailed\": true, \n", + " \"output_file\": null\n", + "}\n", + "[2023-10-28 20:48:56,102] [INFO] [config.py:964:print] fp16_auto_cast ............... False\n", + "[2023-10-28 20:48:56,103] [INFO] [config.py:964:print] fp16_enabled ................. True\n", + "[2023-10-28 20:48:56,103] [INFO] [config.py:964:print] fp16_master_weights_and_gradients False\n", + "[2023-10-28 20:48:56,104] [INFO] [config.py:964:print] global_rank .................. 0\n", + "[2023-10-28 20:48:56,105] [INFO] [config.py:964:print] grad_accum_dtype ............. None\n", + "[2023-10-28 20:48:56,105] [INFO] [config.py:964:print] gradient_accumulation_steps .. 1\n", + "[2023-10-28 20:48:56,106] [INFO] [config.py:964:print] gradient_clipping ............ 1.0\n", + "[2023-10-28 20:48:56,107] [INFO] [config.py:964:print] gradient_predivide_factor .... 1.0\n", + "[2023-10-28 20:48:56,107] [INFO] [config.py:964:print] hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8\n", + "[2023-10-28 20:48:56,108] [INFO] [config.py:964:print] initial_dynamic_scale ........ 65536\n", + "[2023-10-28 20:48:56,109] [INFO] [config.py:964:print] load_universal_checkpoint .... False\n", + "[2023-10-28 20:48:56,109] [INFO] [config.py:964:print] loss_scale ................... 0\n", + "[2023-10-28 20:48:56,110] [INFO] [config.py:964:print] memory_breakdown ............. False\n", + "[2023-10-28 20:48:56,111] [INFO] [config.py:964:print] mics_hierarchial_params_gather False\n", + "[2023-10-28 20:48:56,111] [INFO] [config.py:964:print] mics_shard_size .............. -1\n", + "[2023-10-28 20:48:56,112] [INFO] [config.py:964:print] monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') enabled=False\n", + "[2023-10-28 20:48:56,113] [INFO] [config.py:964:print] nebula_config ................ {\n", + " \"enabled\": false, \n", + " \"persistent_storage_path\": null, \n", + " \"persistent_time_interval\": 100, \n", + " \"num_of_version_in_retention\": 2, \n", + " \"enable_nebula_load\": true, \n", + " \"load_path\": null\n", + "}\n", + "[2023-10-28 20:48:56,113] [INFO] [config.py:964:print] optimizer_legacy_fusion ...... False\n", + "[2023-10-28 20:48:56,114] [INFO] [config.py:964:print] optimizer_name ............... None\n", + "[2023-10-28 20:48:56,115] [INFO] [config.py:964:print] optimizer_params ............. None\n", + "[2023-10-28 20:48:56,115] [INFO] [config.py:964:print] pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}\n", + "[2023-10-28 20:48:56,116] [INFO] [config.py:964:print] pld_enabled .................. False\n", + "[2023-10-28 20:48:56,117] [INFO] [config.py:964:print] pld_params ................... False\n", + "[2023-10-28 20:48:56,117] [INFO] [config.py:964:print] prescale_gradients ........... False\n", + "[2023-10-28 20:48:56,118] [INFO] [config.py:964:print] scheduler_name ............... None\n", + "[2023-10-28 20:48:56,119] [INFO] [config.py:964:print] scheduler_params ............. None\n", + "[2023-10-28 20:48:56,119] [INFO] [config.py:964:print] sparse_attention ............. None\n", + "[2023-10-28 20:48:56,120] [INFO] [config.py:964:print] sparse_gradients_enabled ..... False\n", + "[2023-10-28 20:48:56,121] [INFO] [config.py:964:print] steps_per_print .............. inf\n", + "[2023-10-28 20:48:56,121] [INFO] [config.py:964:print] train_batch_size ............. 128\n", + "[2023-10-28 20:48:56,122] [INFO] [config.py:964:print] train_micro_batch_size_per_gpu 128\n", + "[2023-10-28 20:48:56,123] [INFO] [config.py:964:print] use_node_local_storage ....... False\n", + "[2023-10-28 20:48:56,123] [INFO] [config.py:964:print] wall_clock_breakdown ......... False\n", + "[2023-10-28 20:48:56,124] [INFO] [config.py:964:print] world_size ................... 1\n", + "[2023-10-28 20:48:56,125] [INFO] [config.py:964:print] zero_allow_untested_optimizer True\n", + "[2023-10-28 20:48:56,125] [INFO] [config.py:964:print] zero_config .................. stage=2 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=10000000 allgather_partitions=True allgather_bucket_size=500,000,000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=DeepSpeedZeroOffloadParamConfig(device='none', nvme_path=PosixPath('/scratch'), buffer_count=5, buffer_size=4000000000, max_in_cpu=1,000,000,000, pin_memory=True) offload_optimizer=DeepSpeedZeroOffloadOptimizerConfig(device='none', nvme_path=PosixPath('/scratch'), buffer_count=4, pin_memory=True, pipeline=False, pipeline_read=False, pipeline_write=False, fast_init=False) sub_group_size=1000000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=10000000 param_persistence_threshold=100000 model_persistence_threshold=sys.maxsize max_live_parameters=1000000000 max_reuse_distance=1000000000 gather_16bit_weights_on_model_save=True stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True\n", + "[2023-10-28 20:48:56,126] [INFO] [config.py:964:print] zero_enabled ................. True\n", + "[2023-10-28 20:48:56,127] [INFO] [config.py:964:print] zero_force_ds_cpu_optimizer .. True\n", + "[2023-10-28 20:48:56,127] [INFO] [config.py:964:print] zero_optimization_stage ...... 2\n", + "[2023-10-28 20:48:56,128] [INFO] [config.py:950:print_user_config] json = {\n", + " \"bf16\": {\n", + " \"enabled\": false\n", + " }, \n", + " \"fp16\": {\n", + " \"enabled\": true\n", + " }, \n", + " \"zero_optimization\": {\n", + " \"stage\": 2, \n", + " \"contiguous_gradients\": true, \n", + " \"stage3_gather_16bit_weights_on_model_save\": true, \n", + " \"stage3_max_live_parameters\": 1.000000e+09, \n", + " \"stage3_max_reuse_distance\": 1.000000e+09, \n", + " \"stage3_prefetch_bucket_size\": 1.000000e+07, \n", + " \"stage3_param_persistence_threshold\": 1.000000e+05, \n", + " \"reduce_bucket_size\": 1.000000e+07, \n", + " \"sub_group_size\": 1.000000e+09, \n", + " \"offload_optimizer\": {\n", + " \"device\": \"none\", \n", + " \"nvme_path\": \"/scratch\", \n", + " \"pin_memory\": true\n", + " }, \n", + " \"offload_param\": {\n", + " \"device\": \"none\", \n", + " \"nvme_path\": \"/scratch\", \n", + " \"buffer_size\": 4.000000e+09, \n", + " \"pin_memory\": true\n", + " }\n", + " }, \n", + " \"aio\": {\n", + " \"block_size\": 2.621440e+07, \n", + " \"queue_depth\": 64, \n", + " \"thread_count\": 1, \n", + " \"single_submit\": false, \n", + " \"overlap_events\": true\n", + " }, \n", + " \"gradient_accumulation_steps\": 1, \n", + " \"gradient_clipping\": 1.0, \n", + " \"steps_per_print\": inf, \n", + " \"train_batch_size\": 128, \n", + " \"train_micro_batch_size_per_gpu\": 128, \n", + " \"wall_clock_breakdown\": false, \n", + " \"zero_allow_untested_optimizer\": true\n", + "}\n" + ] + } + ], + "source": [ + "model, optimizer, train_dl, lr_scheduler = accelerator.prepare(\n", + "model, optimizer, train_dl, lr_scheduler\n", + ")\n", + "# leaving out test_dl since we will only have local_rank 0 device do evals" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "469e6313-425f-45ed-875a-ecd5df343e31", + "metadata": {}, + "outputs": [], + "source": [ + "def add_saturation(image, alpha=2):\n", + " gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]\n", + " gray_image = gray_image.unsqueeze(1).expand_as(image)\n", + " saturated_image = alpha * image + (1 - alpha) * gray_image\n", + " return torch.clamp(saturated_image, 0, 1)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "60be0d5f-3e94-4612-9373-61b53d836393", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0qiAxQoaKN_interactive_bsl starting with epoch 0 / 12\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + " 0%| | 0/12 [00:00╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮\n", + "│ in <module>:119 │\n", + "│ │\n", + "│ 116 │ │ │ │ │ # utils.check_loss(pixcorr) │\n", + "│ 117 │ │ │ │\n", + "│ 118 │ │ │ utils.check_loss(loss) │\n", + "│ ❱ 119 │ │ │ accelerator.backward(loss) │\n", + "│ 120 │ │ │ optimizer.step() │\n", + "│ 121 │ │ │ │\n", + "│ 122 │ │ │ losses.append(loss.item()) │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/accelerate/accelerator. │\n", + "│ py:1815 in backward │\n", + "│ │\n", + "│ 1812 │ │ │ # deepspeed handles loss scaling by gradient_accumulation_steps in its `back │\n", + "│ 1813 │ │ │ loss = loss / self.gradient_accumulation_steps │\n", + "│ 1814 │ │ if self.distributed_type == DistributedType.DEEPSPEED: │\n", + "│ ❱ 1815 │ │ │ self.deepspeed_engine_wrapped.backward(loss, **kwargs) │\n", + "│ 1816 │ │ elif self.distributed_type == DistributedType.MEGATRON_LM: │\n", + "│ 1817 │ │ │ return │\n", + "│ 1818 │ │ elif self.scaler is not None: │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/accelerate/utils/deepsp │\n", + "│ eed.py:176 in backward │\n", + "│ │\n", + "│ 173 │ │ # - zero grad │\n", + "│ 174 │ │ # - checking overflow │\n", + "│ 175 │ │ # - lr_scheduler step (only if engine.lr_scheduler is not None) │\n", + "│ ❱ 176 │ │ self.engine.step() │\n", + "│ 177 │ │ # and this plugin overrides the above calls with no-ops when Accelerate runs und │\n", + "│ 178 │ │ # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabli │\n", + "│ 179 │ │ # training loop that works transparently under many training regimes. │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/engin │\n", + "│ e.py:2053 in step │\n", + "│ │\n", + "│ 2050 │ │ │ │ │ and self.quantizer.any_precision_switch()): │\n", + "│ 2051 │ │ │ │ self._take_model_step(lr_kwargs, self.block_eigenvalue) │\n", + "│ 2052 │ │ │ else: │\n", + "│ ❱ 2053 │ │ │ │ self._take_model_step(lr_kwargs) │\n", + "│ 2054 │ │ │ │\n", + "│ 2055 │ │ │ report_progress = self.global_rank == 0 if self.global_rank else True │\n", + "│ 2056 │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/engin │\n", + "│ e.py:1960 in _take_model_step │\n", + "│ │\n", + "│ 1957 │ │ │ │ # https://nvidia.github.io/apex/advanced.html#gradient-clipping │\n", + "│ 1958 │ │ │ │ master_params = amp.master_params(self.optimizer) │\n", + "│ 1959 │ │ │ │ clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clippin │\n", + "│ ❱ 1960 │ │ self.optimizer.step() │\n", + "│ 1961 │ │ │\n", + "│ 1962 │ │ if hasattr(self.optimizer, '_global_grad_norm'): │\n", + "│ 1963 │ │ │ self._global_grad_norm = self.optimizer._global_grad_norm │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/zero/ │\n", + "│ stage_1_and_2.py:1733 in step │\n", + "│ │\n", + "│ 1730 │ │ │ │ │\n", + "│ 1731 │ │ │ │ # Step 3:- run the optimizer if no offloading │\n", + "│ 1732 │ │ │ │ self.start_timers([OPTIMIZER_STEP]) │\n", + "│ ❱ 1733 │ │ │ │ self._optimizer_step(i) │\n", + "│ 1734 │ │ │ │ # Step 4:- get rid of the fp32 gradients. Not needed anymore │\n", + "│ 1735 │ │ │ │ self.single_partition_of_fp32_groups[i].grad = None │\n", + "│ 1736 │ │ │ │ del single_grad_partition │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/zero/ │\n", + "│ stage_1_and_2.py:1638 in _optimizer_step │\n", + "│ │\n", + "│ 1635 │ │ # self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no) │\n", + "│ 1636 │ │ #else: │\n", + "│ 1637 │ │ # self.optimizer.step() │\n", + "│ ❱ 1638 │ │ self.optimizer.step() │\n", + "│ 1639 │ │ self.optimizer.param_groups = original_param_groups │\n", + "│ 1640 │ │\n", + "│ 1641 │ def step(self, closure=None): │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/lr_schedule │\n", + "│ r.py:69 in wrapper │\n", + "│ │\n", + "│ 66 │ │ │ │ instance = instance_ref() │\n", + "│ 67 │ │ │ │ instance._step_count += 1 │\n", + "│ 68 │ │ │ │ wrapped = func.__get__(instance, cls) │\n", + "│ ❱ 69 │ │ │ │ return wrapped(*args, **kwargs) │\n", + "│ 70 │ │ �� │\n", + "│ 71 │ │ │ # Note that the returned function here is no longer a bound method, │\n", + "│ 72 │ │ │ # so attributes like `__func__` and `__self__` no longer exist. │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/optimizer.p │\n", + "│ y:280 in wrapper │\n", + "│ │\n", + "│ 277 │ │ │ │ │ │ │ raise RuntimeError(f\"{func} must return None or a tuple of ( │\n", + "│ 278 │ │ │ │ │ │ │ │ │ │ │ f\"but got {result}.\") │\n", + "│ 279 │ │ │ │ │\n", + "│ ❱ 280 │ │ │ │ out = func(*args, **kwargs) │\n", + "│ 281 │ │ │ │ self._optimizer_step_code() │\n", + "│ 282 │ │ │ │ │\n", + "│ 283 │ │ │ │ # call optimizer step post hooks │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/optimizer.p │\n", + "│ y:33 in _use_grad │\n", + "│ │\n", + "│ 30 │ │ prev_grad = torch.is_grad_enabled() │\n", + "│ 31 │ │ try: │\n", + "│ 32 │ │ │ torch.set_grad_enabled(self.defaults['differentiable']) │\n", + "│ ❱ 33 │ │ │ ret = func(self, *args, **kwargs) │\n", + "│ 34 │ │ finally: │\n", + "│ 35 │ │ │ torch.set_grad_enabled(prev_grad) │\n", + "│ 36 │ │ return ret │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/adamw.py:17 │\n", + "│ 1 in step │\n", + "│ │\n", + "│ 168 │ │ │ │ state_steps, │\n", + "│ 169 │ │ │ ) │\n", + "│ 170 │ │ │ │\n", + "│ ❱ 171 │ │ │ adamw( │\n", + "│ 172 │ │ │ │ params_with_grad, │\n", + "│ 173 │ │ │ │ grads, │\n", + "│ 174 │ │ │ │ exp_avgs, │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/adamw.py:32 │\n", + "│ 1 in adamw │\n", + "│ │\n", + "│ 318 │ else: │\n", + "│ 319 │ │ func = _single_tensor_adamw │\n", + "│ 320 │ │\n", + "│ ❱ 321 │ func( │\n", + "│ 322 │ │ params, │\n", + "│ 323 │ │ grads, │\n", + "│ 324 │ │ exp_avgs, │\n", + "│ │\n", + "│ /admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/adamw.py:56 │\n", + "│ 6 in _multi_tensor_adamw │\n", + "│ │\n", + "│ 563 │ │ │ else: │\n", + "│ 564 │ │ │ │ exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) │\n", + "│ 565 │ │ │ │ torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt) │\n", + "│ ❱ 566 │ │ │ │ denom = torch._foreach_add(exp_avg_sq_sqrt, eps) │\n", + "│ 567 │ │ │ │\n", + "│ 568 │ │ │ torch._foreach_addcdiv_(device_params, device_exp_avgs, denom, step_size) │\n", + "│ 569 │\n", + "╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n", + "OutOfMemoryError: CUDA out of memory. Tried to allocate 3.54 GiB (GPU 0; 39.56 GiB total capacity; 26.52 GiB \n", + "already allocated; 2.00 GiB free; 35.94 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory\n", + "try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and \n", + "PYTORCH_CUDA_ALLOC_CONF\n", + "\n" + ], + "text/plain": [ + "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n", + "\u001b[31m│\u001b[0m in \u001b[92m\u001b[0m:\u001b[94m119\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m116 \u001b[0m\u001b[2m│ │ │ │ │ \u001b[0m\u001b[2m# utils.check_loss(pixcorr)\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m117 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m118 \u001b[0m\u001b[2m│ │ │ \u001b[0mutils.check_loss(loss) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m119 \u001b[2m│ │ │ \u001b[0maccelerator.backward(loss) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m120 \u001b[0m\u001b[2m│ │ │ \u001b[0moptimizer.step() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m121 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m122 \u001b[0m\u001b[2m│ │ │ \u001b[0mlosses.append(loss.item()) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/accelerate/\u001b[0m\u001b[1;33maccelerator.\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33mpy\u001b[0m:\u001b[94m1815\u001b[0m in \u001b[92mbackward\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1812 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[2m# deepspeed handles loss scaling by gradient_accumulation_steps in its `back\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1813 \u001b[0m\u001b[2m│ │ │ \u001b[0mloss = loss / \u001b[96mself\u001b[0m.gradient_accumulation_steps \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1814 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m.distributed_type == DistributedType.DEEPSPEED: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1815 \u001b[2m│ │ │ \u001b[0m\u001b[96mself\u001b[0m.deepspeed_engine_wrapped.backward(loss, **kwargs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1816 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94melif\u001b[0m \u001b[96mself\u001b[0m.distributed_type == DistributedType.MEGATRON_LM: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1817 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[94mreturn\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1818 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94melif\u001b[0m \u001b[96mself\u001b[0m.scaler \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/accelerate/utils/\u001b[0m\u001b[1;33mdeepsp\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33meed.py\u001b[0m:\u001b[94m176\u001b[0m in \u001b[92mbackward\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m173 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# - zero grad\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m174 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# - checking overflow\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m175 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# - lr_scheduler step (only if engine.lr_scheduler is not None)\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m176 \u001b[2m│ │ \u001b[0m\u001b[96mself\u001b[0m.engine.step() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m177 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# and this plugin overrides the above calls with no-ops when Accelerate runs und\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m178 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabli\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m179 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# training loop that works transparently under many training regimes.\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/\u001b[0m\u001b[1;33mengin\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33me.py\u001b[0m:\u001b[94m2053\u001b[0m in \u001b[92mstep\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2050 \u001b[0m\u001b[2m│ │ │ │ │ \u001b[0m\u001b[95mand\u001b[0m \u001b[96mself\u001b[0m.quantizer.any_precision_switch()): \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2051 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m._take_model_step(lr_kwargs, \u001b[96mself\u001b[0m.block_eigenvalue) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2052 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[94melse\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m2053 \u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m._take_model_step(lr_kwargs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2054 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2055 \u001b[0m\u001b[2m│ │ │ \u001b[0mreport_progress = \u001b[96mself\u001b[0m.global_rank == \u001b[94m0\u001b[0m \u001b[94mif\u001b[0m \u001b[96mself\u001b[0m.global_rank \u001b[94melse\u001b[0m \u001b[94mTrue\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m2056 \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/\u001b[0m\u001b[1;33mengin\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33me.py\u001b[0m:\u001b[94m1960\u001b[0m in \u001b[92m_take_model_step\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1957 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[2m# https://nvidia.github.io/apex/advanced.html#gradient-clipping\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1958 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mmaster_params = amp.master_params(\u001b[96mself\u001b[0m.optimizer) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1959 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mclip_grad_norm_(parameters=master_params, max_norm=\u001b[96mself\u001b[0m.gradient_clippin \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1960 \u001b[2m│ │ \u001b[0m\u001b[96mself\u001b[0m.optimizer.step() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1961 \u001b[0m\u001b[2m│ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1962 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94mif\u001b[0m \u001b[96mhasattr\u001b[0m(\u001b[96mself\u001b[0m.optimizer, \u001b[33m'\u001b[0m\u001b[33m_global_grad_norm\u001b[0m\u001b[33m'\u001b[0m): \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1963 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[96mself\u001b[0m._global_grad_norm = \u001b[96mself\u001b[0m.optimizer._global_grad_norm \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/zero/\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33mstage_1_and_2.py\u001b[0m:\u001b[94m1733\u001b[0m in \u001b[92mstep\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1730 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1731 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[2m# Step 3:- run the optimizer if no offloading\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1732 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m.start_timers([OPTIMIZER_STEP]) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1733 \u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m._optimizer_step(i) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1734 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[2m# Step 4:- get rid of the fp32 gradients. Not needed anymore\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1735 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m.single_partition_of_fp32_groups[i].grad = \u001b[94mNone\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1736 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[94mdel\u001b[0m single_grad_partition \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/deepspeed/runtime/zero/\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33mstage_1_and_2.py\u001b[0m:\u001b[94m1638\u001b[0m in \u001b[92m_optimizer_step\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1635 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1636 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m#else:\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1637 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[2m# self.optimizer.step()\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1638 \u001b[2m│ │ \u001b[0m\u001b[96mself\u001b[0m.optimizer.step() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1639 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[96mself\u001b[0m.optimizer.param_groups = original_param_groups \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1640 \u001b[0m\u001b[2m│ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m1641 \u001b[0m\u001b[2m│ \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mstep\u001b[0m(\u001b[96mself\u001b[0m, closure=\u001b[94mNone\u001b[0m): \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33mlr_schedule\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33mr.py\u001b[0m:\u001b[94m69\u001b[0m in \u001b[92mwrapper\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 66 \u001b[0m\u001b[2m│ │ │ │ \u001b[0minstance = instance_ref() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 67 \u001b[0m\u001b[2m│ │ │ │ \u001b[0minstance._step_count += \u001b[94m1\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 68 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mwrapped = func.\u001b[92m__get__\u001b[0m(instance, \u001b[96mcls\u001b[0m) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 69 \u001b[2m│ │ │ │ \u001b[0m\u001b[94mreturn\u001b[0m wrapped(*args, **kwargs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 70 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 71 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[2m# Note that the returned function here is no longer a bound method,\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 72 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[2m# so attributes like `__func__` and `__self__` no longer exist.\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33moptimizer.p\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33my\u001b[0m:\u001b[94m280\u001b[0m in \u001b[92mwrapper\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m277 \u001b[0m\u001b[2m│ │ │ │ │ │ │ \u001b[0m\u001b[94mraise\u001b[0m \u001b[96mRuntimeError\u001b[0m(\u001b[33mf\u001b[0m\u001b[33m\"\u001b[0m\u001b[33m{\u001b[0mfunc\u001b[33m}\u001b[0m\u001b[33m must return None or a tuple of (\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m278 \u001b[0m\u001b[2m│ │ │ │ │ │ │ │ │ │ │ \u001b[0m\u001b[33mf\u001b[0m\u001b[33m\"\u001b[0m\u001b[33mbut got \u001b[0m\u001b[33m{\u001b[0mresult\u001b[33m}\u001b[0m\u001b[33m.\u001b[0m\u001b[33m\"\u001b[0m) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m279 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m280 \u001b[2m│ │ │ │ \u001b[0mout = func(*args, **kwargs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m281 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[96mself\u001b[0m._optimizer_step_code() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m282 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m283 \u001b[0m\u001b[2m│ │ │ │ \u001b[0m\u001b[2m# call optimizer step post hooks\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33moptimizer.p\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[1;33my\u001b[0m:\u001b[94m33\u001b[0m in \u001b[92m_use_grad\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 30 \u001b[0m\u001b[2m│ │ \u001b[0mprev_grad = torch.is_grad_enabled() \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 31 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94mtry\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 32 \u001b[0m\u001b[2m│ │ │ \u001b[0mtorch.set_grad_enabled(\u001b[96mself\u001b[0m.defaults[\u001b[33m'\u001b[0m\u001b[33mdifferentiable\u001b[0m\u001b[33m'\u001b[0m]) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 33 \u001b[2m│ │ │ \u001b[0mret = func(\u001b[96mself\u001b[0m, *args, **kwargs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 34 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94mfinally\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 35 \u001b[0m\u001b[2m│ │ │ \u001b[0mtorch.set_grad_enabled(prev_grad) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m 36 \u001b[0m\u001b[2m│ │ \u001b[0m\u001b[94mreturn\u001b[0m ret \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33madamw.py\u001b[0m:\u001b[94m17\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[94m1\u001b[0m in \u001b[92mstep\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m168 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mstate_steps, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m169 \u001b[0m\u001b[2m│ │ │ \u001b[0m) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m170 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m171 \u001b[2m│ │ │ \u001b[0madamw( \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m172 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mparams_with_grad, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m173 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mgrads, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m174 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mexp_avgs, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33madamw.py\u001b[0m:\u001b[94m32\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[94m1\u001b[0m in \u001b[92madamw\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m318 \u001b[0m\u001b[2m│ \u001b[0m\u001b[94melse\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m319 \u001b[0m\u001b[2m│ │ \u001b[0mfunc = _single_tensor_adamw \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m320 \u001b[0m\u001b[2m│ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m321 \u001b[2m│ \u001b[0mfunc( \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m322 \u001b[0m\u001b[2m│ │ \u001b[0mparams, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m323 \u001b[0m\u001b[2m│ │ \u001b[0mgrads, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m324 \u001b[0m\u001b[2m│ │ \u001b[0mexp_avgs, \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2;33m/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/optim/\u001b[0m\u001b[1;33madamw.py\u001b[0m:\u001b[94m56\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[94m6\u001b[0m in \u001b[92m_multi_tensor_adamw\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m563 \u001b[0m\u001b[2m│ │ │ \u001b[0m\u001b[94melse\u001b[0m: \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m564 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mexp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m565 \u001b[0m\u001b[2m│ │ │ │ \u001b[0mtorch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m566 \u001b[2m│ │ │ │ \u001b[0mdenom = torch._foreach_add(exp_avg_sq_sqrt, eps) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m567 \u001b[0m\u001b[2m│ │ │ \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m568 \u001b[0m\u001b[2m│ │ │ \u001b[0mtorch._foreach_addcdiv_(device_params, device_exp_avgs, denom, step_size) \u001b[31m│\u001b[0m\n", + "\u001b[31m│\u001b[0m \u001b[2m569 \u001b[0m \u001b[31m│\u001b[0m\n", + "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n", + "\u001b[1;91mOutOfMemoryError: \u001b[0mCUDA out of memory. Tried to allocate \u001b[1;36m3.54\u001b[0m GiB \u001b[1m(\u001b[0mGPU \u001b[1;36m0\u001b[0m; \u001b[1;36m39.56\u001b[0m GiB total capacity; \u001b[1;36m26.52\u001b[0m GiB \n", + "already allocated; \u001b[1;36m2.00\u001b[0m GiB free; \u001b[1;36m35.94\u001b[0m GiB reserved in total by PyTorch\u001b[1m)\u001b[0m If reserved memory is >> allocated memory\n", + "try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and \n", + "PYTORCH_CUDA_ALLOC_CONF\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "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", + "\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", + " loss_clip_total = 0.\n", + " loss_blurry_total = 0.\n", + " loss_depth_total = 0.\n", + " test_loss_clip_total = 0.\n", + " test_loss_blurry_total = 0.\n", + " test_loss_depth_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", + " for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n", + " with torch.cuda.amp.autocast(dtype=data_type):\n", + " optimizer.zero_grad()\n", + " \n", + " voxel = voxels[behav[:,0,5].cpu().long()].to(device)\n", + " image = images[behav[:,0,0].cpu().long()].to(device).float()\n", + "\n", + " for past in range(1):\n", + " past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)\n", + " \n", + " if blurry_recon:\n", + " # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215\n", + " blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215\n", + "\n", + " if depth_recon:\n", + " # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)\n", + " depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)\n", + " depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()\n", + " depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215\n", + " \n", + " if use_image_aug: \n", + " image = img_augment(image)\n", + " \n", + " clip_target = clip_model.embed_image(image)\n", + " assert not torch.any(torch.isnan(clip_target))\n", + " \n", + " if epoch < int(mixup_pct * num_epochs):\n", + " voxel, perm, betas, select = utils.mixco(voxel)\n", + " past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)\n", + " \n", + " voxel_ridge = model.ridge(voxel).unsqueeze(1)\n", + " \n", + " # past_voxel_ridge = model.ridge(past_voxel)\n", + " # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1), past_voxel_ridge.unsqueeze(1)), axis=1)\n", + " \n", + " clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)\n", + " \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 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", + " downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)\n", + " re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))\n", + " re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215\n", + " \n", + " loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))\n", + " loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))\n", + " loss_blurry_total += loss_blurry.item()\n", + " loss_blurry *= blur_scale\n", + " loss += loss_blurry\n", + "\n", + " if depth_recon:\n", + " loss_depth = l1(depth_image_enc_, depth_image_enc)\n", + " # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))\n", + " loss_depth_total += loss_depth.item()\n", + " loss_depth *= depth_scale\n", + " loss += loss_depth\n", + " \n", + " # forward and backward top 1 accuracy \n", + " labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) \n", + " fwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm)), labels, k=1).item()\n", + " bwd_percent_correct += utils.topk(torch.abs(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(voxel)), size=batch_size//5, replace=False)\n", + " # random_samps = np.arange(batch_size//5)\n", + " blurry_recon_images = (autoenc.decode(blurry_image_enc_[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)\n", + " # pixcorr_origsize_nanmean is computationally less intense than utils.pixcorr and uses nanmean instead of mean\n", + " pixcorr = utils.pixcorr_origsize_nanmean(image[random_samps], blurry_recon_images)\n", + " # pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n", + " # loss += (1 - pixcorr)\n", + " blurry_pixcorr += pixcorr.item()\n", + " # utils.check_loss(pixcorr)\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[behav[:,0,5].cpu().long()]\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 test_image is None:\n", + " test_image = images[im][None]\n", + " test_voxel = torch.mean(voxel[locs],axis=0)[None]\n", + " else:\n", + " test_image = torch.vstack((test_image, images[im][None]))\n", + " test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))\n", + " \n", + " # random sample of 300\n", + " random_indices = torch.arange(len(test_voxel))[:300]\n", + " voxel = test_voxel[random_indices].to(device)\n", + " image = test_image[random_indices].to(device)\n", + " assert len(image) == 300\n", + "\n", + " if blurry_recon:\n", + " # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215\n", + " blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215\n", + "\n", + " if depth_recon:\n", + " # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)\n", + " depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)\n", + " depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()\n", + " depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215\n", + " \n", + " clip_target = clip_model.embed_image(image.float())\n", + " \n", + " voxel_ridge = model.ridge(voxel).unsqueeze(1)\n", + "\n", + " # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1),voxel_ridge.unsqueeze(1)),axis=1)\n", + " \n", + " clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)\n", + " \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", + " loss_clip = utils.soft_clip_loss(\n", + " clip_voxels_norm,\n", + " clip_target_norm,\n", + " temp=.006)\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", + " downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)\n", + " re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))\n", + " re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215\n", + " \n", + " loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))\n", + " loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))\n", + " test_loss_blurry_total += loss_blurry.item()\n", + " loss_blurry *= blur_scale\n", + " loss += loss_blurry\n", + " \n", + " # halving the batch size because the decoder is computationally heavy\n", + " blurry_recon_images = (autoenc.decode(blurry_image_enc_[:len(voxel)//2]/0.18215).sample / 2 + 0.5).clamp(0,1)\n", + " blurry_recon_images = torch.vstack((blurry_recon_images, (autoenc.decode(blurry_image_enc_[len(voxel)//2:]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", + " pixcorr = utils.pixcorr(image, blurry_recon_images)\n", + " loss += (1 - pixcorr)\n", + " test_blurry_pixcorr += pixcorr.item()\n", + "\n", + " if depth_recon:\n", + " loss_depth = l1(depth_image_enc_, depth_image_enc)\n", + " # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))\n", + " test_loss_depth_total += loss_depth.item()\n", + " loss_depth *= depth_scale\n", + " loss += loss_depth\n", + " \n", + " # forward and backward top 1 accuracy \n", + " labels = torch.arange(len(clip_target_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", + " utils.check_loss(loss) \n", + " test_losses.append(loss.item())\n", + "\n", + " # if utils.is_interactive(): clear_output(wait=True)\n", + " print(\"---\")\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", + " \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n", + " \"test/loss_blurry_total\": test_loss_blurry_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/loss_depth_total\": loss_depth_total / (train_i + 1),\n", + " \"test/loss_depth_total\": test_loss_depth_total / (test_i + 1),\n", + " }\n", + " \n", + " if blurry_recon: \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(blurry_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(blurry_image_enc_[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", + " axes[jj].axis('off')\n", + " \n", + " if wandb_log:\n", + " logs[f\"test/recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\")\n", + " plt.close()\n", + " else:\n", + " plt.show()\n", + "\n", + " if depth_recon:\n", + " # transform blurry recon latents to images and plot it\n", + " fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n", + " # axes[0].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", + " # axes[1].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc_[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n", + " jj=-1\n", + " for j in [0,1,2,3]:\n", + " jj+=1\n", + " axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc[[j]].view(1,1,32,32).clamp(0,1), 224)))\n", + " axes[jj].axis('off')\n", + " jj+=1\n", + " axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc_[[j]].view(1,1,32,32).clamp(0,1), 224)))\n", + " axes[jj].axis('off')\n", + " if wandb_log:\n", + " logs[f\"test/depth_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", + " # Save model checkpoint and reconstruct\n", + " if epoch % ckpt_interval == 0:\n", + " if not utils.is_interactive():\n", + " save_ckpt(f'last')\n", + " \n", + " if wandb_log: wandb.log(logs)\n", + "\n", + " # wait for other GPUs to catch up if needed\n", + " accelerator.wait_for_everyone()\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", + "\n", + "print(\"\\n===Finished!===\\n\")\n", + "if ckpt_saving:\n", + " save_ckpt(f'last')\n", + "if not utils.is_interactive():\n", + " sys.exit(0)" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "35cc1be7-bf76-4ad1-8c6a-de52bd013bf4", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Sat Oct 28 21:13:17 2023 \n", + "+-----------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 |\n", + "|-------------------------------+----------------------+----------------------+\n", + "| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n", + "| | | MIG M. |\n", + "|===============================+======================+======================|\n", + "| 0 NVIDIA A100-SXM... On | 00000000:10:1C.0 Off | 0 |\n", + "| N/A 33C P0 50W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 1 NVIDIA A100-SXM... On | 00000000:10:1D.0 Off | 0 |\n", + "| N/A 30C P0 53W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 2 NVIDIA A100-SXM... On | 00000000:20:1C.0 Off | 0 |\n", + "| N/A 34C P0 54W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 3 NVIDIA A100-SXM... On | 00000000:20:1D.0 Off | 0 |\n", + "| N/A 30C P0 53W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 4 NVIDIA A100-SXM... On | 00000000:90:1C.0 Off | 0 |\n", + "| N/A 36C P0 53W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 5 NVIDIA A100-SXM... On | 00000000:90:1D.0 Off | 0 |\n", + "| N/A 35C P0 72W / 400W | 38467MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 6 NVIDIA A100-SXM... On | 00000000:A0:1C.0 Off | 0 |\n", + "| N/A 33C P0 50W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + "| 7 NVIDIA A100-SXM... On | 00000000:A0:1D.0 Off | 0 |\n", + "| N/A 31C P0 51W / 400W | 3MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-------------------------------+----------------------+----------------------+\n", + " \n", + "+-----------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=============================================================================|\n", + "| 5 N/A N/A 1896724 C ...3/envs/mindeye/bin/python 38464MiB |\n", + "+-----------------------------------------------------------------------------+\n" + ] + } + ], + "source": [ + "!nvidia-smi" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "93e87fde-815d-4452-9915-f5f5dacf7c2a", + "metadata": {}, + "outputs": [], + "source": [ + "plt.plot(losses)\n", + "plt.show()\n", + "plt.plot(test_losses)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "f2690877-a431-44e8-a2ca-61f4b7397070", + "metadata": {}, + "source": [ + "# Retrieve nearest neighbor in the training set using test set data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b6b8feb-391d-437e-a5d9-a2088f1b1149", + "metadata": {}, + "outputs": [], + "source": [ + "annots = np.load(\"/fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_annots_curated.npy\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "612ac5aa-6f0f-45ad-809e-03df905d184c", + "metadata": {}, + "outputs": [], + "source": [ + "ii=2\n", + "all_indices = np.unique(train_73k_images) #np.hstack((test_vox_indices[ii],train_vox_indices))\n", + "with torch.no_grad(), torch.cuda.amp.autocast():\n", + " for batch in tqdm(range(0,len(all_indices),512)):\n", + " if batch==0:\n", + " clip_target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()\n", + " else:\n", + " target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()\n", + " clip_target = torch.vstack((clip_target,target))\n", + " clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n", + "\n", + " voxel = test_voxel[[ii]].to(device)\n", + " image = test_image[[ii]].to(device)\n", + "\n", + " print(\"Original Image (test set)\")\n", + " display(utils.torch_to_Image(image))\n", + " \n", + " clip_target = clip_model.embed_image(image).cpu()\n", + " # clip_target_norm = torch.vstack((clip_target_norm, nn.functional.normalize(clip_target.flatten(1), dim=-1)))\n", + " \n", + " voxel_ridge = model.ridge(voxel).unsqueeze(1)\n", + " clip_voxels, _, _ = model.backbone(voxel_ridge) \n", + " clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n", + " clip_voxels_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n", + "\n", + " print(\"clip_voxels_norm\", clip_voxels_norm.shape)\n", + " print(\"clip_target_norm\", clip_target_norm.shape)\n", + " \n", + " sortt = torch.argsort(utils.batchwise_cosine_similarity(clip_voxels_norm.cpu(), \n", + " clip_target_norm).flatten()).flip(0)\n", + " picks = all_indices[sortt[:5]]\n", + "\n", + " print(\"\\nNearest neighbors in training set\")\n", + " for ip,p in enumerate(picks):\n", + " display(utils.torch_to_Image(images[[p]]))\n", + " # print(utils.select_annotations([annots[int(p)]]))\n", + " if ip==0: predicted_caption = utils.select_annotations([annots[int(p)]])[0]\n", + "\n", + "print(\"\\n=====\\npredicted_caption:\\n\", predicted_caption)" + ] + }, + { + "cell_type": "markdown", + "id": "1473ddaa-5f2b-4448-9194-c7b0801d05db", + "metadata": {}, + "source": [ + "# Feed into Stable Diffusion XL for reconstructions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "70e50e0d-c44f-4d56-939a-2943535e1747", + "metadata": {}, + "outputs": [], + "source": [ + "from diffusers import StableDiffusionXLPipeline\n", + "pipe = StableDiffusionXLPipeline.from_pretrained(\n", + " \"/fsx/proj-fmri/shared/cache/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/f898a3e026e802f68796b95e9702464bac78d76f\", torch_dtype=torch.float16, variant=\"fp16\", use_safetensors=True\n", + ")\n", + "pipe.to(\"cuda\")\n", + "pass" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "479e6994-3eaa-47d2-89a3-422c464fab36", + "metadata": {}, + "outputs": [], + "source": [ + "prompt = predicted_caption\n", + "recon = pipe(prompt=prompt).images[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9dc48e1b-5842-4a29-963a-6469d943a72c", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "print(\"Seen image\")\n", + "display(utils.torch_to_Image(image))\n", + "\n", + "print(\"Reconstruction\")\n", + "utils.torch_to_Image(utils.resize(transforms.ToTensor()(recon),224))" + ] + } + ], + "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.10.8" + }, + "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 +}