convexNet-based/v5/half_training_after_first_model-with-optimized-train-loop.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2024-11-10T08:15:54.955903Z",
     "iopub.status.busy": "2024-11-10T08:15:54.955496Z",
     "iopub.status.idle": "2024-11-10T08:16:14.690228Z",
     "shell.execute_reply": "2024-11-10T08:16:14.689269Z",
     "shell.execute_reply.started": "2024-11-10T08:15:54.955862Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, models\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os\n",
    "from PIL import Image\n",
    "from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "import cv2\n",
    "from transformers import AutoImageProcessor,ConvNextModel\n",
    "from tqdm import tqdm \n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from itertools import cycle\n",
    "import math\n",
    "from tqdm import tqdm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:14.693069Z",
     "iopub.status.busy": "2024-11-10T08:16:14.692037Z",
     "iopub.status.idle": "2024-11-10T08:16:14.699425Z",
     "shell.execute_reply": "2024-11-10T08:16:14.698544Z",
     "shell.execute_reply.started": "2024-11-10T08:16:14.693020Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# pre-defined\n",
    "category_dict = {'c1':'Men Tshirts', 'c2':'Sarees', 'c3': 'Kurtis', 'c4': 'Women Tshirts', 'c5': 'Women Tops & Tunics'}\n",
    "semi_classes_dict = {'c1':[4, 2, 2, 3, 2], 'c2':[4, 6, 3, 8, 4, 3, 4, 5, 9, 2], 'c3': [13, 2, 2, 2, 2, 2, 2, 3, 2], 'c4': [7, 3, 3, 3, 6, 3, 2, 2], 'c5': [12, 4, 2, 7, 2, 3, 6, 4, 4, 6]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:14.700931Z",
     "iopub.status.busy": "2024-11-10T08:16:14.700483Z",
     "iopub.status.idle": "2024-11-10T08:16:14.738717Z",
     "shell.execute_reply": "2024-11-10T08:16:14.737832Z",
     "shell.execute_reply.started": "2024-11-10T08:16:14.700898Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# set paths as per the set-up and Hyperparameters\n",
    "input_path = \"/kaggle/input/meesho\"\n",
    "working_path = \"/kaggle/working\"\n",
    "test_c_name = \"c4\"\n",
    "\n",
    "BEST_MODEL_FROM_BASE_FIRST_TRAINING = f'{working_path}/best_model_{model_type}.pth'\n",
    "NUM_EPOCH = 4\n",
    "NUM_ATTR_EPOCHS = 1\n",
    "LEARNING_RATE = 0.001\n",
    "NUM_OF_SEMI_CLASSES_OF_COLUMNS = semi_classes_dict[test_c_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:14.740914Z",
     "iopub.status.busy": "2024-11-10T08:16:14.740572Z",
     "iopub.status.idle": "2024-11-10T08:16:14.976438Z",
     "shell.execute_reply": "2024-11-10T08:16:14.975645Z",
     "shell.execute_reply.started": "2024-11-10T08:16:14.740871Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "df_train = pd.read_csv(f'{input_path}/train.csv')\n",
    "df_test = pd.read_csv(f'{input_path}/test.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:14.977886Z",
     "iopub.status.busy": "2024-11-10T08:16:14.977543Z",
     "iopub.status.idle": "2024-11-10T08:16:14.996194Z",
     "shell.execute_reply": "2024-11-10T08:16:14.995414Z",
     "shell.execute_reply.started": "2024-11-10T08:16:14.977850Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "df_train['Category'].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:14.997866Z",
     "iopub.status.busy": "2024-11-10T08:16:14.997415Z",
     "iopub.status.idle": "2024-11-10T08:16:15.006210Z",
     "shell.execute_reply": "2024-11-10T08:16:15.005373Z",
     "shell.execute_reply.started": "2024-11-10T08:16:14.997823Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def do_preprocessing(test_c_name, test_category):\n",
    "    df_temp = df_train[df_train['Category'] == test_category]\n",
    "    df_sub = df_train[df_train['Category'] == test_category]\n",
    "    df_sub.dropna(axis=1, how='all', inplace=True)\n",
    "    temp = []\n",
    "    for i in range(1,len(df_sub.columns)-2):\n",
    "        temp.append(len(df_sub[f'attr_{i}'].unique().tolist())-1)\n",
    "    print(f\"Number of features for original df: {temp}\")\n",
    "    print()\n",
    "    df_sub.dropna(axis=0, how='any', inplace=True)\n",
    "    temp = []\n",
    "    for i in range(1,len(df_sub.columns)-2):\n",
    "        temp.append(len(df_sub[f'attr_{i}'].unique().tolist()))\n",
    "    print(f\"Number of features for new df: {temp}\")\n",
    "    return df_sub, df_temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.008455Z",
     "iopub.status.busy": "2024-11-10T08:16:15.008101Z",
     "iopub.status.idle": "2024-11-10T08:16:15.031485Z",
     "shell.execute_reply": "2024-11-10T08:16:15.030645Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.008413Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class LabelEncoderDict:\n",
    "    def __init__(self):\n",
    "        self.encoders = {}\n",
    "        \n",
    "    def fit(self, df, columns):\n",
    "        \"\"\"Fit label encoders for each column\"\"\"\n",
    "        for col in columns:\n",
    "            le = LabelEncoder()\n",
    "            valid_labels = df[col].unique().tolist()\n",
    "            valid_labels = [x for x in valid_labels if not (isinstance(x, float) and math.isnan(x))]\n",
    "            le.fit(valid_labels)\n",
    "            self.encoders[col] = le\n",
    "            \n",
    "    def transform(self, df, columns):\n",
    "        \"\"\"Transform labels using fitted encoders\"\"\"\n",
    "        encoded = np.zeros((len(df), len(columns)))\n",
    "        for i, col in enumerate(columns):\n",
    "            series = df[col].copy()\n",
    "            encoded[:, i] = self.encoders[col].transform(series)\n",
    "        return encoded\n",
    "    \n",
    "    def get_num_classes(self, column):\n",
    "        \"\"\"Get number of classes for a specific column\"\"\"\n",
    "        return len(self.encoders[column].classes_)\n",
    "\n",
    "\n",
    "class MultiLabelImageDataset(Dataset):\n",
    "    def __init__(self, df, image_dir, transform_basic=None, transform_augmented=None, attr_columns=10, do_transform=True):\n",
    "        self.df = df\n",
    "        self.image_dir = image_dir\n",
    "        self.transform_basic = transform_basic  # Basic transform without augmentation\n",
    "        self.transform_augmented = transform_augmented  # Augmented transform with augmentation\n",
    "        self.attr_columns = attr_columns\n",
    "        self.do_transform = do_transform\n",
    "        \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        # Get image path\n",
    "        img_name = str(self.df.iloc[idx]['id']).zfill(6)\n",
    "        img_path = os.path.join(self.image_dir, f\"{img_name}.jpg\")\n",
    "\n",
    "        # Load image\n",
    "        try:\n",
    "            image = Image.open(img_path).convert('RGB')\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading image {img_path}: {e}\")\n",
    "            image = Image.new('RGB', (512, 512))\n",
    "\n",
    "\n",
    "        if  self.do_transform and (random.random() > 0.5):\n",
    "            if self.transform_augmented:\n",
    "                image = self.transform_augmented(image)         \n",
    "        else:\n",
    "            if self.transform_basic:\n",
    "                image = self.transform_basic(image)\n",
    "        # Ensure labels are integers and convert to tensor\n",
    "        labels = torch.tensor(self.df.iloc[idx][self.attr_columns].astype(int).values, dtype=torch.long)\n",
    "\n",
    "        return image, labels\n",
    "\n",
    "\n",
    "\n",
    "def prepare_data(df,label_encoders, image_dir, batch_size=32, test_size=0.1, num_attr_columns=10):\n",
    "    \"\"\"\n",
    "    Prepare data loaders and label encoders\n",
    "    \"\"\"\n",
    "      # TODO: Adjust number of columns\n",
    "    \n",
    "    # Transform labels\n",
    "    attr_columns = [f'attr_{i}' for i in range(1, num_attr_columns+1)]\n",
    "    encoded_labels = label_encoders.transform(df, attr_columns)\n",
    "    df_encoded = df.copy()\n",
    "    for i, col in enumerate(attr_columns):\n",
    "        df_encoded[col] = encoded_labels[:, i]\n",
    "    \n",
    "    # Split data\n",
    "    train_df, val_df = train_test_split(df_encoded, test_size=test_size, random_state=24)\n",
    "    \n",
    "    # Define transforms\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((512, 512)),  # Resize to 512x512\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                             std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    \n",
    "    transform_augmented = transforms.Compose([\n",
    "        transforms.Resize((512, 512)),  # Resize to 512x512\n",
    "        transforms.RandomHorizontalFlip(p=0.9),  # 90% chance of horizontal flipping\n",
    "        transforms.RandomRotation(degrees=5),  # Rotate by up to 20 degrees\n",
    "        transforms.RandomResizedCrop(size=(512, 512), scale=(0.8, 1.0)),  # Randomly crop and resize\n",
    "        transforms.RandomPerspective(distortion_scale=0.1, p=0.5),  # Apply perspective distortion\n",
    "        transforms.RandomAffine(degrees=0, translate=(0.15, 0.15), scale=(0.85, 1.15), shear=5),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "\n",
    "    # Create datasets\n",
    "    train_dataset = MultiLabelImageDataset(\n",
    "        train_df,\n",
    "        image_dir,\n",
    "        transform_basic=transform,\n",
    "        transform_augmented=transform_augmented,\n",
    "        attr_columns=attr_columns\n",
    "    )\n",
    "    \n",
    "    val_dataset = MultiLabelImageDataset(\n",
    "        val_df,\n",
    "        image_dir,\n",
    "        transform_basic=transform,\n",
    "        transform_augmented=transform_augmented,\n",
    "        attr_columns=attr_columns,\n",
    "        do_transform=False\n",
    "    )\n",
    "    \n",
    "    # Create dataloaders\n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=batch_size, drop_last=True)\n",
    "    \n",
    "    # Get number of classes for each attribute\n",
    "    num_classes_per_attr = [label_encoders.get_num_classes(col) for col in attr_columns]\n",
    "    \n",
    "    return train_loader, val_loader, label_encoders, num_classes_per_attr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.033447Z",
     "iopub.status.busy": "2024-11-10T08:16:15.033047Z",
     "iopub.status.idle": "2024-11-10T08:16:15.054740Z",
     "shell.execute_reply": "2024-11-10T08:16:15.053797Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.033405Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelClassifier(nn.Module):\n",
    "    def __init__(self, num_classes_per_attr):\n",
    "        super(MultiLabelClassifier, self).__init__()\n",
    "        \n",
    "        # Use ConvNeXt-Base with unfrozen backbone\n",
    "        self.backbone = ConvNextModel.from_pretrained(\"facebook/convnext-base-384-22k-1k\")\n",
    "        backbone_features = self.backbone.config.hidden_sizes[-1]  # 1024 for base model\n",
    "        \n",
    "        # Modified feature processing without fixed dimensions\n",
    "        self.feature_processor = nn.Sequential(\n",
    "            nn.Conv2d(backbone_features, 1024, kernel_size=1),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(0.1)\n",
    "        )\n",
    "        \n",
    "        # Keep original ModuleList structure\n",
    "        self.classifier_heads = nn.ModuleList()\n",
    "        for num_classes in num_classes_per_attr:\n",
    "            classifier_head = nn.Sequential(\n",
    "                # First branch - Spatial attention\n",
    "                nn.Sequential(\n",
    "                    nn.Conv2d(1024, 512, kernel_size=3, padding=1, groups=32),\n",
    "                    nn.GELU(),\n",
    "                    nn.Conv2d(512, 512, kernel_size=3, padding=1, groups=32),\n",
    "                    nn.GELU(),\n",
    "                ),\n",
    "                \n",
    "                # Second branch - Channel attention (SE-like module)\n",
    "                nn.Sequential(\n",
    "                    nn.AdaptiveAvgPool2d(1),\n",
    "                    nn.Flatten(),\n",
    "                    nn.Linear(512, 128),\n",
    "                    nn.GELU(),\n",
    "                    nn.Linear(128, 512),\n",
    "                    nn.Sigmoid(),\n",
    "                ),\n",
    "                \n",
    "                # Combine branches and final classification\n",
    "                nn.Sequential(\n",
    "                    nn.AdaptiveAvgPool2d(1),\n",
    "                    nn.Flatten(),\n",
    "                    nn.Linear(512, 1024),\n",
    "                    nn.LayerNorm(1024),\n",
    "                    nn.GELU(),\n",
    "                    nn.Dropout(0.2),\n",
    "                    nn.Linear(1024, 512),\n",
    "                    nn.LayerNorm(512),\n",
    "                    nn.Sigmoid(),\n",
    "                    nn.Dropout(0.1),\n",
    "                    nn.Linear(512, num_classes)\n",
    "                )\n",
    "            )\n",
    "            self.classifier_heads.append(classifier_head)\n",
    "\n",
    "    def freeze_backbone(self):\n",
    "        \"\"\"Freeze the backbone model\"\"\"\n",
    "        for param in self.backbone.parameters():\n",
    "            param.requires_grad = False\n",
    "            \n",
    "    def unfreeze_backbone(self):\n",
    "        \"\"\"Unfreeze the backbone model\"\"\"\n",
    "        for param in self.backbone.parameters():\n",
    "            param.requires_grad = True\n",
    "            \n",
    "    def freeze_feature_processor(self):\n",
    "        \"\"\"Freeze the feature processor\"\"\"\n",
    "        for param in self.feature_processor.parameters():\n",
    "            param.requires_grad = False\n",
    "            \n",
    "    def unfreeze_feature_processor(self):\n",
    "        \"\"\"Unfreeze the feature processor\"\"\"\n",
    "        for param in self.feature_processor.parameters():\n",
    "            param.requires_grad = True\n",
    "\n",
    "    def set_classifier_head_trainable(self, attr_index):\n",
    "        \"\"\"\n",
    "        Freeze all classifier heads except the specified one\n",
    "        Args:\n",
    "            attr_index: index of the attribute head to train (0 for attr_1, 1 for attr_2, etc.)\n",
    "        \"\"\"\n",
    "        for i, head in enumerate(self.classifier_heads):\n",
    "            for param in head.parameters():\n",
    "                param.requires_grad = (i == attr_index)\n",
    "\n",
    "    def freeze_all_except_head(self, attr_index):\n",
    "        \"\"\"\n",
    "        Freeze everything except the specified classifier head\n",
    "        Args:\n",
    "            attr_index: index of the attribute head to train (0 for attr_1, 1 for attr_2, etc.)\n",
    "        \"\"\"\n",
    "        self.freeze_backbone()\n",
    "        self.freeze_feature_processor()\n",
    "        self.set_classifier_head_trainable(attr_index)\n",
    "        \n",
    "    def unfreeze_all(self):\n",
    "        \"\"\"Unfreeze all model components\"\"\"\n",
    "        self.unfreeze_backbone()\n",
    "        self.unfreeze_feature_processor()\n",
    "        for head in self.classifier_heads:\n",
    "            for param in head.parameters():\n",
    "                param.requires_grad = True\n",
    "\n",
    "    def forward(self, x, attr_index=None, return_features=False):\n",
    "        \"\"\"\n",
    "        Forward pass with optional attribute-specific output\n",
    "        Args:\n",
    "            x: input tensor\n",
    "            attr_index: specific attribute index to get output for (0 for attr_1, etc.)\n",
    "            return_features: whether to return processed features\n",
    "        \"\"\"\n",
    "        # Extract features from ConvNeXt backbone\n",
    "        features = self.backbone(x).last_hidden_state\n",
    "        \n",
    "        # Process features\n",
    "        processed_features = self.feature_processor(features)\n",
    "        \n",
    "        if attr_index is not None:\n",
    "            # Get output for specific attribute\n",
    "            classifier_head = self.classifier_heads[attr_index]\n",
    "            # Spatial attention branch\n",
    "            spatial_features = classifier_head[0](processed_features)\n",
    "            \n",
    "            # Channel attention branch\n",
    "            channel_attention = classifier_head[1](spatial_features)\n",
    "            channel_attention = channel_attention.view(-1, 512, 1, 1)\n",
    "            \n",
    "            # Apply channel attention and get final output\n",
    "            attended_features = spatial_features * channel_attention\n",
    "            output = classifier_head[2](attended_features)\n",
    "            \n",
    "            if return_features:\n",
    "                return output, processed_features\n",
    "            return output\n",
    "        \n",
    "        # Get outputs for all attributes\n",
    "        outputs = []\n",
    "        for classifier_head in self.classifier_heads:\n",
    "            # Spatial attention branch\n",
    "            spatial_features = classifier_head[0](processed_features)\n",
    "            \n",
    "            # Channel attention branch\n",
    "            channel_attention = classifier_head[1](spatial_features)\n",
    "            channel_attention = channel_attention.view(-1, 512, 1, 1)\n",
    "            \n",
    "            # Apply channel attention and get final output\n",
    "            attended_features = spatial_features * channel_attention\n",
    "            outputs.append(classifier_head[2](attended_features))\n",
    "            \n",
    "        if return_features:\n",
    "            return outputs, processed_features\n",
    "        return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.056123Z",
     "iopub.status.busy": "2024-11-10T08:16:15.055849Z",
     "iopub.status.idle": "2024-11-10T08:16:15.102107Z",
     "shell.execute_reply": "2024-11-10T08:16:15.101155Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.056093Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class MultiLabelCELoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiLabelCELoss, self).__init__()\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, outputs, targets):\n",
    "        loss = 0\n",
    "        for i, output in enumerate(outputs):\n",
    "            loss += self.criterion(output, targets[:, i])\n",
    "        return loss / len(outputs)\n",
    "    \n",
    "class CombinedLoss(nn.Module):\n",
    "    def __init__(self, lambda_mmd=0.1, chunk_size=1024):\n",
    "        super(CombinedLoss, self).__init__()\n",
    "        self.lambda_mmd = lambda_mmd\n",
    "        self.chunk_size = chunk_size\n",
    "        self.cross_entropy_loss = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, outputs, labels, source_features, target_features):\n",
    "        # Calculate Cross-Entropy Loss\n",
    "        if isinstance(outputs, list):\n",
    "            ce_loss = 0\n",
    "            for i, output in enumerate(outputs):\n",
    "                if isinstance(labels, list):\n",
    "                    label = labels[i]\n",
    "                else:\n",
    "                    label = labels[:, i] if labels.dim() > 1 else labels\n",
    "                ce_loss += self.cross_entropy_loss(output, label)\n",
    "            ce_loss = ce_loss / len(outputs)\n",
    "        else:\n",
    "            ce_loss = self.cross_entropy_loss(outputs, labels)\n",
    "\n",
    "        # Calculate MMD Loss\n",
    "        mmd_loss = self.maximum_mean_discrepancy(source_features, target_features)\n",
    "        \n",
    "        # Combine losses (without CORAL)\n",
    "        total_loss = ce_loss + self.lambda_mmd * mmd_loss\n",
    "        \n",
    "        # Return zero for coral_loss to maintain compatibility\n",
    "        return total_loss, ce_loss, mmd_loss, torch.tensor(0.0, device=ce_loss.device)\n",
    "\n",
    "    def gaussian_kernel(self, x, y, bandwidth):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        y = y.view(y.size(0), -1)\n",
    "        \n",
    "        x_size = x.size(0)\n",
    "        y_size = y.size(0)\n",
    "        dim = x.size(1)\n",
    "        \n",
    "        x = x.unsqueeze(1)  # (x_size, 1, dim)\n",
    "        y = y.unsqueeze(0)  # (1, y_size, dim)\n",
    "        \n",
    "        kernel_input = (x - y).pow(2).sum(2).div(2 * bandwidth * bandwidth)\n",
    "        return torch.exp(-kernel_input)  # (x_size, y_size)\n",
    "\n",
    "    def maximum_mean_discrepancy(self, source_features, target_features):\n",
    "        # Ensure inputs are 2D tensors\n",
    "        if source_features.dim() > 2:\n",
    "            source_features = source_features.view(source_features.size(0), -1)\n",
    "        if target_features.dim() > 2:\n",
    "            target_features = target_features.view(target_features.size(0), -1)\n",
    "\n",
    "        # Get sizes\n",
    "        batch_source = source_features.size(0)\n",
    "        batch_target = target_features.size(0)\n",
    "        dim = source_features.size(1)\n",
    "\n",
    "        # Initialize MMD\n",
    "        mmd = torch.tensor(0., device=source_features.device)\n",
    "\n",
    "        # Use multiple kernel bandwidths\n",
    "        bandwidths = [dim * (2 ** i) for i in range(-3, 3)]\n",
    "\n",
    "        for bandwidth in bandwidths:\n",
    "            # Process source-source\n",
    "            source_sum = 0\n",
    "            for i in range(0, batch_source, self.chunk_size):\n",
    "                end = min(i + self.chunk_size, batch_source)\n",
    "                chunk = source_features[i:end]\n",
    "                kernel = self.gaussian_kernel(chunk, source_features, bandwidth)\n",
    "                source_sum += kernel.sum().item()\n",
    "\n",
    "            # Process target-target\n",
    "            target_sum = 0\n",
    "            for i in range(0, batch_target, self.chunk_size):\n",
    "                end = min(i + self.chunk_size, batch_target)\n",
    "                chunk = target_features[i:end]\n",
    "                kernel = self.gaussian_kernel(chunk, target_features, bandwidth)\n",
    "                target_sum += kernel.sum().item()\n",
    "\n",
    "            # Process source-target\n",
    "            cross_sum = 0\n",
    "            for i in range(0, batch_source, self.chunk_size):\n",
    "                s_end = min(i + self.chunk_size, batch_source)\n",
    "                s_chunk = source_features[i:s_end]\n",
    "                \n",
    "                for j in range(0, batch_target, self.chunk_size):\n",
    "                    t_end = min(j + self.chunk_size, batch_target)\n",
    "                    t_chunk = target_features[j:t_end]\n",
    "                    \n",
    "                    kernel = self.gaussian_kernel(s_chunk, t_chunk, bandwidth)\n",
    "                    cross_sum += kernel.sum().item()\n",
    "\n",
    "            # Calculate bandwidth contribution to MMD\n",
    "            source_term = source_sum / (batch_source * batch_source)\n",
    "            target_term = target_sum / (batch_target * batch_target)\n",
    "            cross_term = 2 * cross_sum / (batch_source * batch_target)\n",
    "            \n",
    "            mmd = mmd + torch.tensor(source_term + target_term - cross_term, \n",
    "                                   device=source_features.device)\n",
    "\n",
    "            # Clear cache\n",
    "            if hasattr(torch.cuda, 'empty_cache'):\n",
    "                torch.cuda.empty_cache()\n",
    "\n",
    "        return mmd / len(bandwidths)\n",
    "\n",
    "\n",
    "\n",
    "def train_model(model, train_loader, val_loader, num_epochs, num_classes_per_attr, model_type):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model = model.to(device)\n",
    "    model = torch.nn.DataParallel(model)\n",
    "    criterion = CombinedLoss()\n",
    "    criterion2 = MultiLabelCELoss()\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)\n",
    "\n",
    "    # Reduce learning rate on plateau\n",
    "    lr_scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=2, verbose=True)\n",
    "\n",
    "    # Early stopping params\n",
    "    early_stopping_patience = 4\n",
    "    early_stopping_counter = 0\n",
    "    best_val_overall_acc = 0.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        train_loss = 0\n",
    "        correct_predictions = [0] * len(num_classes_per_attr)\n",
    "        total_predictions = 0\n",
    "        overall_correct = 0\n",
    "\n",
    "        # Store true and predicted labels for metrics calculation\n",
    "        train_true_labels = [[] for _ in range(len(num_classes_per_attr))]\n",
    "        train_predicted_labels = [[] for _ in range(len(num_classes_per_attr))]\n",
    "\n",
    "        # Create cyclic iterator for validation data during training\n",
    "        val_cycle = cycle(val_loader)\n",
    "\n",
    "        with tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} [Training]\", unit=\"batch\") as t:\n",
    "            for images, labels in t:\n",
    "                # Get target domain batch\n",
    "                target_images, _ = next(val_cycle)\n",
    "                \n",
    "                # Move data to device\n",
    "                images = images.to(device)\n",
    "                labels = labels.to(device)\n",
    "                target_images = target_images.to(device)\n",
    "\n",
    "                # Forward pass on source domain (training data)\n",
    "                outputs, source_features = model(images, return_features=True)\n",
    "                \n",
    "                # Forward pass on target domain (validation data)\n",
    "                _, target_features = model(target_images, return_features=True)\n",
    "\n",
    "                # Calculate combined loss\n",
    "                loss, ce_loss, mmd_loss, coral_loss = criterion(outputs, labels, source_features, target_features)\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                optimizer.step()\n",
    "\n",
    "                train_loss += loss.item()\n",
    "                t.set_postfix(loss=loss.item())\n",
    "\n",
    "                all_labels_match = torch.ones(labels.size(0), dtype=torch.bool, device=device)\n",
    "                for i, output in enumerate(outputs):\n",
    "                    _, predicted = torch.max(output, 1)\n",
    "                    correct_predictions[i] += (predicted == labels[:, i]).sum().item()\n",
    "                    all_labels_match &= (predicted == labels[:, i])\n",
    "\n",
    "                    train_true_labels[i].extend(labels[:, i].cpu().numpy())\n",
    "                    train_predicted_labels[i].extend(predicted.cpu().numpy())\n",
    "\n",
    "                overall_correct += all_labels_match.sum().item()\n",
    "                total_predictions += labels.size(0)\n",
    "\n",
    "        # Validation phase\n",
    "        model.eval()\n",
    "        val_loss = 0\n",
    "        val_correct_predictions = [0] * len(num_classes_per_attr)\n",
    "        val_total_predictions = 0\n",
    "        val_overall_correct = 0\n",
    "\n",
    "        val_true_labels = [[] for _ in range(len(num_classes_per_attr))]\n",
    "        val_predicted_labels = [[] for _ in range(len(num_classes_per_attr))]\n",
    "\n",
    "        with tqdm(val_loader, desc=f\"Epoch {epoch+1}/{num_epochs} [Validation]\", unit=\"batch\") as v:\n",
    "            with torch.no_grad():\n",
    "                for images, labels in v:\n",
    "                    images, labels = images.to(device), labels.to(device)\n",
    "                    outputs = model(images,return_features=False)\n",
    "                    loss = criterion2(outputs, labels)\n",
    "                    val_loss += loss.item()\n",
    "                    v.set_postfix(loss=loss.item())\n",
    "\n",
    "                    all_labels_match_val = torch.ones(labels.size(0), dtype=torch.bool, device=device)\n",
    "                    for i, output in enumerate(outputs):\n",
    "                        _, predicted = torch.max(output, 1)\n",
    "                        val_correct_predictions[i] += (predicted == labels[:, i]).sum().item()\n",
    "                        all_labels_match_val &= (predicted == labels[:, i])\n",
    "\n",
    "                        # Store validation labels for precision, recall, f1-score calculations\n",
    "                        val_true_labels[i].extend(labels[:, i].cpu().numpy())\n",
    "                        val_predicted_labels[i].extend(predicted.cpu().numpy())\n",
    "\n",
    "                    val_overall_correct += all_labels_match_val.sum().item()\n",
    "                    val_total_predictions += labels.size(0)\n",
    "                    \n",
    "        print(f'Epoch {epoch+1}/{num_epochs}')\n",
    "        print(f'Training Loss: {train_loss/len(train_loader):.4f}')\n",
    "        print(f'Validation Loss: {val_loss/len(val_loader):.4f}')\n",
    "        \n",
    "        for i in range(len(num_classes_per_attr)):\n",
    "            train_acc = 100 * correct_predictions[i] / total_predictions\n",
    "            val_acc = 100 * val_correct_predictions[i] / val_total_predictions\n",
    "\n",
    "            train_precision = precision_score(train_true_labels[i], train_predicted_labels[i], average='weighted')\n",
    "            train_recall = recall_score(train_true_labels[i], train_predicted_labels[i], average='weighted')\n",
    "            train_f1 = f1_score(train_true_labels[i], train_predicted_labels[i], average='weighted')\n",
    "\n",
    "            val_precision = precision_score(val_true_labels[i], val_predicted_labels[i], average='weighted')\n",
    "            val_recall = recall_score(val_true_labels[i], val_predicted_labels[i], average='weighted')\n",
    "            val_f1 = f1_score(val_true_labels[i], val_predicted_labels[i], average='weighted')\n",
    "\n",
    "            print(f'Attribute {i+1} - Train Acc: {train_acc:.2f}%, Val Acc: {val_acc:.2f}%')\n",
    "            print(f'Attribute {i+1} - Train Precision: {train_precision:.2f}, Train Recall: {train_recall:.2f}, Train F1-Score: {train_f1:.2f}')\n",
    "            print(f'Attribute {i+1} - Val Precision: {val_precision:.2f}, Val Recall: {val_recall:.2f}, Val F1-Score: {val_f1:.2f}')\n",
    "            print()\n",
    "\n",
    "        overall_train_acc = 100 * overall_correct / total_predictions\n",
    "        overall_val_acc = 100 * val_overall_correct / val_total_predictions\n",
    "        print(f'Overall Train Accuracy: {overall_train_acc:.2f}%')\n",
    "        print(f'Overall Validation Accuracy: {overall_val_acc:.2f}%')\n",
    "\n",
    "        # Early stopping logic\n",
    "        if overall_val_acc >= best_val_overall_acc:\n",
    "            best_val_overall_acc = overall_val_acc\n",
    "            torch.save(model.module.state_dict(), f'best_model_{model_type}.pth')\n",
    "            early_stopping_counter = 0\n",
    "        else:\n",
    "            early_stopping_counter += 1\n",
    "        \n",
    "        lr_scheduler.step(overall_val_acc)\n",
    "        \n",
    "        if early_stopping_counter >= early_stopping_patience:\n",
    "            print(\"Early stopping triggered\")\n",
    "    torch.save(model.module.state_dict(), f'best_model_end_{model_type}.pth')\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.105689Z",
     "iopub.status.busy": "2024-11-10T08:16:15.105216Z",
     "iopub.status.idle": "2024-11-10T08:16:15.126408Z",
     "shell.execute_reply": "2024-11-10T08:16:15.125691Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.105650Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class SingleAttributeDataset(Dataset):\n",
    "    def __init__(self, df, image_dir, attribute, transform_basic=None, transform_augmented=None, do_transform=True):\n",
    "        self.df = df[df[attribute].notna()].reset_index(drop=True)\n",
    "        self.image_dir = image_dir\n",
    "        self.transform_basic = transform_basic\n",
    "        self.transform_augmented = transform_augmented\n",
    "        self.attribute = attribute\n",
    "        self.do_transform = do_transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # Get image path\n",
    "        img_name = str(self.df.iloc[idx]['id']).zfill(6)\n",
    "        img_path = os.path.join(self.image_dir, f\"{img_name}.jpg\")\n",
    "\n",
    "        # Load image\n",
    "        try:\n",
    "            image = Image.open(img_path).convert('RGB')\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading image {img_path}: {e}\")\n",
    "            image = Image.new('RGB', (512, 512))\n",
    "\n",
    "        # Apply transforms\n",
    "        if self.do_transform and (random.random() > 0.5):\n",
    "            if self.transform_augmented:\n",
    "                image = self.transform_augmented(image)\n",
    "        else:\n",
    "            if self.transform_basic:\n",
    "                image = self.transform_basic(image)\n",
    "\n",
    "        # Get label for this attribute\n",
    "        label = torch.tensor(self.df.iloc[idx][self.attribute], dtype=torch.long)\n",
    "\n",
    "        return image, label\n",
    "\n",
    "def prepare_attribute_data(df, image_dir,label_encoders, batch_size=32, num_attr_columns=10, test_size=0.1):\n",
    "    \"\"\"\n",
    "    Prepare separate data loaders for each attribute\n",
    "    \"\"\"\n",
    "    # Define attribute columns\n",
    "    attr_columns = [f'attr_{i}' for i in range(1, num_attr_columns+1)]\n",
    "    \n",
    "    # Transform labels for each attribute\n",
    "    df_encoded = df.copy()\n",
    "    for col in attr_columns:\n",
    "        # Only encode non-null values\n",
    "        mask = df_encoded[col].notna()\n",
    "        df_encoded.loc[mask, col] = label_encoders.encoders[col].transform(df_encoded.loc[mask, col])\n",
    "    \n",
    "    # Define transforms\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((512, 512)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                           std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    \n",
    "    transform_augmented = transforms.Compose([\n",
    "        transforms.Resize((512, 512)),\n",
    "        transforms.RandomHorizontalFlip(p=0.9),\n",
    "        transforms.RandomRotation(degrees=5),\n",
    "        transforms.RandomResizedCrop(size=(512, 512), scale=(0.8, 1.0)),\n",
    "        transforms.RandomPerspective(distortion_scale=0.1, p=0.5),\n",
    "        transforms.RandomAffine(degrees=0, translate=(0.15, 0.15), scale=(0.85, 1.15), shear=5),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "\n",
    "    # Create separate dataloaders for each attribute\n",
    "    attribute_loaders = {}\n",
    "    num_classes_per_attr = {}\n",
    "    \n",
    "    for attr in attr_columns:\n",
    "        # Get data where this attribute is not null\n",
    "        attr_df = df_encoded[df_encoded[attr].notna()].copy()\n",
    "        \n",
    "        # Split data for this attribute\n",
    "        train_df, val_df = train_test_split(\n",
    "            attr_df, \n",
    "            test_size=test_size, \n",
    "            stratify=attr_df[attr],  # Stratify by this attribute\n",
    "            random_state=24\n",
    "        )\n",
    "        \n",
    "        # Create datasets\n",
    "        train_dataset = SingleAttributeDataset(\n",
    "            train_df,\n",
    "            image_dir,\n",
    "            attr,\n",
    "            transform_basic=transform,\n",
    "            transform_augmented=transform_augmented,\n",
    "            do_transform=True\n",
    "        )\n",
    "        \n",
    "        val_dataset = SingleAttributeDataset(\n",
    "            val_df,\n",
    "            image_dir,\n",
    "            attr,\n",
    "            transform_basic=transform,\n",
    "            transform_augmented=None, \n",
    "            do_transform=False\n",
    "        )\n",
    "        \n",
    "        # Create dataloaders\n",
    "        train_loader = DataLoader(\n",
    "            train_dataset, \n",
    "            batch_size=batch_size, \n",
    "            shuffle=True, \n",
    "            drop_last=True\n",
    "        )\n",
    "        \n",
    "        val_loader = DataLoader(\n",
    "            val_dataset, \n",
    "            batch_size=batch_size, \n",
    "            shuffle=False,\n",
    "            drop_last=True\n",
    "        )\n",
    "        \n",
    "        # Store loaders for this attribute\n",
    "        attribute_loaders[attr] = {\n",
    "            'train': train_loader,\n",
    "            'val': val_loader,\n",
    "            'num_samples': len(attr_df)\n",
    "        }\n",
    "        \n",
    "        # Store number of classes for this attribute\n",
    "        num_classes_per_attr[attr] = label_encoders.get_num_classes(attr)\n",
    "        \n",
    "        print(f\"{attr}: {len(train_df)} train samples, {len(val_df)} val samples\")\n",
    "    \n",
    "    return attribute_loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.128264Z",
     "iopub.status.busy": "2024-11-10T08:16:15.127796Z",
     "iopub.status.idle": "2024-11-10T08:16:15.145424Z",
     "shell.execute_reply": "2024-11-10T08:16:15.144657Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.128222Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def train_attribute_model(model, dataloaders, num_epochs, model_type, device='cuda', max_grad_norm=1.0):\n",
    "    \"\"\"\n",
    "    Train a model for multi-attribute classification, handling layer freezing/unfreezing per attribute.\n",
    "    \n",
    "    Parameters:\n",
    "    - model: The model wrapped in torch.nn.DataParallel (must have a method to freeze/unfreeze layers).\n",
    "    - dataloaders: Dictionary containing 'train' and 'val' DataLoaders for each attribute.\n",
    "    - num_epochs: Number of epochs to train the model.\n",
    "    - model_type: Type of model being trained (used for saving checkpoints).\n",
    "    - device: The device to run the model on (default is 'cuda').\n",
    "    - max_grad_norm: Maximum norm for gradient clipping (default is 1.0).\n",
    "    \"\"\"\n",
    "    for attr_index, attr in enumerate(dataloaders.keys()):\n",
    "        train_loader = dataloaders[attr]['train']\n",
    "        val_loader = dataloaders[attr]['val']\n",
    "        \n",
    "        best_val_acc = 0.0\n",
    "        \n",
    "        # Ensure model is detached from DataParallel before modifying layers\n",
    "        model.module.freeze_all_except_head(attr_index)\n",
    "        model.to(device)\n",
    "\n",
    "        optimizer = torch.optim.AdamW(model.parameters())\n",
    "        # Add ReduceLROnPlateau scheduler\n",
    "        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "            optimizer, \n",
    "            mode='min',\n",
    "            factor=0.1,\n",
    "            patience=3,\n",
    "            verbose=True\n",
    "        )\n",
    "        criterion = torch.nn.CrossEntropyLoss()\n",
    "        \n",
    "        for epoch in range(num_epochs):\n",
    "            # Training phase\n",
    "            model.train()\n",
    "            train_loss = 0\n",
    "            all_train_preds = []\n",
    "            all_train_labels = []\n",
    "\n",
    "            for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1} train {attr}\"):\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "                \n",
    "                # Forward pass for the specific attribute\n",
    "                outputs = model(images, attr_index=attr_index)\n",
    "                loss = criterion(outputs, labels)\n",
    "                train_loss += loss.item()\n",
    "\n",
    "                # Backward pass and optimization\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                \n",
    "                # Add gradient clipping\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n",
    "                \n",
    "                optimizer.step()\n",
    "\n",
    "                # Store predictions and labels for metrics\n",
    "                _, preds = torch.max(outputs, 1)\n",
    "                all_train_preds.extend(preds.cpu().numpy())\n",
    "                all_train_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "            # Calculate training metrics\n",
    "            train_accuracy = accuracy_score(all_train_labels, all_train_preds)\n",
    "            train_precision = precision_score(all_train_labels, all_train_preds, average='weighted')\n",
    "            train_recall = recall_score(all_train_labels, all_train_preds, average='weighted')\n",
    "            train_f1 = f1_score(all_train_labels, all_train_preds, average='weighted')\n",
    "\n",
    "            # Validation phase\n",
    "            model.eval()\n",
    "            val_loss = 0\n",
    "            all_val_preds = []\n",
    "            all_val_labels = []\n",
    "\n",
    "            with torch.no_grad():\n",
    "                for images, labels in tqdm(val_loader, desc=f\"Epoch {epoch+1} val {attr}\"):\n",
    "                    images, labels = images.to(device), labels.to(device)\n",
    "                    outputs = model(images, attr_index=attr_index)\n",
    "                    val_loss += criterion(outputs, labels).item()\n",
    "\n",
    "                    # Store predictions and labels for metrics\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    all_val_preds.extend(preds.cpu().numpy())\n",
    "                    all_val_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "            # Calculate average validation loss for the scheduler\n",
    "            avg_val_loss = val_loss / len(val_loader)\n",
    "            \n",
    "            # Update learning rate based on validation loss\n",
    "            scheduler.step(avg_val_loss)\n",
    "\n",
    "            # Calculate validation metrics\n",
    "            val_accuracy = accuracy_score(all_val_labels, all_val_preds)\n",
    "            val_precision = precision_score(all_val_labels, all_val_preds, average='weighted')\n",
    "            val_recall = recall_score(all_val_labels, all_val_preds, average='weighted')\n",
    "            val_f1 = f1_score(all_val_labels, all_val_preds, average='weighted')\n",
    "\n",
    "            # Save model weights if validation accuracy improves for this attribute\n",
    "            if val_accuracy >= best_val_acc:\n",
    "                best_val_acc = val_accuracy\n",
    "                # Save the model's state dictionary using .module\n",
    "                torch.save(model.module.state_dict(), f'best_model_attr_{model_type}.pth')\n",
    "\n",
    "            # Print current learning rate\n",
    "            current_lr = optimizer.param_groups[0]['lr']\n",
    "            \n",
    "            print(f\"Epoch {epoch+1}/{num_epochs}, Attr: {attr}, Learning Rate: {current_lr:.6f}\")\n",
    "            print(f\"Train Loss: {train_loss/len(train_loader):.4f}, Train Accuracy: {train_accuracy:.4f}, \"\n",
    "                  f\"Precision: {train_precision:.4f}, Recall: {train_recall:.4f}, F1-score: {train_f1:.4f}\")\n",
    "            print(f\"Val Loss: {avg_val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}, \"\n",
    "                  f\"Precision: {val_precision:.4f}, Recall: {val_recall:.4f}, F1-score: {val_f1:.4f}\")\n",
    "            \n",
    "    torch.save(model.module.state_dict(), f'best_model_attr_end_{model_type}.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.147052Z",
     "iopub.status.busy": "2024-11-10T08:16:15.146783Z",
     "iopub.status.idle": "2024-11-10T08:16:15.158217Z",
     "shell.execute_reply": "2024-11-10T08:16:15.157462Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.147022Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "import gc\n",
    "\n",
    "def main(test_c_name):\n",
    "    # Set image directory\n",
    "    image_dir = f'{input_path}/train_images'\n",
    "    # Initialize device\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    # Define a function to train each model sequentially\n",
    "    def train_single_model(data1,data2, num_attr_columns, model_type):\n",
    "        attr_columns = [f'attr_{i}' for i in range(1, num_attr_columns+1)]\n",
    "        print(f\"Preparing data for {model_type}\")\n",
    "        label_encoders = LabelEncoderDict()\n",
    "        label_encoders.fit(data2, attr_columns)\n",
    "        \n",
    "        train_loader, val_loader, label_encoders, num_classes_per_attr = prepare_data(data1,label_encoders, image_dir, batch_size=4, num_attr_columns=num_attr_columns)\n",
    "        world_size = torch.cuda.device_count()\n",
    "        # Initialize the model\n",
    "        print(f\"Initializing model {model_type}\")\n",
    "        model = MultiLabelClassifier(semi_classes_dict[test_c_name]).to(device)\n",
    "        model = torch.nn.DataParallel(model)\n",
    "        # Train the model\n",
    "        # print(f\"Training model with base model {model_type}\")\n",
    "        # model = train_model(model, train_loader, val_loader, num_epochs=NUM_EPOCH, num_classes_per_attr=num_classes_per_attr, model_type=model_type)\n",
    "        print(f\"\\nTraining model with attribute only {model_type}\")\n",
    "\n",
    "        # ERROR IN CODE: TOBE REMOVED\n",
    "        # model.module.load_state_dict(torch.load(BEST_MODEL_FROM_BASE_FIRST_TRAINING, map_location=device), strict=False)\n",
    "        \n",
    "        # Updated statement, TOBE: Runned\n",
    "        model.load_state_dict(torch.load(BEST_MODEL_FROM_BASE_FIRST_TRAINING, map_location=device))\n",
    "        \n",
    "        attr_loaders = prepare_attribute_data(data2,image_dir,label_encoders, batch_size=4,num_attr_columns=num_attr_columns)\n",
    "        train_attribute_model(model,attr_loaders,NUM_ATTR_EPOCHS,model_type)\n",
    "        # Save label encoders\n",
    "        with open(f'label_encoders_{model_type}.pkl', 'wb') as f:\n",
    "            pickle.dump(label_encoders, f)\n",
    "        print(\"---------------------------------------------\")\n",
    "        print(f\"number of classes for {model_type} is {num_classes_per_attr}\")\n",
    "        print(\"---------------------------------------------\")\n",
    "        # Free up memory\n",
    "        del model, train_loader, val_loader, label_encoders, num_classes_per_attr\n",
    "        torch.cuda.empty_cache()\n",
    "        gc.collect()\n",
    "\n",
    "    df1,df2 = do_preprocessing(test_c_name, category_dict[test_c_name])\n",
    "    train_single_model(df1,df2, num_attr_columns=len(semi_classes_dict[test_c_name]), model_type=test_c_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-10T08:16:15.159491Z",
     "iopub.status.busy": "2024-11-10T08:16:15.159218Z",
     "iopub.status.idle": "2024-11-10T08:22:04.455392Z",
     "shell.execute_reply": "2024-11-10T08:22:04.452645Z",
     "shell.execute_reply.started": "2024-11-10T08:16:15.159461Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "main(test_c_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-09T05:45:12.979885Z",
     "iopub.status.busy": "2024-11-09T05:45:12.979565Z",
     "iopub.status.idle": "2024-11-09T05:45:13.004318Z",
     "shell.execute_reply": "2024-11-09T05:45:13.003395Z",
     "shell.execute_reply.started": "2024-11-09T05:45:12.979854Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "test_df_semi = df_test[df_test['Category'] == category_dict[test_c_name]]\n",
    "test_df_semi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-09T05:45:13.005825Z",
     "iopub.status.busy": "2024-11-09T05:45:13.005523Z",
     "iopub.status.idle": "2024-11-09T05:45:13.018505Z",
     "shell.execute_reply": "2024-11-09T05:45:13.017750Z",
     "shell.execute_reply.started": "2024-11-09T05:45:13.005790Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "\n",
    "def load_model(model_path, num_classes_per_attr, device):\n",
    "    # Initialize the model architecture and load the saved weights\n",
    "    model = MultiLabelClassifier(num_classes_per_attr)\n",
    "    model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "    model = torch.nn.DataParallel(model)\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "def load_label_encoders(encoder_path):\n",
    "    with open(encoder_path, 'rb') as f:\n",
    "        encoders = pickle.load(f)\n",
    "    return encoders\n",
    "\n",
    "def preprocess_image(image_path, image_size=(512,512)):\n",
    "    # Define image transformations (same as used during training)\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((512,512)),  # TODO: Change with (512, 512)\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                             std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    \n",
    "    # Open image and apply transformations\n",
    "    image = Image.open(image_path).convert('RGB')\n",
    "    image = transform(image)\n",
    "    image = image.unsqueeze(0)  # Add batch dimension\n",
    "    return image\n",
    "\n",
    "def inference(images, model, label_encoders):\n",
    "    # Perform inference\n",
    "    with torch.no_grad():\n",
    "        outputs = model(x=images,return_features=False)\n",
    "\n",
    "    # Decode predictions\n",
    "    predicted_labels = []\n",
    "    for i, output in enumerate(outputs):\n",
    "        _, predicted = torch.max(output, 1)\n",
    "\n",
    "        attr_name = f'attr_{i + 1}'\n",
    "        if attr_name in label_encoders.encoders:\n",
    "            decoded_label = label_encoders.encoders[attr_name].inverse_transform([predicted.item()])[0]\n",
    "            predicted_labels.append(decoded_label)\n",
    "        else:\n",
    "            raise KeyError(f\"Encoder for {attr_name} not found in the loaded label encoders.\")\n",
    "    \n",
    "    return predicted_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-11-09T05:46:43.901793Z",
     "iopub.status.idle": "2024-11-09T05:46:43.902187Z",
     "shell.execute_reply": "2024-11-09T05:46:43.901994Z",
     "shell.execute_reply.started": "2024-11-09T05:46:43.901975Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "model_path = f'best_model_attr_{test_c_name}.pth'\n",
    "encoder_path = f\"{working_path}/label_encoders_{test_c_name}.pkl\"\n",
    "num_classes_per_attr = NUM_OF_SEMI_CLASSES_OF_COLUMNS\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model = load_model(model_path, num_classes_per_attr, device)\n",
    "label_encoders = load_label_encoders(encoder_path)\n",
    "image_dir = f\"{input_path}/test_images\"\n",
    "interval = len(semi_classes_dict[test_c_name]) + 1\n",
    "preds = {f'attr_{i}': [] for i in range(1, interval)}\n",
    "t1 = time.time()\n",
    "for val in tqdm(test_df_semi['id'], desc='Processing Images', total=len(test_df_semi)):\n",
    "    image_path = f\"{image_dir}/{str(val).zfill(6)}.jpg\"\n",
    "    image = preprocess_image(image_path).to(device)  # Preprocess and send image to device\n",
    "    predictions = inference(image, model, label_encoders)  # Use the already loaded model and encoders\n",
    "    for i in range(1, interval):\n",
    "        preds[f'attr_{i}'].append(predictions[i-1])\n",
    "print(f'Time taken to process images is {time.time() - t1} seconds, which is {len(test_df_semi) / (time.time() - t1)} images per second')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-11-09T05:46:43.903586Z",
     "iopub.status.idle": "2024-11-09T05:46:43.903919Z",
     "shell.execute_reply": "2024-11-09T05:46:43.903764Z",
     "shell.execute_reply.started": "2024-11-09T05:46:43.903747Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "for i in range(1,interval):\n",
    "    test_df_semi[f'attr_{i}'] = preds[f'attr_{i}']\n",
    "test_df_semi    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.status.busy": "2024-11-09T05:46:43.905273Z",
     "iopub.status.idle": "2024-11-09T05:46:43.905656Z",
     "shell.execute_reply": "2024-11-09T05:46:43.905486Z",
     "shell.execute_reply.started": "2024-11-09T05:46:43.905467Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "test_df_semi.to_csv(f'test_attr_validation_df_{test_c_name}.csv',index=False)"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "datasetId": 5923426,
     "sourceId": 9689319,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30787,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}