ExampleNotebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "475c651d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Imports\n",
    "import os, sys\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import functools\n",
    "\n",
    "from VICReg.vicreg        import GeneralMultipleVICReg\n",
    "from VICReg.vicreg_utils  import create_projector, create_resnet, create_adam_opt\n",
    "from VICReg.dataset_utils import preprocess_ds\n",
    "from VICReg.augmentations import custom_augment_image\n",
    "from VICReg.warmup_learning_rate import WarmUpLR\n",
    "from VICReg.warmupcosine import WarmUpCosine\n",
    "from VICReg.classifier    import ClusterClassifier, classifier_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0077d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Hyperparameters\n",
    "\n",
    "AUTO = tf.data.AUTOTUNE\n",
    "SEED = 42\n",
    "\n",
    "PROJECT_DIM = 2048\n",
    "BATCH_SIZE  = 128\n",
    "EPOCHS      = 100\n",
    "IM_SIZE = 224\n",
    "DATASET_SIZE = 100000\n",
    "STEPS_PER_EPOCH = DATASET_SIZE//BATCH_SIZE\n",
    "WARMUP_EPOCHS = EPOCHS * 0.001\n",
    "WARMUP_STEPS = int(WARMUP_EPOCHS * STEPS_PER_EPOCH)\n",
    "SHUFFLE_BUFFER = 2**10\n",
    "\n",
    "path_train = \"\"\n",
    "path_test = \"\"\n",
    "MODEL_SAVE_PATH = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df93965",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Functions\n",
    "\n",
    "lr_decayed_fn = WarmUpCosine(learning_rate_base=1e-4,\n",
    "                             total_steps=EPOCHS*STEPS_PER_EPOCH,\n",
    "                             warmup_learning_rate=0.0,\n",
    "                             warmup_steps=WARMUP_STEPS\n",
    "                             )\n",
    "\n",
    "augment_im = lambda x: custom_augment_image(x, input_shape=(IM_SIZE,IM_SIZE,3), output_shape=(IM_SIZE,IM_SIZE,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55cecec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dataset\n",
    "\n",
    "image_train_ds = load_dataset(path_train)\n",
    "image_valid_ds = load_dataset(path_test)\n",
    "patch_valid_ds = image_valid_ds.map(lambda x: (augment_im(x['original_images']), augment_im(tf.stack([x['masked_images']for _ in range(3)],-1))), num_parallel_calls=AUTO)\n",
    "patch_valid_ds = patch_valid_ds.batch(BATCH_SIZE).prefetch(AUTO)\n",
    "patch_train_ds = image_train_ds.map(lambda x: (augment_im(x['original_images']), augment_im(tf.stack([x['masked_images']for _ in range(3)],-1))), num_parallel_calls=AUTO)\n",
    "patch_train_ds = preprocess_ds(patch_train_ds, batch_size=BATCH_SIZE, seed=SEED, pre=AUTO, shuffle_no=SHUFFLE_BUFFER, rei=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56fec9a7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "input_shape = (IM_SIZE, IM_SIZE, 3)\n",
    "encoder1    = create_resnet(input_shape)\n",
    "projector1  = create_projector()\n",
    "encoder2    = create_resnet((IM_SIZE,IM_SIZE,3))\n",
    "projector2      = create_projector()\n",
    "\n",
    "optimizer_enc1  = create_adam_opt(lr_decayed_fn)\n",
    "optimizer_proj1 = create_adam_opt(lr_decayed_fn)\n",
    "optimizer_enc2  = create_adam_opt(lr_decayed_fn)\n",
    "optimizer_proj2 = create_adam_opt(lr_decayed_fn)\n",
    "\n",
    "enc_list            = [encoder1, encoder2]\n",
    "proj_list           = [projector1, projector2]\n",
    "optimizer_list_enc  = [optimizer_enc1, optimizer_enc2]\n",
    "optimizer_list_proj = [optimizer_proj1, optimizer_proj2]\n",
    "optimizer_list      = [optimizer_list_enc, optimizer_list_proj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89e1d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, enc in enumerate(enc_list):\n",
    "    enc.load_weights(os.path.join(MODEL_SAVE_PATH, f'encoder_weights_{i}'))\n",
    "for i, enc in enumerate(proj_list):\n",
    "    enc.load_weights(os.path.join(MODEL_SAVE_PATH, f'projector_weights_{i}'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfe5b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "vicreg = GeneralMultipleVICReg(encoder_list=enc_list, projector_list=proj_list, encoder_indices=[0,1], projector_indices=[0,1])\n",
    "vicreg.compile(optimizer=optimizer_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf97a719",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vicreg.fit(patch_train_ds,\n",
    "           epochs=EPOCHS,\n",
    "           callbacks=[],\n",
    "           validation_data=patch_valid_ds\n",
    "          )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b6134b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, enc in enumerate(enc_list):\n",
    "    enc.save_weights(os.path.join(MODEL_SAVE_PATH, f'encoder_weights_{i}'))\n",
    "for i, enc in enumerate(proj_list):\n",
    "    enc.save_weights(os.path.join(MODEL_SAVE_PATH, f'projector_weights_{i}'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "032de42a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tissue Classification Patches\n",
    "\n",
    "task_train_ds = image_train_ds.map(lambda x: (x['original_images'], x['tissue_types']))\n",
    "task_train_ds = task_train_ds.map(lambda x, y: (augment_im(x), y), num_parallel_calls=AUTO)\n",
    "task_train_ds = preprocess_ds(task_train_ds, batch_size=BATCH_SIZE, seed=SEED, pre=AUTO, shuffle_no=SHUFFLE_BUFFER, rei=True)\n",
    "task_train_ds = task_train_ds.map(lambda x, y: (x, tf.reshape(y, shape=(-1,1))), num_parallel_calls=AUTO)\n",
    "task_test_ds = image_valid_ds.map(lambda x: (tf.cast(x['original_images'], tf.float32),  x['tissue_types']), num_parallel_calls=AUTO).batch(BATCH_SIZE).prefetch(AUTO)\n",
    "task_test_ds = task_test_ds.map(lambda x, y: (x, tf.reshape(y, shape=(-1,1))), num_parallel_calls=AUTO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bad0c00f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "classifier_head = classifier_class()\n",
    "classifier = ClusterClassifier(vicreg.encoder_list[0], classifier_head)\n",
    "classifier.compile(optimizer='adam',\n",
    "                   loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),\n",
    "                   metrics=[\"accuracy\"]\n",
    "                  )\n",
    "classifier.fit(task_train_ds, epochs=100, validation_data=task_test_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98857d45",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier.evaluate(task_test_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "283c47c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier.save_weights(os.path.join(MODEL_SAVE_PATH, 'classifier_weights'))"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}