run_dann.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/20], Step [10/22], Clf Loss: 0.4569, Domain Loss: 0.6016, Total Loss: 1.0585\n",
      "Epoch [1/20], Step [20/22], Clf Loss: 0.3445, Domain Loss: 0.3576, Total Loss: 0.7021\n",
      "Epoch [1/20], Train Accuracy: 81.68%\n",
      "Epoch [2/20], Step [10/22], Clf Loss: 0.2281, Domain Loss: 0.3660, Total Loss: 0.5940\n",
      "Epoch [2/20], Step [20/22], Clf Loss: 0.1543, Domain Loss: 0.2313, Total Loss: 0.3856\n",
      "Epoch [2/20], Train Accuracy: 96.24%\n",
      "Epoch [3/20], Step [10/22], Clf Loss: 0.1320, Domain Loss: 0.3486, Total Loss: 0.4807\n",
      "Epoch [3/20], Step [20/22], Clf Loss: 0.1062, Domain Loss: 0.3126, Total Loss: 0.4188\n",
      "Epoch [3/20], Train Accuracy: 98.93%\n",
      "Epoch [4/20], Step [10/22], Clf Loss: 0.0849, Domain Loss: 0.2209, Total Loss: 0.3058\n",
      "Epoch [4/20], Step [20/22], Clf Loss: 0.1027, Domain Loss: 0.1878, Total Loss: 0.2906\n",
      "Epoch [4/20], Train Accuracy: 99.72%\n",
      "Epoch [5/20], Step [10/22], Clf Loss: 0.0683, Domain Loss: 0.2021, Total Loss: 0.2704\n",
      "Epoch [5/20], Step [20/22], Clf Loss: 0.0796, Domain Loss: 0.1838, Total Loss: 0.2635\n",
      "Epoch [5/20], Train Accuracy: 99.72%\n",
      "Epoch [6/20], Step [10/22], Clf Loss: 0.0554, Domain Loss: 0.1437, Total Loss: 0.1991\n",
      "Epoch [6/20], Step [20/22], Clf Loss: 0.0475, Domain Loss: 0.2300, Total Loss: 0.2775\n",
      "Epoch [6/20], Train Accuracy: 99.86%\n",
      "Epoch [7/20], Step [10/22], Clf Loss: 0.0409, Domain Loss: 0.2415, Total Loss: 0.2824\n",
      "Epoch [7/20], Step [20/22], Clf Loss: 0.0382, Domain Loss: 0.0980, Total Loss: 0.1362\n",
      "Epoch [7/20], Train Accuracy: 99.93%\n",
      "Epoch [8/20], Step [10/22], Clf Loss: 0.0422, Domain Loss: 0.1273, Total Loss: 0.1696\n",
      "Epoch [8/20], Step [20/22], Clf Loss: 0.0335, Domain Loss: 0.1965, Total Loss: 0.2300\n",
      "Epoch [8/20], Train Accuracy: 99.93%\n",
      "Epoch [9/20], Step [10/22], Clf Loss: 0.0303, Domain Loss: 0.0496, Total Loss: 0.0799\n",
      "Epoch [9/20], Step [20/22], Clf Loss: 0.0309, Domain Loss: 0.0514, Total Loss: 0.0822\n",
      "Epoch [9/20], Train Accuracy: 100.00%\n",
      "Epoch [10/20], Step [10/22], Clf Loss: 0.0304, Domain Loss: 0.0836, Total Loss: 0.1140\n",
      "Epoch [10/20], Step [20/22], Clf Loss: 0.0249, Domain Loss: 0.0526, Total Loss: 0.0774\n",
      "Epoch [10/20], Train Accuracy: 99.93%\n",
      "Epoch [11/20], Step [10/22], Clf Loss: 0.0228, Domain Loss: 0.0690, Total Loss: 0.0918\n",
      "Epoch [11/20], Step [20/22], Clf Loss: 0.0208, Domain Loss: 0.0682, Total Loss: 0.0891\n",
      "Epoch [11/20], Train Accuracy: 100.00%\n",
      "Epoch [12/20], Step [10/22], Clf Loss: 0.0259, Domain Loss: 0.0227, Total Loss: 0.0486\n",
      "Epoch [12/20], Step [20/22], Clf Loss: 0.0185, Domain Loss: 0.0450, Total Loss: 0.0634\n",
      "Epoch [12/20], Train Accuracy: 100.00%\n",
      "Epoch [13/20], Step [10/22], Clf Loss: 0.0179, Domain Loss: 0.0264, Total Loss: 0.0443\n",
      "Epoch [13/20], Step [20/22], Clf Loss: 0.0224, Domain Loss: 0.0320, Total Loss: 0.0544\n",
      "Epoch [13/20], Train Accuracy: 100.00%\n",
      "Epoch [14/20], Step [10/22], Clf Loss: 0.0170, Domain Loss: 0.0395, Total Loss: 0.0566\n",
      "Epoch [14/20], Step [20/22], Clf Loss: 0.0214, Domain Loss: 0.0175, Total Loss: 0.0389\n",
      "Epoch [14/20], Train Accuracy: 100.00%\n",
      "Epoch [15/20], Step [10/22], Clf Loss: 0.0144, Domain Loss: 0.0286, Total Loss: 0.0429\n",
      "Epoch [15/20], Step [20/22], Clf Loss: 0.0143, Domain Loss: 0.0171, Total Loss: 0.0314\n",
      "Epoch [15/20], Train Accuracy: 99.93%\n",
      "Epoch [16/20], Step [10/22], Clf Loss: 0.0139, Domain Loss: 0.0345, Total Loss: 0.0485\n",
      "Epoch [16/20], Step [20/22], Clf Loss: 0.0141, Domain Loss: 0.0264, Total Loss: 0.0405\n",
      "Epoch [16/20], Train Accuracy: 100.00%\n",
      "Epoch [17/20], Step [10/22], Clf Loss: 0.0179, Domain Loss: 0.0145, Total Loss: 0.0324\n",
      "Epoch [17/20], Step [20/22], Clf Loss: 0.0145, Domain Loss: 0.0250, Total Loss: 0.0394\n",
      "Epoch [17/20], Train Accuracy: 100.00%\n",
      "Epoch [18/20], Step [10/22], Clf Loss: 0.0155, Domain Loss: 0.0259, Total Loss: 0.0414\n",
      "Epoch [18/20], Step [20/22], Clf Loss: 0.0159, Domain Loss: 0.0108, Total Loss: 0.0267\n",
      "Epoch [18/20], Train Accuracy: 100.00%\n",
      "Epoch [19/20], Step [10/22], Clf Loss: 0.0132, Domain Loss: 0.0090, Total Loss: 0.0222\n",
      "Epoch [19/20], Step [20/22], Clf Loss: 0.0112, Domain Loss: 0.0116, Total Loss: 0.0228\n",
      "Epoch [19/20], Train Accuracy: 100.00%\n",
      "Epoch [20/20], Step [10/22], Clf Loss: 0.0101, Domain Loss: 0.0065, Total Loss: 0.0166\n",
      "Epoch [20/20], Step [20/22], Clf Loss: 0.0183, Domain Loss: 0.0066, Total Loss: 0.0249\n",
      "Epoch [20/20], Train Accuracy: 100.00%\n",
      "Epoch [20/20], Test Accuracy: 88.54%\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms\n",
    "from dataset import ImageDataset, DomainDataset\n",
    "from torch.utils.data import DataLoader,random_split\n",
    "from dann import LeNetFeatureExtractor, LabelPredictor, DomainDiscriminator, DANN\n",
    "from itertools import zip_longest\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "config={\n",
    "    \"mode\": \"train\",\n",
    "    \"batch\": 64,\n",
    "    \"epoch\": 20,\n",
    "    \"lr\": 1e-4,\n",
    "    \"cuda\": 0,\n",
    "    \"norm\": 2,\n",
    "    \"norm_type\": \"BN\",\n",
    "    \"dropout\": False,\n",
    "    \"weight_decay\": False,\n",
    "    \"opt\": \"adam\",\n",
    "    \"activation\": \"leaky_relu\",\n",
    "    \"data_augmentation\":False,\n",
    "    \"save\":False,\n",
    "    \"dropout_rate\":0.7,\n",
    "    \"alpha\": 0.1,  # You can set your desired initial alpha value here.\n",
    "    \"domain_weight\": 0.7,\n",
    "    \"source_domain_label\": 0,\n",
    "}\n",
    "\n",
    "base_dir=\"dataset/\"\n",
    "device=config[\"cuda\"]\n",
    "device=torch.device(f\"cuda:{device}\")\n",
    "\n",
    "# Define the hyperparameters\n",
    "# batch_size = 128\n",
    "# lr = 0.001\n",
    "# num_epochs = 10\n",
    "batch_size = config[\"batch\"]\n",
    "lr = config[\"lr\"]\n",
    "num_epochs = config[\"epoch\"]\n",
    "lambda_val = 0.1  # domain adversarial loss weight\n",
    "ratio=0.2\n",
    "# Define the datasets and data loaders\n",
    "trainset = DomainDataset(base_dir + config[\"mode\"], device=device, config=config, domain=0, train=True)\n",
    "test_dataset = DomainDataset(base_dir + 'test', device=device, config=config, domain=1, train=False)\n",
    "target_domain_dataset, testdataset = random_split(test_dataset, [int(ratio * len(test_dataset)), len(test_dataset)-int(ratio * len(test_dataset))])\n",
    "trainloader = DataLoader(torch.utils.data.ConcatDataset([trainset, target_domain_dataset]), batch_size=config[\"batch\"], shuffle=True, num_workers=16, drop_last=True)\n",
    "testloader = DataLoader(testdataset, batch_size=config[\"batch\"], shuffle=False, num_workers=16, drop_last=True)\n",
    "\n",
    "# Define the DANN model and the optimizer\n",
    "feature_extractor = LeNetFeatureExtractor(config=config)\n",
    "label_predictor = LabelPredictor(num_classes=2,config=config)\n",
    "domain_discriminator = DomainDiscriminator(config=config)\n",
    "dann = DANN(feature_extractor, label_predictor, domain_discriminator)\n",
    "\n",
    "optimizer = optim.Adam(dann.parameters(), lr=lr)\n",
    "\n",
    "# Define the loss functions\n",
    "clf_loss_fn = nn.CrossEntropyLoss()\n",
    "domain_loss_fn = nn.BCELoss()\n",
    "\n",
    "# Train the DANN model\n",
    "for epoch in range(num_epochs):\n",
    "    dann.train()\n",
    "    num_correct_train = 0\n",
    "    num_total_train = 0\n",
    "    for i, (inputs, labels, domains) in enumerate(trainloader):\n",
    "        # Set the domain labels (0 for source, 1 for target)\n",
    "        # source_domain_labels = torch.zeros(inputs.size(0))\n",
    "        # target_domain_labels = torch.ones(inputs.size(0))\n",
    "        # domain_labels = torch.cat((source_domain_labels, target_domain_labels)).unsqueeze(1)\n",
    "        domain_labels = domains.unsqueeze(1).float()\n",
    "        # print(\"domain_labels: \",domain_labels.shape,type(domain_labels))\n",
    "\n",
    "        # Zero the gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Extract features and predict labels\n",
    "        features = dann.feature_extractor(inputs)\n",
    "        label_preds = dann.label_predictor(features)\n",
    "\n",
    "        # Count Acc\n",
    "        preds = F.softmax(label_preds, dim=1)\n",
    "        pred_labels = torch.argmax(preds, dim=1)\n",
    "        # print(pred_labels.shape,pred_labels,labels, labels.shape)\n",
    "        num_correct_train += (pred_labels == labels).sum().item()\n",
    "        num_total_train += labels.size(0)\n",
    "\n",
    "        # Compute the label prediction loss\n",
    "        clf_loss = clf_loss_fn(label_preds, labels)\n",
    "\n",
    "        # Compute the domain classification loss\n",
    "        domain_preds = dann.domain_discriminator(features)\n",
    "        domain_loss = domain_loss_fn(domain_preds, domain_labels)\n",
    "\n",
    "        # Compute the total loss and update the parameters\n",
    "        total_loss = clf_loss + domain_loss\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Print the training statistics\n",
    "        if (i+1) % 10 == 0:\n",
    "            print(\"Epoch [{}/{}], Step [{}/{}], Clf Loss: {:.4f}, Domain Loss: {:.4f}, Total Loss: {:.4f}\"\n",
    "                  .format(epoch+1, num_epochs, i+1, len(trainloader), clf_loss.item(), domain_loss.item(), total_loss.item()))\n",
    "        \n",
    "    print(\"Epoch [{}/{}], Train Accuracy: {:.2f}%\".format(epoch+1, num_epochs, 100 * num_correct_train / num_total_train))\n",
    "\n",
    "# Evaluate the model on the test set\n",
    "dann.eval()\n",
    "with torch.no_grad():\n",
    "    num_correct = 0\n",
    "    num_total = 0\n",
    "    for i, (inputs, labels, domains) in enumerate(testloader):\n",
    "        features = dann.feature_extractor(inputs)\n",
    "        logits = dann.label_predictor(features)\n",
    "        preds = F.softmax(logits, dim=1)\n",
    "        pred_labels = torch.argmax(preds, dim=1)\n",
    "        # print(pred_labels.shape,pred_labels,labels, labels.shape)\n",
    "        num_correct += (pred_labels == labels).sum().item()\n",
    "        num_total += labels.size(0)\n",
    "    accuracy = 100 * num_correct / num_total\n",
    "    print(\"Epoch [{}/{}], Test Accuracy: {:.2f}%\".format(epoch+1, num_epochs, accuracy))\n",
    "\n",
    "    # Set the model back to training mode\n",
    "    dann.train()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}