train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "import time\n",
    "import torch\n",
    "from torch.cuda.amp import GradScaler, autocast\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "from data_loader import *\n",
    "from encoder_decoder import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_location =  \"dataa\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_dataset = data_location\n",
    "\n",
    "# Path to the directory containing image files\n",
    "img_dir = f\"{path_to_dataset}/Flickr8k_Dataset/Flicker8k_Dataset/\"\n",
    "\n",
    "# Path to the file containing captions\n",
    "captions_file = f\"{path_to_dataset}/Flickr8k_text/Flickr8k.token.txt\"\n",
    "\n",
    "# Path to the split files\n",
    "train_file = f\"{path_to_dataset}/Flickr8k_text/Flickr_8k.trainImages.txt\"\n",
    "val_file = f\"{path_to_dataset}/Flickr8k_text/Flickr_8k.devImages.txt\"\n",
    "test_file = f\"{path_to_dataset}/Flickr8k_text/Flickr_8k.testImages.txt\"\n",
    "\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Initialize datasets for training, validation, and testing\n",
    "train_dataset = Flickr8kDataset(img_dir, captions_file, train_file, transform=transform)\n",
    "val_dataset = Flickr8kDataset(img_dir, captions_file, val_file, transform=transform)\n",
    "test_dataset = Flickr8kDataset(img_dir, captions_file, test_file, transform=transform)\n",
    "\n",
    "# Create DataLoader instances\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=32,\n",
    "    shuffle=True,\n",
    "    collate_fn=Flickr8kDataset.collate_fn,\n",
    "    num_workers=4,\n",
    ")\n",
    "val_loader = DataLoader(\n",
    "    dataset=val_dataset,\n",
    "    batch_size=32,\n",
    "    shuffle=False,\n",
    "    collate_fn=Flickr8kDataset.collate_fn,\n",
    "    num_workers=4,\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    dataset=test_dataset,\n",
    "    batch_size=32,\n",
    "    shuffle=False,\n",
    "    collate_fn=Flickr8kDataset.collate_fn,\n",
    "    num_workers=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(train_dataset.stoi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Hyperparams\n",
    "embed_size=300\n",
    "attention_dim=256\n",
    "encoder_dim=2048\n",
    "decoder_dim=512\n",
    "learning_rate = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#init model\n",
    "model = EncoderDecoder(\n",
    "    embed_size=embed_size,\n",
    "    vocab_size = vocab_size,\n",
    "    attention_dim=attention_dim,\n",
    "    encoder_dim=encoder_dim,\n",
    "    decoder_dim=decoder_dim\n",
    ").to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=train_dataset.stoi[\"<PAD>\"])\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#helper function to save the model\n",
    "def save_model(model,num_epochs):\n",
    "    model_state = {\n",
    "        'num_epochs':num_epochs,\n",
    "        'embed_size':embed_size,\n",
    "        'vocab_size':vocab_size,\n",
    "        'attention_dim':attention_dim,\n",
    "        'encoder_dim':encoder_dim,\n",
    "        'decoder_dim':decoder_dim,\n",
    "        'state_dict':model.state_dict()\n",
    "    }\n",
    "\n",
    "    torch.save(model_state,'attention_model_state.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "from torch.cuda.amp import GradScaler, autocast\n",
    "from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_epochs = 15\n",
    "print_every = 100\n",
    "validate_every = 1\n",
    "scaler = GradScaler() \n",
    "\n",
    "best_val_loss = float(\"inf\")\n",
    "\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "bleu_scores = []\n",
    "\n",
    "for epoch in range(1, num_epochs + 1):\n",
    "    start_time = time.time()  \n",
    "\n",
    "    model.train()\n",
    "    total_train_loss = 0\n",
    "    for idx, (images, captions) in enumerate(train_loader):\n",
    "        images, captions = images.to(device, non_blocking=True), captions.to(\n",
    "            device, non_blocking=True\n",
    "        )\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        with autocast():\n",
    "            outputs, attentions = model(images, captions)\n",
    "            targets = captions[:, 1:]\n",
    "            loss = criterion(outputs.view(-1, vocab_size), targets.reshape(-1))\n",
    "            total_train_loss += loss.item()\n",
    "\n",
    "        scaler.scale(loss).backward()\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "\n",
    "        if (idx + 1) % print_every == 0:\n",
    "            print(f\"Epoch: {epoch}, Step: {idx + 1}, Loss: {loss.item():.5f}\")\n",
    "\n",
    "    avg_train_loss = total_train_loss / len(train_loader)\n",
    "    train_losses.append(avg_train_loss)\n",
    "\n",
    "    # Validation loop\n",
    "    if epoch % validate_every == 0:\n",
    "        model.eval()\n",
    "        total_val_loss = 0\n",
    "        references = []\n",
    "        hypotheses = []\n",
    "        with torch.no_grad():\n",
    "            for images, captions in val_loader:\n",
    "                images, captions = images.to(device, non_blocking=True), captions.to(\n",
    "                    device, non_blocking=True\n",
    "                )\n",
    "\n",
    "                with autocast():\n",
    "                    outputs, _ = model(images, captions)\n",
    "                    targets = captions[:, 1:]\n",
    "                    val_loss = criterion(\n",
    "                        outputs.view(-1, vocab_size), targets.reshape(-1)\n",
    "                    )\n",
    "                    total_val_loss += val_loss.item()\n",
    "\n",
    "                for output_batch in outputs.argmax(dim=-1):\n",
    "                    words_preds = [\n",
    "                        train_dataset.itos[idx.item()] for idx in output_batch\n",
    "                    ]\n",
    "                    hypotheses.append(words_preds)\n",
    "\n",
    "                for target_batch in targets:\n",
    "                    words_targets = [\n",
    "                        train_dataset.itos[idx.item()] for idx in target_batch\n",
    "                    ]\n",
    "                    references.append([words_targets])\n",
    "\n",
    "        avg_val_loss = total_val_loss / len(val_loader)\n",
    "        val_losses.append(avg_val_loss)\n",
    "        print(f\"Validation Loss after Epoch {epoch}: {avg_val_loss:.5f}\")\n",
    "\n",
    "        # Compute BLEU scores\n",
    "        bleu_1 = corpus_bleu(\n",
    "            references,\n",
    "            hypotheses,\n",
    "            weights=(1, 0, 0, 0),\n",
    "            smoothing_function=SmoothingFunction().method1,\n",
    "        )\n",
    "        bleu_2 = corpus_bleu(\n",
    "            references,\n",
    "            hypotheses,\n",
    "            weights=(0.5, 0.5, 0, 0),\n",
    "            smoothing_function=SmoothingFunction().method1,\n",
    "        )\n",
    "        bleu_3 = corpus_bleu(\n",
    "            references,\n",
    "            hypotheses,\n",
    "            weights=(0.33, 0.33, 0.33, 0),\n",
    "            smoothing_function=SmoothingFunction().method1,\n",
    "        )\n",
    "        bleu_4 = corpus_bleu(\n",
    "            references, hypotheses, smoothing_function=SmoothingFunction().method1\n",
    "        )\n",
    "        bleu_scores.append((bleu_1, bleu_2, bleu_3, bleu_4))\n",
    "        if avg_val_loss < best_val_loss:\n",
    "            best_val_loss = avg_val_loss\n",
    "\n",
    "        model.train() \n",
    "\n",
    "    end_time = time.time() \n",
    "    epoch_duration = end_time - start_time\n",
    "    print(f\"Time for epoch {epoch}: {epoch_duration:.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "bleu_1_scores = [score[0] for score in bleu_scores]\n",
    "bleu_2_scores = [score[1] for score in bleu_scores]\n",
    "bleu_3_scores = [score[2] for score in bleu_scores]\n",
    "bleu_4_scores = [score[3] for score in bleu_scores]\n",
    "\n",
    "epochs = range(1, len(bleu_scores) + 1)\n",
    "\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.plot(epochs, bleu_1_scores, label=\"BLEU-1\", marker=\"o\")\n",
    "plt.plot(epochs, bleu_2_scores, label=\"BLEU-2\", marker=\"o\")\n",
    "plt.plot(epochs, bleu_3_scores, label=\"BLEU-3\", marker=\"o\")\n",
    "plt.plot(epochs, bleu_4_scores, label=\"BLEU-4\", marker=\"o\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"BLEU Score\")\n",
    "plt.title(\"BLEU Scores Over Training Epochs\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(train_losses, label='Training Loss')\n",
    "plt.plot(range(0, len(train_losses), validate_every), val_losses, label='Validation Loss', linestyle='--')\n",
    "plt.title('Training and Validation Loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ece176",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}