stereoscope.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The StereoScope\n",
    "built for Hugging Face's AI or Not competition.<br>\n",
    "https://huggingface.co/spaces/competitions/aiornot <br>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2023, [Jozsef Szalma](https://www.linkedin.com/in/szalma/)<br>\n",
    "Creative Commons Attribution-NonCommercial 4.0 International Public License <br>\n",
    "https://creativecommons.org/licenses/by-nc/4.0/legalcode <br>\n",
    "If you want to use the below code commercially then do contact me for licensing or with other offers of collaboration ;)<br>\n",
    "Also keep in mind, the weights of pretrained ConvNeXt-V2 are also on CC BY-NC 4.0 as of writing."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HW requirements: I built this notebook on an nvidia GPU with 24GB memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#important! The below code is relying on timm models not available in timm 0.6.x, I used timm 0.8.12.dev0 installed from git\n",
    "%pip install git+https://github.com/rwightman/pytorch-image-models.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "from random import randint\n",
    "from tqdm import tqdm\n",
    "from tqdm.auto import tqdm\n",
    "import pandas as pd\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import timm\n",
    "import torchmetrics\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import autocast \n",
    "from torch.cuda.amp import GradScaler\n",
    "\n",
    "from timm.models.layers import trunc_normal_\n",
    "\n",
    "from huggingface_hub import login\n",
    "from datasets import load_dataset\n",
    "\n",
    "import mlflow"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters<br>\n",
    "credentials and working directories are set up in environment variables HF_KEY, WORKING_DIR_WIN and CHECKPOINT_DIR_WIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ConvNeXt\n",
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# All rights reserved.\n",
    "# This source code is licensed under the MIT license\n",
    "\n",
    "# ConvNeXt-V2\n",
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# All rights reserved.\n",
    "# This source code is licensed under the license found in the\n",
    "# LICENSE file in the root directory of this source tree (Attribution-NonCommercial 4.0 International (CC BY-NC 4.0))\n",
    "# No code was used directly from ConvNeXt-V2, however the weights are CC BY-NC 4.0 so beware if using commercially.\n",
    "\n",
    "\n",
    "args = {\n",
    "        'random_seed'           : 42,\n",
    "        'rgb_target_size'       : (384,384), \n",
    "        'fft_target_size'       : (384,384),  \n",
    "        'rgb_mean'              : (0.485, 0.456, 0.406), \n",
    "        'rgb_std'               : (0.229, 0.224, 0.225), \n",
    "        'rgb_model_name'        : 'convnextv2_huge.fcmae_ft_in22k_in1k_384',\n",
    "        'fft_model_name'        : 'convnext_xlarge.fb_in22k_ft_in1k_384',\n",
    "        'rgb_model_pretrained'  : True,\n",
    "        'fft_model_pretrained'  : True, \n",
    "        'rgb_model_frozen'      : False,\n",
    "        'fft_model_frozen'      : False, \n",
    "\n",
    "        'batch_size'            : 2, \n",
    "        'num_workers'           : 0, #multithreading is a tad problematic in Jupyter on Windows\n",
    "        'num_epochs'            : 100, \n",
    "        'weight_decay'          : 1e-3,\n",
    "        'resume_from_checkpoint': None, #'DualDetector_23a34a28b8024139bf238ce289204654_epoch=29_val_loss=1e-05.pth', \n",
    "        'validation_size'       : 0.1, \n",
    "\n",
    "        'model_default_lr'      : 4e-3 * 0.8,\n",
    "        'model_default_batch'   : 4096,\n",
    "        #LR is scaled from the model_default_lr with (batch_size / default_batch_size)\n",
    "\n",
    "        'gradient_clip'         : None, \n",
    "        'gradient_accum'        : 4        \n",
    "\n",
    "}\n",
    "\n",
    "hf_key = os.getenv(\"HF_KEY\")\n",
    "working_dir = os.getenv(\"WORKING_DIR_WIN\")\n",
    "checkpoint_dir = os.getenv(\"CHECKPOINT_DIR_WIN\")\n",
    "\n",
    "torch.manual_seed(args['random_seed'])\n",
    "np.random.seed(args['random_seed'])\n",
    "random.seed(args['random_seed'])\n",
    "np_generator = np.random.default_rng(args['random_seed'])\n",
    "\n",
    "torch_device = 'cuda'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-Stream architecture\n",
    "- the model is composed of two pretrained networks joined on the last fully connected layer\n",
    "- the RGB network is exposed to random 384x384 crops \n",
    "- the FFT network is exposed to 2D spectrograms of the same cropped RGB image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DualDetector (nn.Module):\n",
    "    def __init__(self, rgb_model_name='convnext_xlarge_in22k',fft_model_name='convnext_xlarge_in22k', rgb_model_pretrained=True, fft_model_pretrained=True,rgb_model_frozen=True, fft_model_frozen=True): \n",
    "        super(DualDetector,self).__init__()\n",
    "        self.model_rgb = timm.create_model(rgb_model_name, pretrained = rgb_model_pretrained)\n",
    "        self.model_fft = timm.create_model(fft_model_name, pretrained = fft_model_pretrained)\n",
    "        \n",
    "        if rgb_model_frozen:\n",
    "            for p in self.model_rgb.parameters():\n",
    "                p.requires_grad = False\n",
    "        if fft_model_frozen:\n",
    "            for p in self.model_fft.parameters():\n",
    "                p.requires_grad = False\n",
    "\n",
    "        embedding_size = self.model_rgb.head.fc.in_features + self.model_fft.head.fc.in_features\n",
    "        \n",
    "        self.model_rgb.head.fc = nn.Identity()\n",
    "        self.model_fft.head.fc = nn.Identity()\n",
    "\n",
    "        self.head_fc1 = nn.Linear(embedding_size, 1)\n",
    "\n",
    "        self.criterion = torch.nn.BCEWithLogitsLoss()\n",
    "        self.val_accuracy = torchmetrics.Accuracy('binary')\n",
    "\n",
    "\n",
    "    def forward(self, rgb, fft):\n",
    "        \n",
    "        rgb = self.model_rgb(rgb)\n",
    "        fft = self.model_fft(fft)\n",
    "        embeddings = torch.cat([rgb,fft],dim=1)\n",
    "  \n",
    "        x = self.head_fc1(embeddings)\n",
    "\n",
    "        return x      \n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "dualdetector = DualDetector(\n",
    "                            rgb_model_name=args['rgb_model_name'],\n",
    "                            fft_model_name=args['fft_model_name'],\n",
    "                            rgb_model_pretrained=args['rgb_model_pretrained'], \n",
    "                            fft_model_pretrained=args['fft_model_pretrained'],\n",
    "                            rgb_model_frozen=args['rgb_model_frozen'], \n",
    "                            fft_model_frozen=args['fft_model_frozen']\n",
    "                            )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### dataset from Hugging Face<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "login(token=hf_key,add_to_git_credential=True)\n",
    "\n",
    "os.environ['HF_HOME']=working_dir\n",
    "os.chdir(working_dir)\n",
    "\n",
    "ds = load_dataset('competitions/aiornot')\n",
    "\n",
    "split = ds[\"train\"].train_test_split(args['validation_size'],generator=np_generator)\n",
    "ds[\"train\"] = split[\"train\"]\n",
    "ds[\"validation\"] = split[\"test\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocessing\n",
    "I feed the network two images; a cropped RGB image and its 2D spectrogram.<br>\n",
    "I intentionally avoid downsampling images.<br>\n",
    "Normalizing is done with the mean and std specific to the pretrained network.<br>\n",
    "Data augmentation is a trade-off between disrupting the fake image signals (repeating patterns, color-shifts, changes in sharpness) and having a large enough training set.<br>\n",
    "At inference time the images are center cropped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_transform = transforms.Compose([                \n",
    "                                        transforms.RandomVerticalFlip(),\n",
    "                                        transforms.RandomHorizontalFlip(),\n",
    "                                        transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),\n",
    "                                        transforms.RandomGrayscale(),\n",
    "                                   \n",
    "                                        transforms.Normalize(args['rgb_mean'], args['rgb_std'])\n",
    "                                    ])\n",
    "\n",
    "\n",
    "rgb_transform_val = transforms.Compose([\n",
    "                               \n",
    "                                        transforms.Normalize(args['rgb_mean'], args['rgb_std'])\n",
    "                                    ])\n",
    "\n",
    "\n",
    "def simple_crop_tensor(target_size, tensor_img, how='random'):\n",
    "    #a simple cropping method of mine that A: does not need or do padding and B: can be either random or centered and C: won't do any rescaling\n",
    "    width, height = tensor_img.shape[2], tensor_img.shape[1]\n",
    "    target_width, target_height = target_size\n",
    "    width_range = width - target_width\n",
    "    height_range = height - target_height\n",
    "\n",
    "    if ((width_range > 0) or (height_range > 0)) :\n",
    "        if how == 'random':\n",
    "            x_movement = randint(0, width_range-1)\n",
    "            y_movement = randint(0, height_range-1)\n",
    "        elif how == 'center':\n",
    "            x_movement = int(width_range/2)-1\n",
    "            y_movement = int(height_range/2)-1\n",
    "        else:\n",
    "            raise NotImplemented\n",
    "            \n",
    "        bbox = (y_movement, x_movement, y_movement + target_height, x_movement + target_width)\n",
    "\n",
    "        cropped_tensor = tensor_img[:, bbox[0]:bbox[2], bbox[1]:bbox[3]]\n",
    "\n",
    "        return cropped_tensor\n",
    "    else:\n",
    "        return tensor_img\n",
    "\n",
    "\n",
    "def tensor_spectrogram(im):\n",
    "    #generates a 2D spectrogram with Fast Fourier transform (phase discarded) for each channel of an RGB image\n",
    "\n",
    "    for i in range(3):\n",
    "        img = im[i,:,:]\n",
    "        fft_img = torch.fft.fft2(img)\n",
    "        fft_img = torch.log(torch.abs(fft_img) + 1e-3)\n",
    "        \n",
    "        fft_img_np = fft_img.cpu().numpy()\n",
    "        fft_min = np.percentile(fft_img_np,5) + 1e-8 \n",
    "        fft_max = np.percentile(fft_img_np,95) + 1e-7\n",
    "        \n",
    "        fft_img = (fft_img - fft_min)/(fft_max - fft_min)\n",
    "        fft_img = (fft_img-0.5)*2\n",
    "\n",
    "        fft_img = torch.clamp(fft_img, min=-1, max=1) \n",
    "        im[i,:,:] = fft_img\n",
    "\n",
    "    return im          \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The preprocessing code is bit of a spagetti and could be certainly improved.<br>\n",
    "<br>\n",
    "I'm off-loading preprocessing to the GPU at the earliest opportunity, as I run into a CPU bottleneck here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train_transform_fn(example_batch):\n",
    "\n",
    "    tensor_lst_rgb = []\n",
    "    tensor_lst_fft = []\n",
    "    for im in example_batch['image']:\n",
    "        \n",
    "        im_tensor = transforms.functional.pil_to_tensor(im).to(torch_device).float() / 255.0\n",
    "        \n",
    "        im_tensor = rgb_transform(im_tensor)\n",
    "        #TODO move this into a transforms.Lambda\n",
    "        im_tensor_rgb = simple_crop_tensor(target_size=args['rgb_target_size'], tensor_img=im_tensor, how='random')\n",
    "        im_tensor_fft = tensor_spectrogram(im=im_tensor_rgb.detach().clone()) \n",
    "        \n",
    "        tensor_lst_fft.append(im_tensor_fft)\n",
    "        tensor_lst_rgb.append(im_tensor_rgb)\n",
    "    x_rgb = torch.stack(tensor_lst_rgb)\n",
    "    x_fft = torch.stack(tensor_lst_fft)\n",
    "\n",
    "    \n",
    "    y = example_batch['label']\n",
    "    return {'x_rgb': x_rgb,'x_fft': x_fft, 'y': y}\n",
    "    \n",
    "\n",
    "def val_transform_fn(example_batch):\n",
    "    \n",
    "    tensor_lst_rgb = []\n",
    "    tensor_lst_fft = []\n",
    "    for im in example_batch['image']:\n",
    "\n",
    "        im_tensor = transforms.functional.pil_to_tensor(im).to(torch_device).float() / 255.0\n",
    "        \n",
    "        im_tensor = rgb_transform(im_tensor)\n",
    "     \n",
    "        im_tensor_rgb = simple_crop_tensor(target_size=args['rgb_target_size'], tensor_img=im_tensor, how='center')\n",
    "        im_tensor_fft = tensor_spectrogram(im=im_tensor_rgb.detach().clone()) \n",
    "        \n",
    "        tensor_lst_fft.append(im_tensor_fft)\n",
    "        tensor_lst_rgb.append(im_tensor_rgb)\n",
    "    x_rgb = torch.stack(tensor_lst_rgb)\n",
    "    x_fft = torch.stack(tensor_lst_fft)\n",
    "\n",
    "    y = example_batch['label']\n",
    "    return {'x_rgb': x_rgb,'x_fft': x_fft, 'y': y}\n",
    "    \n",
    "\n",
    "\n",
    "def test_transform_fn(example_batch):\n",
    "    \n",
    "    tensor_lst_rgb = []\n",
    "    tensor_lst_fft = []\n",
    "    for im in example_batch['image']:\n",
    "\n",
    "        im_tensor = transforms.functional.pil_to_tensor(im).to(torch_device).float() / 255.0\n",
    "        \n",
    "        im_tensor = rgb_transform(im_tensor)\n",
    "      \n",
    "        im_tensor_rgb = simple_crop_tensor(target_size=args['rgb_target_size'], tensor_img=im_tensor, how='center')\n",
    "        im_tensor_fft = tensor_spectrogram(im=im_tensor_rgb.detach().clone()) \n",
    "        \n",
    "        tensor_lst_fft.append(im_tensor_fft)\n",
    "        tensor_lst_rgb.append(im_tensor_rgb)\n",
    "    x_rgb = torch.stack(tensor_lst_rgb)\n",
    "    x_fft = torch.stack(tensor_lst_fft)\n",
    "\n",
    "    return {'x_rgb': x_rgb,'x_fft': x_fft}\n",
    "    \n",
    "\n",
    "\n",
    "def collate_fn(batch):\n",
    "    x_rgb = torch.stack([ex['x_rgb'] for ex in batch])\n",
    "    x_fft = torch.stack([ex['x_fft'] for ex in batch])\n",
    "    y = torch.tensor([ex['y'] for ex in batch]).float().to(torch_device)\n",
    "    return x_rgb,x_fft, y\n",
    "    \n",
    "\n",
    "def collate_fn_test(batch):\n",
    "    x_rgb = torch.stack([ex['x_rgb'] for ex in batch])\n",
    "    x_fft = torch.stack([ex['x_fft'] for ex in batch])\n",
    "\n",
    "    return x_rgb,x_fft"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MLFlow \n",
    "I'm an avid practitioner of \"lossing\" and I can't always get Tensorboard to show me data during training and that makes me sad.<br>\n",
    "https://twitter.com/JozsefSzalma/status/1621060399956213761?s=20&t=3xoxeOdMjZTM1kdGzvLADw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "try: \n",
    "    experiment_id = mlflow.create_experiment(working_dir)\n",
    "except:\n",
    "    experiment_id = mlflow.get_experiment_by_name(working_dir)\n",
    "\n",
    "experiment = mlflow.set_experiment(working_dir)\n",
    "\n",
    "mlflow.pytorch.autolog()\n",
    "mlflow.start_run()\n",
    "run = mlflow.active_run()\n",
    "run_id = run.info.run_id"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Loop<br>\n",
    "\n",
    "a couple of design choices:\n",
    "- data loaders pin_memory=False: the data is already on the GPU so this won't work otherwise\n",
    "- the training loop is a bit artisanal on purpose; my apologies to all the Pytorch Lightning fans.\n",
    "- mixed precision to fit the architecture into the GPU and to speed up the training (GradScaler and autocast)\n",
    "- a bit of gradient accumulation as I don't have much trust in batch size of 2 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#need to scale learning rate due to small batch sizes\n",
    "learning_rate = args['model_default_lr'] * (args['batch_size'] / args['model_default_batch'])\n",
    "args['learning_rate'] = learning_rate\n",
    "\n",
    "train_loader = DataLoader(ds['train'].with_transform(train_transform_fn), batch_size=args['batch_size'], shuffle=True, collate_fn=collate_fn, num_workers=args['num_workers'], pin_memory=False,drop_last=True)\n",
    "val_loader = DataLoader(ds['validation'].with_transform(val_transform_fn), batch_size=args['batch_size'], shuffle=False, collate_fn=collate_fn, num_workers=args['num_workers'], pin_memory=False)\n",
    "\n",
    "optimizer = optim.AdamW(dualdetector.parameters(), lr=learning_rate,weight_decay=args['weight_decay'])\n",
    "scheduler = CosineAnnealingLR(optimizer, T_max=args['num_epochs'], eta_min=0, last_epoch=-1)\n",
    "\n",
    "#amp is all you need\n",
    "scaler = GradScaler()\n",
    "\n",
    "start_epoch = 0\n",
    "\n",
    "if (args['resume_from_checkpoint'] is not None):\n",
    "    checkpoint = torch.load(checkpoint_dir + args['resume_from_checkpoint'],map_location='cpu') #load checkpoint to CPU first and only load the model to GPU once the state dicionary is restored\n",
    "    dualdetector.load_state_dict(checkpoint['model_state_dict'])\n",
    "    dualdetector.to(torch_device)\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict']) #the model should be on GPU already when the optimizer and the scheduler are built\n",
    "    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
    "    torch.set_rng_state(checkpoint['torch_rng_state']) \n",
    "    start_epoch = checkpoint['epoch']+1\n",
    "    print('resuming from epoch ', start_epoch)\n",
    "else:\n",
    "    dualdetector.to(torch_device)\n",
    "\n",
    "args['start_epoch'] = start_epoch\n",
    "torch.set_float32_matmul_precision('medium') #RTX3090 specific optimization; \n",
    "\n",
    "mlflow.log_params(args)\n",
    "\n",
    "\n",
    "for epoch in range(start_epoch,args['num_epochs']):\n",
    "    running_loss = 0.0\n",
    "    dualdetector.train()\n",
    "    with torch.set_grad_enabled(True):\n",
    "        with tqdm(total=len(train_loader)) as pbar:\n",
    "            for i, (x_rgb,x_fft, labels) in enumerate(train_loader):\n",
    "                \n",
    "                #using autocast + scaler for mixed precision \n",
    "                with autocast(device_type=torch_device, dtype=torch.float16):\n",
    "\n",
    "                    outputs = dualdetector(x_rgb,x_fft).squeeze(1)\n",
    "                    \n",
    "                    loss = dualdetector.criterion(outputs, labels)\n",
    "                \n",
    "                scaler.scale(loss).backward()\n",
    "                \n",
    "                if (i+1) % args['gradient_accum'] == 0:\n",
    "                    \n",
    "                    if (args['gradient_clip'] is not None):\n",
    "                        scaler.unscale_(optimizer)\n",
    "                        torch.nn.utils.clip_grad_norm_(dualdetector.parameters(), max_norm=args['gradient_clip'])\n",
    "                    \n",
    "                    scaler.step(optimizer)\n",
    "                    scaler.update()\n",
    "                    optimizer.zero_grad()\n",
    "\n",
    "                running_loss +=loss.item()\n",
    "                pbar.set_description(f\"running_loss:{(running_loss/(i+1)):.6f}\")\n",
    "                pbar.update(1)\n",
    "            \n",
    "        current_lr = optimizer.param_groups[0]['lr']\n",
    "        scheduler.step()\n",
    "\n",
    "        #eval step\n",
    "        dualdetector.eval()\n",
    "        with torch.no_grad():\n",
    "            running_val_loss = 0.0        \n",
    "            for x_rgb,x_fft, labels in tqdm(val_loader):\n",
    "                with autocast(device_type=torch_device, dtype=torch.float16):\n",
    "\n",
    "                    outputs = dualdetector(x_rgb,x_fft).squeeze(1)\n",
    "                    \n",
    "                    val_loss = dualdetector.criterion(outputs, labels)\n",
    "                running_val_loss += val_loss\n",
    "                dualdetector.val_accuracy.update(outputs.sigmoid(), labels)\n",
    "\n",
    "    \n",
    "    #loggig\n",
    "    running_loss = running_loss/len(train_loader)\n",
    "    running_val_loss = running_val_loss/len(val_loader)\n",
    "    val_accuracy = dualdetector.val_accuracy.compute()\n",
    "    print(f\"Epoch: {epoch}/{args['num_epochs']}, Loss: {running_loss}, Validation Loss: {running_val_loss}, Validation Accuracy: {val_accuracy}, Learning Rate: {current_lr} \")\n",
    "\n",
    "    mlflow.log_metric('lr',current_lr,step=epoch) \n",
    "    mlflow.log_metric('loss',running_loss,step=epoch) \n",
    "    mlflow.log_metric('val_loss',running_val_loss,step=epoch) \n",
    "    mlflow.log_metric('val_accuracy',val_accuracy,step=epoch) \n",
    "\n",
    "    #checkpointing\n",
    "    checkpoint_path = checkpoint_dir + f\"DualDetector_{run_id}_epoch={epoch}_val_loss={round(float(running_val_loss),5)}.pth\"\n",
    "    torch.save({\n",
    "                'epoch': epoch,\n",
    "                'model_state_dict': dualdetector.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'scheduler_state_dict': scheduler.state_dict(),\n",
    "                'torch_rng_state': torch.get_rng_state(),\n",
    "                'loss': loss,\n",
    "                'batch_size': args['batch_size'],\n",
    "                'learning_rate': learning_rate,\n",
    "                'current_lr': current_lr,\n",
    "                'weight_decay': args['weight_decay']\n",
    "                }, \n",
    "                checkpoint_path)\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.end_run()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "checkpoint_name = 'DualDetector_23a34a28b8024139bf238ce289204654_epoch=29_val_loss=1e-05.pth'\n",
    "\n",
    "dualdetector = DualDetector(\n",
    "                            rgb_model_name=args['rgb_model_name'],\n",
    "                            fft_model_name=args['fft_model_name'],\n",
    "                            rgb_model_pretrained=args['rgb_model_pretrained'], \n",
    "                            fft_model_pretrained=args['fft_model_pretrained'],\n",
    "                            rgb_model_frozen=args['rgb_model_frozen'], \n",
    "                            fft_model_frozen=args['fft_model_frozen']                            \n",
    "                            )\n",
    "\n",
    "checkpoint = torch.load(checkpoint_dir + checkpoint_name)\n",
    "dualdetector.load_state_dict(checkpoint['model_state_dict'])\n",
    "\n",
    "dualdetector.to(torch_device)\n",
    "dualdetector.eval()\n",
    "torch.set_float32_matmul_precision('medium') #RTX3090 specific optimization; \n",
    "\n",
    "test_loader = DataLoader(ds['test'].with_transform(test_transform_fn), batch_size=16, shuffle=False, collate_fn=collate_fn_test, num_workers=0, pin_memory=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "all_preds_sigm = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    with autocast(device_type=torch_device, dtype=torch.float16):\n",
    "        for x_rgb,x_fft in tqdm(test_loader):\n",
    "            x_rgb = x_rgb.to(torch_device)\n",
    "            x_fft = x_fft.to(torch_device)\n",
    "            out = dualdetector(x_rgb,x_fft)\n",
    "            all_preds_sigm.extend(out.squeeze().sigmoid().detach().cpu().numpy().tolist())\n",
    "           \n",
    "\n",
    "\n",
    "\n",
    "df_sigm = pd.DataFrame(ds['test'].remove_columns(['image']))     \n",
    "     \n",
    "df_sigm['label'] = all_preds_sigm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sigm.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sigm.to_csv(checkpoint_name + '.csv', index=None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aiornot",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2f89f0ba50e4a3d9eef35b528037a641bcd2b049665a29bbd49302dabc94daad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}