src/train_ProbVLM_CLIP.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2986807e-7850-404f-a957-eaeb16371d24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from os.path import join as ospj\n",
    "from os.path import expanduser\n",
    "from munch import Munch as mch\n",
    "import numpy as np\n",
    "\n",
    "from ds import prepare_coco_dataloaders, prepare_flickr_dataloaders, prepare_cub_dataloaders, prepare_flo_dataloaders\n",
    "\n",
    "from utils import *\n",
    "from networks import *\n",
    "from train_ProbVLM import *\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbdb7f39-293e-48fe-bd3b-da616157f4a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'CUB' # coco or flickr\n",
    "data_dir = ospj('/mnt/Datasets/', dataset) # e.g. ospj(expanduser('~'), 'Documents', 'jm', 'data', dataset)\n",
    "dataloader_config = mch({\n",
    "    'batch_size': 64,\n",
    "    'random_erasing_prob': 0.,\n",
    "    'traindata_shuffle': True\n",
    "})\n",
    "loaders,vocab = load_data_loader(dataset, data_dir, dataloader_config)\n",
    "cub_train_loader, cub_valid_loader, cub_test_loader = loaders['train'], loaders['val'], loaders['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7da29997-5483-4fa0-b927-74b38dc36cdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# clip_net = load_model('cuda')\n",
    "CLIP_Net = load_model(device='cuda', model_path=None)\n",
    "ProbVLM_Net = BayesCap_for_CLIP(\n",
    "    inp_dim=512,\n",
    "    out_dim=512,\n",
    "    hid_dim=256,\n",
    "    num_layers=3,\n",
    "    p_drop=0.05,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1513f40-28d1-4d4d-8d9b-e113c8cd4183",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ProbVLM(\n",
    "    CLIP_Net,\n",
    "    ProbVLM_Net,\n",
    "    cub_train_loader,\n",
    "    cub_valid_loader,\n",
    "    Cri = TempCombLoss(),\n",
    "    device='cuda',\n",
    "    dtype=torch.cuda.FloatTensor,\n",
    "    init_lr=8e-5,\n",
    "    num_epochs=500,\n",
    "    eval_every=5,\n",
    "    ckpt_path='../ckpt/ProbVLM_Net',\n",
    "    T1=1e0,\n",
    "    T2=1e-4\n",
    ")"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}