model_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.10.1 True\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np \n",
    "import json\n",
    "import pickle\n",
    "import cv2\n",
    "import re\n",
    "from pprint import pprint\n",
    "\n",
    "from PIL import Image, ImageDraw\n",
    "\n",
    "import os\n",
    "\n",
    "print(torch.__version__,torch.cuda.is_available())\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "torch.cuda.device_count()\n",
    "\n",
    "import detectron2\n",
    "from detectron2.utils.logger import setup_logger\n",
    "from detectron2.data.datasets import register_coco_instances, load_coco_json\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.engine import DefaultPredictor, default_argument_parser, \\\n",
    "                              default_setup, hooks, launch, DefaultTrainer, HookBase, default_writers\n",
    "from detectron2.config import CfgNode as CN, get_cfg\n",
    "from detectron2.modeling import build_model\n",
    "import detectron2.utils.comm as comm\n",
    "from detectron2.utils.visualizer import Visualizer, ColorMode\n",
    "from detectron2.solver import build_lr_scheduler, build_optimizer\n",
    "from detectron2.data import detection_utils as utils, build_detection_test_loader, DatasetMapper, \\\n",
    "                            build_detection_train_loader, MetadataCatalog, DatasetCatalog\n",
    "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
    "from detectron2.checkpoint import DetectionCheckpointer, PeriodicCheckpointer\n",
    "from detectron2.utils.events import EventStorage\n",
    "\n",
    "logger = setup_logger()\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "\n",
    "from modules import CBMCOCOEvaluator, CustomMapper, add_cbm_config, add_uhcc_config, MyVisualizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# register sample dataset in Detectron2 style \n",
    "# as per Detectron2, the first argument in load_coco_json has to be a fully-qualified path \n",
    "DatasetCatalog.register(\"sample\", lambda: load_coco_json(\"/YOUR_FILE_PATH/sample_data/sample_data.json\", \"\", \"sample\", ['region_shape', 'region_orientation', 'region_margin', 'region_echo', 'region_posterior', 'region_cancer']))\n",
    "\n",
    "MetadataCatalog.get(\"sample\").shape_classes = ['oval', 'not oval']\n",
    "MetadataCatalog.get(\"sample\").orientation_classes = ['parallel', 'not parallel']\n",
    "MetadataCatalog.get(\"sample\").margin_classes = ['circumscribed', 'not circumscribed']\n",
    "MetadataCatalog.get(\"sample\").echo_classes = ['anechoic', 'not anechoic']\n",
    "MetadataCatalog.get(\"sample\").posterior_classes = ['no features', 'features']\n",
    "MetadataCatalog.get(\"sample\").cancer_classes = ['benign', 'malignant']\n",
    "\n",
    "MetadataCatalog.get(\"sample\").thing_colors = [\"g\", \"m\"]\n",
    "MetadataCatalog.get(\"sample\").thing_classes = ['lesion']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Cancer Model(s)\n",
    "As training the lesion-only and concept prediction models can be handled by the standard Detectron2 training loop (just remember to freeze the backbone when training the concept predictions on top), we only provide an example for training the different cancer head configurations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GeneralizedRCNN(\n",
       "  (backbone): FPN(\n",
       "    (fpn_lateral2): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (fpn_output2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (fpn_lateral3): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (fpn_output3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (fpn_lateral4): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (fpn_output4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (fpn_lateral5): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (fpn_output5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (top_block): LastLevelMaxPool()\n",
       "    (bottom_up): ResNet(\n",
       "      (stem): BasicStem(\n",
       "        (conv1): Conv2d(\n",
       "          3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
       "          (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "        )\n",
       "      )\n",
       "      (res2): Sequential(\n",
       "        (0): BottleneckBlock(\n",
       "          (shortcut): Conv2d(\n",
       "            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv1): Conv2d(\n",
       "            64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (1): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (2): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=64, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (res3): Sequential(\n",
       "        (0): BottleneckBlock(\n",
       "          (shortcut): Conv2d(\n",
       "            256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv1): Conv2d(\n",
       "            256, 128, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (1): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (2): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (3): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=128, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (res4): Sequential(\n",
       "        (0): BottleneckBlock(\n",
       "          (shortcut): Conv2d(\n",
       "            512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "          (conv1): Conv2d(\n",
       "            512, 256, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (1): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (2): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (3): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (4): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (5): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (6): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (7): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (8): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (9): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (10): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (11): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (12): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (13): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (14): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (15): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (16): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (17): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (18): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (19): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (20): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (21): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (22): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=256, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=1024, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (res5): Sequential(\n",
       "        (0): BottleneckBlock(\n",
       "          (shortcut): Conv2d(\n",
       "            1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)\n",
       "          )\n",
       "          (conv1): Conv2d(\n",
       "            1024, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (1): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "        (2): BottleneckBlock(\n",
       "          (conv1): Conv2d(\n",
       "            2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv2): Conv2d(\n",
       "            512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=512, eps=1e-05)\n",
       "          )\n",
       "          (conv3): Conv2d(\n",
       "            512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "            (norm): FrozenBatchNorm2d(num_features=2048, eps=1e-05)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (proposal_generator): RPN(\n",
       "    (rpn_head): StandardRPNHead(\n",
       "      (conv): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (objectness_logits): Conv2d(256, 3, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (anchor_deltas): Conv2d(256, 12, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "    (anchor_generator): DefaultAnchorGenerator(\n",
       "      (cell_anchors): BufferList()\n",
       "    )\n",
       "  )\n",
       "  (roi_heads): CBMStandardROIHeads(\n",
       "    (box_pooler): ROIPooler(\n",
       "      (level_poolers): ModuleList(\n",
       "        (0): ROIAlign(output_size=(7, 7), spatial_scale=0.25, sampling_ratio=0, aligned=True)\n",
       "        (1): ROIAlign(output_size=(7, 7), spatial_scale=0.125, sampling_ratio=0, aligned=True)\n",
       "        (2): ROIAlign(output_size=(7, 7), spatial_scale=0.0625, sampling_ratio=0, aligned=True)\n",
       "        (3): ROIAlign(output_size=(7, 7), spatial_scale=0.03125, sampling_ratio=0, aligned=True)\n",
       "      )\n",
       "    )\n",
       "    (box_head): FastRCNNConvFCHead(\n",
       "      (conv1): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (conv2): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (conv3): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (conv4): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "      (fc1): Linear(in_features=12544, out_features=1024, bias=True)\n",
       "      (fc_relu1): ReLU()\n",
       "      (fc2): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "      (fc_relu2): ReLU()\n",
       "      (fc3): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "      (fc_relu3): ReLU()\n",
       "    )\n",
       "    (box_predictor): FastRCNNOutputLayers(\n",
       "      (cls_score): Linear(in_features=1024, out_features=2, bias=True)\n",
       "      (bbox_pred): Linear(in_features=1024, out_features=4, bias=True)\n",
       "    )\n",
       "    (mask_pooler): ROIPooler(\n",
       "      (level_poolers): ModuleList(\n",
       "        (0): ROIAlign(output_size=(14, 14), spatial_scale=0.25, sampling_ratio=0, aligned=True)\n",
       "        (1): ROIAlign(output_size=(14, 14), spatial_scale=0.125, sampling_ratio=0, aligned=True)\n",
       "        (2): ROIAlign(output_size=(14, 14), spatial_scale=0.0625, sampling_ratio=0, aligned=True)\n",
       "        (3): ROIAlign(output_size=(14, 14), spatial_scale=0.03125, sampling_ratio=0, aligned=True)\n",
       "      )\n",
       "    )\n",
       "    (mask_head): MaskRCNNConvUpsampleHead(\n",
       "      (mask_fcn1): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (mask_fcn2): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (mask_fcn3): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (mask_fcn4): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (mask_fcn5): Conv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n",
       "        (activation): ReLU()\n",
       "      )\n",
       "      (deconv): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))\n",
       "      (deconv_relu): ReLU()\n",
       "      (predictor): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "    )\n",
       "    (cancer_head): CBMCancerHead(\n",
       "      (first_model): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU()\n",
       "          (shortcut): Sequential(\n",
       "            (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "        (2): Bottleneck(\n",
       "          (conv1): Conv2d(512, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU()\n",
       "          (shortcut): Sequential(\n",
       "            (0): Conv2d(512, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (3): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      )\n",
       "      (transfer_concepts): Sequential(\n",
       "        (0): Bottleneck(\n",
       "          (conv1): Conv2d(64, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn1): BatchNorm2d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (conv2): Conv2d(5, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn2): BatchNorm2d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (relu): ReLU()\n",
       "          (shortcut): Sequential(\n",
       "            (0): Conv2d(64, 5, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "        (2): Flatten(start_dim=1, end_dim=-1)\n",
       "      )\n",
       "      (second_model): Sequential(\n",
       "        (0): Linear(in_features=5, out_features=512, bias=True)\n",
       "        (1): ReLU()\n",
       "        (2): Linear(in_features=512, out_features=512, bias=True)\n",
       "        (3): ReLU()\n",
       "        (4): Linear(in_features=512, out_features=512, bias=True)\n",
       "        (5): ReLU()\n",
       "        (6): Linear(in_features=512, out_features=1, bias=True)\n",
       "      )\n",
       "      (loss): BCEWithLogitsLoss()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cfg3a = get_cfg()\n",
    "cfg3a.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml\"))\n",
    "add_uhcc_config(cfg3a)\n",
    "add_cbm_config(cfg3a)\n",
    "cfg3a.MODEL.WEIGHTS = None\n",
    "cfg3a.merge_from_file(\"configs/stage_3a.yaml\")\n",
    "cfg3a.DATALOADER.NUM_WORKERS = 0\n",
    "#reduce the number of detections for visualization purposes \n",
    "cfg3a.TEST.DETECTIONS_PER_IMAGE = 1\n",
    "\n",
    "model3a.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[03/06 21:49:30 d2.data.datasets.coco]: \u001b[0mLoaded 5 images in COCO format from /YOUR_FILE_PATH/sample_data/sample_data.json\n",
      "\u001b[32m[03/06 21:49:30 d2.data.dataset_mapper]: \u001b[0m[DatasetMapper] Augmentations used in training: [ResizeShortestEdge(short_edge_length=(640, 672, 704, 736, 768, 800), max_size=1333, sample_style='choice'), RandomFlip()]\n",
      "\u001b[32m[03/06 21:49:30 d2.data.common]: \u001b[0mSerializing 5 elements to byte tensors and concatenating them all ...\n",
      "\u001b[32m[03/06 21:49:30 d2.data.common]: \u001b[0mSerialized dataset takes 0.00 MiB\n",
      "\u001b[32m[03/06 21:49:30 detectron2]: \u001b[0mStarting training from iteration 0\n",
      "\u001b[32m[03/06 21:49:30 d2.data.dataset_mapper]: \u001b[0m[DatasetMapper] Augmentations used in training: [ResizeShortestEdge(short_edge_length=(640, 672, 704, 736, 768, 800), max_size=1333, sample_style='choice'), RandomFlip()]\n",
      "\u001b[32m[03/06 21:49:30 d2.data.common]: \u001b[0mSerializing 5 elements to byte tensors and concatenating them all ...\n",
      "\u001b[32m[03/06 21:49:30 d2.data.common]: \u001b[0mSerialized dataset takes 0.00 MiB\n",
      "\u001b[32m[03/06 21:49:30 detectron2]: \u001b[0mStarting training from iteration 0\n",
      "\u001b[32m[03/06 21:49:35 d2.utils.events]: \u001b[0m iter: 19  total_loss: 2.63  loss_cls: 0.03462  loss_box_reg: 0.04154  loss_mask: 0.1658  shape_loss: 0.6075  margin_loss: 0.701  orientation_loss: 0.2335  echo_loss: 0.02818  posterior_loss: 0.2377  cancer_loss: 0.5729  loss_rpn_cls: 0.001968  loss_rpn_loc: 0.002344  lr: 7.9924e-06  max_mem: 6101M\n",
      "\u001b[32m[03/06 21:49:37 d2.utils.events]: \u001b[0m eta: 0:00:00  iter: 24  total_loss: 2.64  loss_cls: 0.03468  loss_box_reg: 0.04153  loss_mask: 0.1656  shape_loss: 0.6259  margin_loss: 0.7069  orientation_loss: 0.2323  echo_loss: 0.02848  posterior_loss: 0.2367  cancer_loss: 0.5632  loss_rpn_cls: 0.001813  loss_rpn_loc: 0.002345  lr: 9.9904e-06  max_mem: 6101M\n"
     ]
    }
   ],
   "source": [
    "# the structure of this train loop is provided in the Detectron2 GitHub repository. \n",
    "# to implement validation loss computation during the training loop, use the ValidationMapper class\n",
    "\n",
    "# we use the existing_weights argument when training the side channel on an existing concept-only model \n",
    "existing_weights = None\n",
    "# resuming from previous training or not \n",
    "resume = False\n",
    "\n",
    "optimizer = build_optimizer(cfg3a, model3a)\n",
    "scheduler = build_lr_scheduler(cfg3a, optimizer)\n",
    "\n",
    "checkpointer = DetectionCheckpointer(\n",
    "    model3a, save_dir='output', optimizer=optimizer, scheduler=scheduler\n",
    ")\n",
    "start_iter = (\n",
    "    checkpointer.resume_or_load(cfg3a.MODEL.WEIGHTS, resume=resume).get(\"iteration\", -1) + 1\n",
    ")\n",
    "max_iter = 25\n",
    "\n",
    "periodic_checkpointer = PeriodicCheckpointer(\n",
    "    checkpointer, cfg3a.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter\n",
    ")\n",
    "\n",
    "writers = default_writers(cfg3a.OUTPUT_DIR, max_iter) if comm.is_main_process() else []\n",
    "\n",
    "data_loader = build_detection_train_loader(DatasetCatalog.get(\"sample\"), mapper=CustomMapper(cfg3a, is_train=True, augmentations=[]), total_batch_size=8)\n",
    "\n",
    "non_cancer_losses = ['loss_cls', 'loss_box_reg', 'loss_mask',  'shape_loss',  'margin_loss',  \n",
    "                        'orientation_loss', 'echo_loss', 'posterior_loss', 'loss_rpn_cls',  'loss_rpn_loc']\n",
    "\n",
    "logger.info(\"Starting training from iteration {}\".format(start_iter))\n",
    "with EventStorage(start_iter) as storage:\n",
    "    for data, iteration in zip(data_loader, range(start_iter, max_iter)):\n",
    "        storage.iter = iteration\n",
    "\n",
    "        loss_dict = model3a(data)\n",
    "        losses = sum(loss_dict.values()) \n",
    "        loss_dict_back = loss_dict.copy()\n",
    "\n",
    "        # popping off all the other losses\n",
    "        for e in non_cancer_losses: \n",
    "            if e in loss_dict_back:\n",
    "                loss_dict_back.pop(e)\n",
    "\n",
    "        losses_back = sum(loss_dict_back.values()) \n",
    "\n",
    "        assert torch.isfinite(losses).all(), loss_dict\n",
    "\n",
    "        loss_dict_reduced = {k: v.item() for k, v in comm.reduce_dict(loss_dict).items()}\n",
    "        losses_reduced = sum(loss for loss in loss_dict_reduced.values())\n",
    "\n",
    "        if comm.is_main_process():\n",
    "            storage.put_scalars(total_loss=losses_reduced, **loss_dict_reduced)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        losses_back.backward()       \n",
    "        optimizer.step()\n",
    "\n",
    "        if cfg3a.MODEL.CBM.SIDE_CHANNEL and (existing_weights is not None):\n",
    "            # making a copy of the adjusted weight \n",
    "            c = copy.deepcopy(model3a.roi_heads.cancer_head.second_model[0].weight)\n",
    "            # slice out the new weights for the side_channel \n",
    "            d = c[:, 5:]\n",
    "            # combine original values with placeholder \n",
    "            weight_val = torch.cat((existing_weights, d), dim=1)\n",
    "            # assign weight to the combined version of old + new\n",
    "            model3a.roi_heads.cancer_head.second_model[0].weight = torch.nn.Parameter(weight_val)\n",
    "        elif not cfg3a.MODEL.CBM.SIDE_CHANNEL and (existing_weights is None):\n",
    "            pass\n",
    "        else:\n",
    "            raise ValueError('Mismatch between side channel and existing weights found')\n",
    "\n",
    "        storage.put_scalar(\"lr\", optimizer.param_groups[0][\"lr\"], smoothing_hint=False)\n",
    "        scheduler.step()\n",
    "\n",
    "        if iteration - start_iter > 5 and (\n",
    "            (iteration + 1) % 20 == 0 or iteration == max_iter - 1\n",
    "        ):\n",
    "            for writer in writers:\n",
    "                writer.write()\n",
    "        periodic_checkpointer.step(iteration)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_detectron_ci",
   "language": "python",
   "name": "torch_detectron_ci"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}