self_bimodal_BYOL_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8127e2a9-a082-403f-ba27-c3217d759fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.BYOL2_model import BYOL2\n",
    "from data.custom_transforms import BatchTransform, ListToTensor, PadToSquare, SelectFromTuple\n",
    "from data.pairs_dataset import PairsDataset, pair_collate_fn\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "import sys\n",
    "import torch\n",
    "from torch.utils.data import Subset\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as T\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "029f13cc-b988-4b8b-a21f-8311dd057a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = PairsDataset(\n",
    "    '/users/jmorales/Shoes/images_train/',\n",
    "    '/users/jmorales/Shoes/images_train_pidinet/'\n",
    ")\n",
    "\n",
    "transforms_1 = T.Compose([\n",
    "    BatchTransform(SelectFromTuple(0)),\n",
    "    BatchTransform(PadToSquare(255)),\n",
    "    BatchTransform(T.Resize((224,224))),\n",
    "    ListToTensor('cuda', torch.float),\n",
    "])\n",
    "transforms_2 = T.Compose([\n",
    "    BatchTransform(SelectFromTuple(1)),\n",
    "    BatchTransform(PadToSquare(255)),\n",
    "    BatchTransform(T.Resize((224,224))),\n",
    "    ListToTensor('cuda', torch.float),\n",
    "])\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=50,\n",
    "    shuffle=True,\n",
    "    collate_fn=pair_collate_fn,\n",
    "    num_workers=4\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d9619e0e-3a22-462b-8463-a531cc6cc5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = models.resnet50(pretrained=False)\n",
    "encoder.load_state_dict(torch.load('../checkpoints/resnet50_byol_quickdraw_128_1000_v3.pt'))\n",
    "empty_transform = T.Compose([])\n",
    "epochs = 5\n",
    "epoch_size = len(train_loader)\n",
    "learner = BYOL2(\n",
    "    encoder,\n",
    "    image_size=224,\n",
    "    hidden_layer='avgpool',\n",
    "    augment_fn=empty_transform,\n",
    "    cosine_ema_steps=epochs*epoch_size\n",
    ")\n",
    "learner.augment1 = transforms_1\n",
    "learner.augment2 = transforms_2\n",
    "opt = torch.optim.Adam(learner.parameters(), lr=3e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "74c4456a-d7c1-49e8-a4a7-25b0a1552e87",
   "metadata": {},
   "outputs": [],
   "source": [
    "learner.load_state_dict(torch.load('../checkpoints/bimodal_byol_resnet50_pretrained_sketchy_v5.pt'))\n",
    "torch.save(learner.state_dict(), '../checkpoints/self_bimodal_byol_sketchy_then_shoes_{}epochs.pt'.format(epochs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "90abbedb-c016-45fa-8cdd-a5aaa6f83a2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, batch 36 - loss 0.2459\n",
      "Epoch 2, batch 36 - loss 0.1377\n",
      "Epoch 3, batch 36 - loss 0.0938\n",
      "Epoch 4, batch 36 - loss 0.0688\n",
      "Epoch 5, batch 36 - loss 0.0543\n"
     ]
    }
   ],
   "source": [
    "learner.load_state_dict(torch.load('../checkpoints/self_bimodal_byol_sketchy_then_shoes_{}epochs.pt'.format(epochs)))\n",
    "learner = learner.to('cuda')\n",
    "learner.train()\n",
    "filehandler = open('../checkpoints/training_bimodal_byol.txt', 'w')\n",
    "with warnings.catch_warnings():\n",
    "    warnings.filterwarnings('ignore')\n",
    "    running_loss = np.array([], dtype=np.float32)\n",
    "    for epoch in range(epochs):\n",
    "        i = 0\n",
    "        for images in train_loader:\n",
    "            loss = learner(images) #.to('cuda', dtype=torch.float))\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            learner.update_moving_average()\n",
    "            running_loss = np.append(running_loss, [loss.item()])\n",
    "            sys.stdout.write('\\rEpoch {}, batch {} - loss {:.4f}'.format(epoch+1, i+1, np.mean(running_loss)))\n",
    "            filehandler.write('Epoch {}, batch {} - loss {:.4f}\\n'.format(epoch+1, i+1, np.mean(running_loss)))\n",
    "            filehandler.flush()\n",
    "            i += 1\n",
    "            if i%(epoch_size/2)==0:\n",
    "                torch.save(learner.state_dict(), '../checkpoints/self_bimodal_byol_sketchy_then_shoes_{}epochs.pt'.format(epochs))\n",
    "        running_loss = np.array([], dtype=np.float32)\n",
    "        sys.stdout.write('\\n')\n",
    "filehandler.close()"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}