{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "567f48ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🧹 Membersihkan memori...\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, models, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import time, os, copy, gc\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')  # ✅ FIX: Use non-interactive backend initially\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "import numpy as np\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# 0. CLEAN MEMORY\n",
    "print(\"🧹 Membersihkan memori...\")\n",
    "gc.collect()\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.empty_cache()\n",
    "print(\"=\"*50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "10035e1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🚀 Device: cuda\n"
     ]
    }
   ],
   "source": [
    "# 1. CONFIG\n",
    "\n",
    "DATA_DIR = 'Dataset3'\n",
    "IMG_SIZE = 256\n",
    "BATCH_SIZE = 32\n",
    "NUM_WORKERS = 4\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "torch.manual_seed(42)\n",
    "torch.backends.cudnn.benchmark = True\n",
    "print(f\"🚀 Device: {DEVICE}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2e77d25b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. TRANSFORM (REALISTIC)\n",
    "# ==========================================\n",
    "data_transforms = {\n",
    "    'train': transforms.Compose([\n",
    "        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
    "        transforms.RandomAffine(15, translate=(0.1,0.1), scale=(0.9,1.1)),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        #transforms.ColorJitter(0.1,0.1),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])\n",
    "    ]),\n",
    "    'val': transforms.Compose([\n",
    "        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])\n",
    "    ])\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e7218bec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 Kelas: ['Kaleng', 'PET_Bening', 'PET_Berwarna', 'Residu']\n",
      "📊 Train: 2141 | Val: 936\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# 3. DATASET\n",
    "\n",
    "image_datasets = {\n",
    "    x: datasets.ImageFolder(os.path.join(DATA_DIR, x), data_transforms[x])\n",
    "    for x in ['train','val']\n",
    "}\n",
    "dataloaders = {\n",
    "    'train': DataLoader(\n",
    "        image_datasets['train'],\n",
    "        batch_size=BATCH_SIZE,\n",
    "        shuffle=True,\n",
    "        num_workers=NUM_WORKERS,\n",
    "        pin_memory=True\n",
    "    ),\n",
    "    'val': DataLoader(\n",
    "        image_datasets['val'],\n",
    "        batch_size=BATCH_SIZE,\n",
    "        shuffle=False,\n",
    "        num_workers=NUM_WORKERS,\n",
    "        pin_memory=True\n",
    "    )\n",
    "}\n",
    "\n",
    "dataset_sizes = {x: len(image_datasets[x]) for x in ['train','val']}\n",
    "class_names = image_datasets['train'].classes\n",
    "\n",
    "print(f\"📊 Kelas: {class_names}\")\n",
    "print(f\"📊 Train: {dataset_sizes['train']} | Val: {dataset_sizes['val']}\")\n",
    "print(\"=\"*50)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "26793eb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. PLOT FUNCTIONS (FIXED VERSION)\n",
    "\n",
    "def save_training_plot(history, save_path='training_progress.png'):\n",
    "    \"\"\"Save training history plot\"\"\"\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n",
    "    \n",
    "    epochs = range(1, len(history['train_acc']) + 1)\n",
    "    \n",
    "    # Plot Accuracy\n",
    "    ax1.plot(epochs, history['train_acc'], 'b-o', label='Train Acc', linewidth=2, markersize=4)\n",
    "    ax1.plot(epochs, history['val_acc'], 'r-s', label='Val Acc', linewidth=2, markersize=4)\n",
    "    ax1.set_xlabel('Epoch', fontsize=12)\n",
    "    ax1.set_ylabel('Accuracy', fontsize=12)\n",
    "    ax1.set_title('Training & Validation Accuracy', fontsize=14, fontweight='bold')\n",
    "    ax1.legend(loc='best', fontsize=10)\n",
    "    ax1.grid(True, alpha=0.3)\n",
    "    \n",
    "    # Plot Loss\n",
    "    ax2.plot(epochs, history['train_loss'], 'b-o', label='Train Loss', linewidth=2, markersize=4)\n",
    "    ax2.plot(epochs, history['val_loss'], 'r-s', label='Val Loss', linewidth=2, markersize=4)\n",
    "    ax2.set_xlabel('Epoch', fontsize=12)\n",
    "    ax2.set_ylabel('Loss', fontsize=12)\n",
    "    ax2.set_title('Training & Validation Loss', fontsize=14, fontweight='bold')\n",
    "    ax2.legend(loc='best', fontsize=10)\n",
    "    ax2.grid(True, alpha=0.3)\n",
    "    \n",
    "    # Add gap annotation on accuracy plot\n",
    "    if len(history['train_acc']) > 0:\n",
    "        latest_gap = history['train_acc'][-1] - history['val_acc'][-1]\n",
    "        gap_color = 'red' if latest_gap > 0.15 else 'green'\n",
    "        ax1.text(0.02, 0.98, f'Latest Gap: {latest_gap:.3f}', \n",
    "                transform=ax1.transAxes, fontsize=10,\n",
    "                verticalalignment='top', bbox=dict(boxstyle='round', \n",
    "                facecolor=gap_color, alpha=0.3))\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=150, bbox_inches='tight')\n",
    "    plt.close()\n",
    "    print(f\"📊 Training plot saved to {save_path}\")\n",
    "\n",
    "\n",
    "def save_confusion_matrix(y_true, y_pred, class_names, epoch, save_path='confusion_matrix.png'):\n",
    "    \"\"\"Save confusion matrix\"\"\"\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    cm_normalized = cm.astype('float') / cm.sum(axis=1, keepdims=True)\n",
    "    \n",
    "    fig, ax = plt.subplots(figsize=(10, 8))\n",
    "    \n",
    "    # Plot with better colors\n",
    "    im = ax.imshow(cm_normalized, interpolation='nearest', cmap='Blues')\n",
    "    ax.set_title(f'Confusion Matrix - Epoch {epoch}', fontsize=16, fontweight='bold', pad=20)\n",
    "    \n",
    "    # Colorbar\n",
    "    cbar = plt.colorbar(im, ax=ax)\n",
    "    cbar.set_label('Normalized Frequency', rotation=270, labelpad=20)\n",
    "    \n",
    "    # Axis labels\n",
    "    tick_marks = np.arange(len(class_names))\n",
    "    ax.set_xticks(tick_marks)\n",
    "    ax.set_yticks(tick_marks)\n",
    "    ax.set_xticklabels(class_names, rotation=45, ha='right', fontsize=10)\n",
    "    ax.set_yticklabels(class_names, fontsize=10)\n",
    "    \n",
    "    ax.set_ylabel('True Label', fontsize=12, fontweight='bold')\n",
    "    ax.set_xlabel('Predicted Label', fontsize=12, fontweight='bold')\n",
    "    \n",
    "    # Add text annotations\n",
    "    thresh = cm_normalized.max() / 2.\n",
    "    for i in range(len(class_names)):\n",
    "        for j in range(len(class_names)):\n",
    "            # Show both count and percentage\n",
    "            text = f'{cm[i, j]}\\n({cm_normalized[i, j]:.2%})'\n",
    "            ax.text(j, i, text,\n",
    "                   ha=\"center\", va=\"center\", fontsize=9,\n",
    "                   color=\"white\" if cm_normalized[i, j] > thresh else \"black\")\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=150, bbox_inches='tight')\n",
    "    plt.close()\n",
    "    print(f\"📊 Confusion matrix saved to {save_path}\")\n",
    "\n",
    "\n",
    "def print_classification_report(y_true, y_pred, class_names):\n",
    "    \"\"\"Print per-class accuracy\"\"\"\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    print(\"📊 PER-CLASS PERFORMANCE\")\n",
    "    print(\"=\"*60)\n",
    "    \n",
    "    for i, class_name in enumerate(class_names):\n",
    "        if cm[i].sum() > 0:\n",
    "            accuracy = cm[i, i] / cm[i].sum()\n",
    "            total = cm[i].sum()\n",
    "            correct = cm[i, i]\n",
    "            print(f\"{class_name:15s} | Acc: {accuracy:6.2%} | Correct: {correct:4d}/{total:4d}\")\n",
    "        else:\n",
    "            print(f\"{class_name:15s} | No samples\")\n",
    "    \n",
    "    print(\"=\"*60 + \"\\n\")\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "daa3d291",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. TRAIN FUNCTION (FIXED)\n",
    "def train_model(model, criterion, optimizer, scheduler, epochs=20, unfreeze_epoch=12):\n",
    "    history = {'train_loss':[], 'val_loss':[], 'train_acc':[], 'val_acc':[]}\n",
    "    best_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "    start_time = time.time()\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        print(f\"\\n{'='*60}\")\n",
    "        print(f\"EPOCH {epoch+1}/{epochs}\")\n",
    "        print('='*60)\n",
    "\n",
    "        # Progressive unfreezing\n",
    "        if epoch == unfreeze_epoch:\n",
    "            print(\"🔓 Unfreezing deeper layers (features[7:])...\")\n",
    "            for p in model.features[7:].parameters():\n",
    "                p.requires_grad = True\n",
    "            print(\"✅ More layers are now trainable\\n\")\n",
    "\n",
    "        for phase in ['train','val']:\n",
    "            model.train() if phase=='train' else model.eval()\n",
    "            loss_sum, correct = 0, 0\n",
    "            all_preds, all_labels = [], []\n",
    "\n",
    "            pbar = tqdm(dataloaders[phase], \n",
    "                       desc=f\"{phase.upper():5s}\",\n",
    "                       leave=False,\n",
    "                       ncols=100)\n",
    "\n",
    "            for x, y in pbar:\n",
    "                x, y = x.to(DEVICE), y.to(DEVICE)\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                with torch.set_grad_enabled(phase=='train'):\n",
    "                    out = model(x)\n",
    "                    loss = criterion(out, y)\n",
    "                    pred = out.argmax(1)\n",
    "                    \n",
    "                    if phase=='train':\n",
    "                        loss.backward()\n",
    "                        # Gradient clipping for stability\n",
    "                        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                        optimizer.step()\n",
    "\n",
    "                loss_sum += loss.item() * x.size(0)\n",
    "                correct += (pred==y).sum().item()\n",
    "\n",
    "                # Update progress bar\n",
    "                pbar.set_postfix({\n",
    "                    'loss': f'{loss.item():.4f}',\n",
    "                    'acc': f'{(pred==y).float().mean():.4f}'\n",
    "                })\n",
    "\n",
    "                if phase=='val':\n",
    "                    all_preds.extend(pred.cpu().numpy())\n",
    "                    all_labels.extend(y.cpu().numpy())\n",
    "\n",
    "            epoch_loss = loss_sum / dataset_sizes[phase]\n",
    "            epoch_acc = correct / dataset_sizes[phase]\n",
    "\n",
    "            history[f'{phase}_loss'].append(epoch_loss)\n",
    "            history[f'{phase}_acc'].append(epoch_acc)\n",
    "\n",
    "            # Print results\n",
    "            print(f\"{phase.upper():5s} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f} ({epoch_acc*100:.2f}%)\")\n",
    "\n",
    "            if phase=='val':\n",
    "                # Update scheduler (for CosineAnnealingLR)\n",
    "                scheduler.step()\n",
    "                \n",
    "                # Calculate and display train-val gap\n",
    "                train_val_gap = history['train_acc'][-1] - history['val_acc'][-1]\n",
    "                gap_status = \"⚠️ OVERFITTING!\" if train_val_gap > 0.15 else \"✅ Good\"\n",
    "                print(f\"\\n📊 Train-Val Gap: {train_val_gap:.4f} {gap_status}\")\n",
    "                \n",
    "                # Save confusion matrix\n",
    "                save_confusion_matrix(all_labels, all_preds, class_names, \n",
    "                                    epoch+1, f'confusion_matrix_epoch_{epoch+1:02d}.png')\n",
    "                \n",
    "                # Print per-class performance\n",
    "                print_classification_report(all_labels, all_preds, class_names)\n",
    "                \n",
    "                # Save best model\n",
    "                if epoch_acc > best_acc:\n",
    "                    best_acc = epoch_acc\n",
    "                    best_wts = copy.deepcopy(model.state_dict())\n",
    "                    print(f\"⭐ NEW BEST MODEL! Val Acc: {best_acc:.4f}\\n\")\n",
    "\n",
    "        # Save training progress plot after each epoch\n",
    "        save_training_plot(history, 'training_progress.png')\n",
    "        \n",
    "        # Memory cleanup\n",
    "        gc.collect()\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "    # Training complete\n",
    "    elapsed = time.time() - start_time\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"⏱️  Training completed in {elapsed//60:.0f}m {elapsed%60:.0f}s\")\n",
    "    print(f\"🏆 Best Validation Accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)\")\n",
    "    print('='*60 + \"\\n\")\n",
    "\n",
    "    # Load best weights\n",
    "    model.load_state_dict(best_wts)\n",
    "    \n",
    "    # Save final plots\n",
    "    save_training_plot(history, 'training_progress_FINAL.png')\n",
    "    \n",
    "    return model, history\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e33a5eb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainable params : 1,281,000\n",
      "Total params     : 3,504,872\n",
      "features.0.0.weight                                | trainable = False\n",
      "features.0.1.weight                                | trainable = False\n",
      "features.0.1.bias                                  | trainable = False\n",
      "features.1.conv.0.0.weight                         | trainable = False\n",
      "features.1.conv.0.1.weight                         | trainable = False\n",
      "features.1.conv.0.1.bias                           | trainable = False\n",
      "features.1.conv.1.weight                           | trainable = False\n",
      "features.1.conv.2.weight                           | trainable = False\n",
      "features.1.conv.2.bias                             | trainable = False\n",
      "features.2.conv.0.0.weight                         | trainable = False\n",
      "features.2.conv.0.1.weight                         | trainable = False\n",
      "features.2.conv.0.1.bias                           | trainable = False\n",
      "features.2.conv.1.0.weight                         | trainable = False\n",
      "features.2.conv.1.1.weight                         | trainable = False\n",
      "features.2.conv.1.1.bias                           | trainable = False\n",
      "features.2.conv.2.weight                           | trainable = False\n",
      "features.2.conv.3.weight                           | trainable = False\n",
      "features.2.conv.3.bias                             | trainable = False\n",
      "features.3.conv.0.0.weight                         | trainable = False\n",
      "features.3.conv.0.1.weight                         | trainable = False\n",
      "features.3.conv.0.1.bias                           | trainable = False\n",
      "features.3.conv.1.0.weight                         | trainable = False\n",
      "features.3.conv.1.1.weight                         | trainable = False\n",
      "features.3.conv.1.1.bias                           | trainable = False\n",
      "features.3.conv.2.weight                           | trainable = False\n",
      "features.3.conv.3.weight                           | trainable = False\n",
      "features.3.conv.3.bias                             | trainable = False\n",
      "features.4.conv.0.0.weight                         | trainable = False\n",
      "features.4.conv.0.1.weight                         | trainable = False\n",
      "features.4.conv.0.1.bias                           | trainable = False\n",
      "features.4.conv.1.0.weight                         | trainable = False\n",
      "features.4.conv.1.1.weight                         | trainable = False\n",
      "features.4.conv.1.1.bias                           | trainable = False\n",
      "features.4.conv.2.weight                           | trainable = False\n",
      "features.4.conv.3.weight                           | trainable = False\n",
      "features.4.conv.3.bias                             | trainable = False\n",
      "features.5.conv.0.0.weight                         | trainable = False\n",
      "features.5.conv.0.1.weight                         | trainable = False\n",
      "features.5.conv.0.1.bias                           | trainable = False\n",
      "features.5.conv.1.0.weight                         | trainable = False\n",
      "features.5.conv.1.1.weight                         | trainable = False\n",
      "features.5.conv.1.1.bias                           | trainable = False\n",
      "features.5.conv.2.weight                           | trainable = False\n",
      "features.5.conv.3.weight                           | trainable = False\n",
      "features.5.conv.3.bias                             | trainable = False\n",
      "features.6.conv.0.0.weight                         | trainable = False\n",
      "features.6.conv.0.1.weight                         | trainable = False\n",
      "features.6.conv.0.1.bias                           | trainable = False\n",
      "features.6.conv.1.0.weight                         | trainable = False\n",
      "features.6.conv.1.1.weight                         | trainable = False\n",
      "features.6.conv.1.1.bias                           | trainable = False\n",
      "features.6.conv.2.weight                           | trainable = False\n",
      "features.6.conv.3.weight                           | trainable = False\n",
      "features.6.conv.3.bias                             | trainable = False\n",
      "features.7.conv.0.0.weight                         | trainable = False\n",
      "features.7.conv.0.1.weight                         | trainable = False\n",
      "features.7.conv.0.1.bias                           | trainable = False\n",
      "features.7.conv.1.0.weight                         | trainable = False\n",
      "features.7.conv.1.1.weight                         | trainable = False\n",
      "features.7.conv.1.1.bias                           | trainable = False\n",
      "features.7.conv.2.weight                           | trainable = False\n",
      "features.7.conv.3.weight                           | trainable = False\n",
      "features.7.conv.3.bias                             | trainable = False\n",
      "features.8.conv.0.0.weight                         | trainable = False\n",
      "features.8.conv.0.1.weight                         | trainable = False\n",
      "features.8.conv.0.1.bias                           | trainable = False\n",
      "features.8.conv.1.0.weight                         | trainable = False\n",
      "features.8.conv.1.1.weight                         | trainable = False\n",
      "features.8.conv.1.1.bias                           | trainable = False\n",
      "features.8.conv.2.weight                           | trainable = False\n",
      "features.8.conv.3.weight                           | trainable = False\n",
      "features.8.conv.3.bias                             | trainable = False\n",
      "features.9.conv.0.0.weight                         | trainable = False\n",
      "features.9.conv.0.1.weight                         | trainable = False\n",
      "features.9.conv.0.1.bias                           | trainable = False\n",
      "features.9.conv.1.0.weight                         | trainable = False\n",
      "features.9.conv.1.1.weight                         | trainable = False\n",
      "features.9.conv.1.1.bias                           | trainable = False\n",
      "features.9.conv.2.weight                           | trainable = False\n",
      "features.9.conv.3.weight                           | trainable = False\n",
      "features.9.conv.3.bias                             | trainable = False\n",
      "features.10.conv.0.0.weight                        | trainable = False\n",
      "features.10.conv.0.1.weight                        | trainable = False\n",
      "features.10.conv.0.1.bias                          | trainable = False\n",
      "features.10.conv.1.0.weight                        | trainable = False\n",
      "features.10.conv.1.1.weight                        | trainable = False\n",
      "features.10.conv.1.1.bias                          | trainable = False\n",
      "features.10.conv.2.weight                          | trainable = False\n",
      "features.10.conv.3.weight                          | trainable = False\n",
      "features.10.conv.3.bias                            | trainable = False\n",
      "features.11.conv.0.0.weight                        | trainable = False\n",
      "features.11.conv.0.1.weight                        | trainable = False\n",
      "features.11.conv.0.1.bias                          | trainable = False\n",
      "features.11.conv.1.0.weight                        | trainable = False\n",
      "features.11.conv.1.1.weight                        | trainable = False\n",
      "features.11.conv.1.1.bias                          | trainable = False\n",
      "features.11.conv.2.weight                          | trainable = False\n",
      "features.11.conv.3.weight                          | trainable = False\n",
      "features.11.conv.3.bias                            | trainable = False\n",
      "features.12.conv.0.0.weight                        | trainable = False\n",
      "features.12.conv.0.1.weight                        | trainable = False\n",
      "features.12.conv.0.1.bias                          | trainable = False\n",
      "features.12.conv.1.0.weight                        | trainable = False\n",
      "features.12.conv.1.1.weight                        | trainable = False\n",
      "features.12.conv.1.1.bias                          | trainable = False\n",
      "features.12.conv.2.weight                          | trainable = False\n",
      "features.12.conv.3.weight                          | trainable = False\n",
      "features.12.conv.3.bias                            | trainable = False\n",
      "features.13.conv.0.0.weight                        | trainable = False\n",
      "features.13.conv.0.1.weight                        | trainable = False\n",
      "features.13.conv.0.1.bias                          | trainable = False\n",
      "features.13.conv.1.0.weight                        | trainable = False\n",
      "features.13.conv.1.1.weight                        | trainable = False\n",
      "features.13.conv.1.1.bias                          | trainable = False\n",
      "features.13.conv.2.weight                          | trainable = False\n",
      "features.13.conv.3.weight                          | trainable = False\n",
      "features.13.conv.3.bias                            | trainable = False\n",
      "features.14.conv.0.0.weight                        | trainable = False\n",
      "features.14.conv.0.1.weight                        | trainable = False\n",
      "features.14.conv.0.1.bias                          | trainable = False\n",
      "features.14.conv.1.0.weight                        | trainable = False\n",
      "features.14.conv.1.1.weight                        | trainable = False\n",
      "features.14.conv.1.1.bias                          | trainable = False\n",
      "features.14.conv.2.weight                          | trainable = False\n",
      "features.14.conv.3.weight                          | trainable = False\n",
      "features.14.conv.3.bias                            | trainable = False\n",
      "features.15.conv.0.0.weight                        | trainable = False\n",
      "features.15.conv.0.1.weight                        | trainable = False\n",
      "features.15.conv.0.1.bias                          | trainable = False\n",
      "features.15.conv.1.0.weight                        | trainable = False\n",
      "features.15.conv.1.1.weight                        | trainable = False\n",
      "features.15.conv.1.1.bias                          | trainable = False\n",
      "features.15.conv.2.weight                          | trainable = False\n",
      "features.15.conv.3.weight                          | trainable = False\n",
      "features.15.conv.3.bias                            | trainable = False\n",
      "features.16.conv.0.0.weight                        | trainable = False\n",
      "features.16.conv.0.1.weight                        | trainable = False\n",
      "features.16.conv.0.1.bias                          | trainable = False\n",
      "features.16.conv.1.0.weight                        | trainable = False\n",
      "features.16.conv.1.1.weight                        | trainable = False\n",
      "features.16.conv.1.1.bias                          | trainable = False\n",
      "features.16.conv.2.weight                          | trainable = False\n",
      "features.16.conv.3.weight                          | trainable = False\n",
      "features.16.conv.3.bias                            | trainable = False\n",
      "features.17.conv.0.0.weight                        | trainable = False\n",
      "features.17.conv.0.1.weight                        | trainable = False\n",
      "features.17.conv.0.1.bias                          | trainable = False\n",
      "features.17.conv.1.0.weight                        | trainable = False\n",
      "features.17.conv.1.1.weight                        | trainable = False\n",
      "features.17.conv.1.1.bias                          | trainable = False\n",
      "features.17.conv.2.weight                          | trainable = False\n",
      "features.17.conv.3.weight                          | trainable = False\n",
      "features.17.conv.3.bias                            | trainable = False\n",
      "features.18.0.weight                               | trainable = False\n",
      "features.18.1.weight                               | trainable = False\n",
      "features.18.1.bias                                 | trainable = False\n",
      "classifier.1.weight                                | trainable = True\n",
      "classifier.1.bias                                  | trainable = True\n"
     ]
    }
   ],
   "source": [
    "model = models.mobilenet_v2(weights='IMAGENET1K_V1')\n",
    "\n",
    "#print(model)\n",
    "\n",
    "#for name, param in model.named_parameters():\n",
    "#    if param.requires_grad:\n",
    "#        print(name)\n",
    "for p in model.features.parameters():\n",
    "    p.requires_grad = False\n",
    "\n",
    "#for p in model.features[-1:].parameters():\n",
    "#    p.requires_grad = True\n",
    "\n",
    "trainable_params = sum(\n",
    "    p.numel() for p in model.parameters() if p.requires_grad\n",
    ")\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "\n",
    "print(f\"Trainable params : {trainable_params:,}\")\n",
    "print(f\"Total params     : {total_params:,}\")\n",
    "\n",
    "for name, param in model.named_parameters():\n",
    "    print(f\"{name:50} | trainable = {param.requires_grad}\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f814a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "🔧 Setting up model...\n",
      "✅ Model loaded to cuda\n",
      "✅ Number of classes: 4\n",
      "✅ Trainable parameters: 891,204 / 2,228,996 (40.0%)\n",
      "==================================================\n"
     ]
    }
   ],
   "source": [
    "# 6. MODEL SETUP\n",
    "\n",
    "print(\"\\n🔧 Setting up model...\")\n",
    "model = models.mobilenet_v2(weights='IMAGENET1K_V1')\n",
    "\n",
    "# Freeze all features initially\n",
    "for p in model.features.parameters():\n",
    "    p.requires_grad = False\n",
    "\n",
    "# Unfreeze only last 2 blocks\n",
    "for p in model.features[-2:].parameters():\n",
    "    p.requires_grad = True\n",
    "\n",
    "# Classifier\n",
    "num_ftrs = model.classifier[1].in_features\n",
    "model.classifier = nn.Sequential(\n",
    "    nn.Dropout(0.4),\n",
    "    nn.Linear(num_ftrs, len(class_names))\n",
    ")\n",
    "model.to(DEVICE)\n",
    "\n",
    "print(f\"✅ Model loaded to {DEVICE}\")\n",
    "print(f\"✅ Number of classes: {len(class_names)}\")\n",
    "\n",
    "# Count trainable parameters\n",
    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"✅ Trainable parameters: {trainable_params:,} / {total_params:,} ({trainable_params/total_params*100:.1f}%)\")\n",
    "print(\"=\"*50)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7c01e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 7. LOSS & OPTIMIZER\n",
    "# ==========================================\n",
    "class_weights = torch.tensor([1.2, 1.5, 1.2, 0.8]).to(DEVICE)\n",
    "criterion = nn.CrossEntropyLoss(label_smoothing=0.3)\n",
    "\n",
    "optimizer = optim.Adam([\n",
    "    {'params': model.features[-2:].parameters(), 'lr': 1e-5},\n",
    "    {'params': model.classifier.parameters(), 'lr': 1e-4}\n",
    "], weight_decay=1e-4)\n",
    "\n",
    "scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)\n",
    "\n",
    "print(\"✅ Optimizer: Adam\")\n",
    "print(\"✅ Scheduler: CosineAnnealingLR\")\n",
    "print(\"✅ Loss: CrossEntropyLoss (label_smoothing=0.3)\")\n",
    "print(\"=\"*50)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94e464b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 8. TRAIN\n",
    "# ==========================================\n",
    "print(\"\\n🚀 STARTING TRAINING...\\n\")\n",
    "model, history = train_model(model, criterion, optimizer, scheduler, \n",
    "                             epochs=20, unfreeze_epoch=12)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "619b3f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 9. SAVE MODEL\n",
    "# ==========================================\n",
    "print(\"\\n💾 Saving model...\")\n",
    "save_dict = {\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'class_names': class_names,\n",
    "    'history': history,\n",
    "    'optimizer_state_dict': optimizer.state_dict()\n",
    "}\n",
    "\n",
    "torch.save(save_dict, \"model_sampah_BEST2.pth\")\n",
    "print(\"✅ Model saved to 'model_sampah_BEST2.pth'\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab728ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 10. FINAL SUMMARY\n",
    "# ==========================================\n",
    "print(\"\\n\" + \"=\"*60)\n",
    "print(\"📊 TRAINING SUMMARY\")\n",
    "print(\"=\"*60)\n",
    "print(f\"Final Train Acc: {history['train_acc'][-1]:.4f} ({history['train_acc'][-1]*100:.2f}%)\")\n",
    "print(f\"Final Val Acc:   {history['val_acc'][-1]:.4f} ({history['val_acc'][-1]*100:.2f}%)\")\n",
    "print(f\"Best Val Acc:    {max(history['val_acc']):.4f} ({max(history['val_acc'])*100:.2f}%)\")\n",
    "print(f\"Train-Val Gap:   {history['train_acc'][-1] - history['val_acc'][-1]:.4f}\")\n",
    "print(\"=\"*60)\n",
    "\n",
    "print(\"\\n📁 Generated files:\")\n",
    "print(\"  - model_sampah_BEST2.pth (saved model)\")\n",
    "print(\"  - training_progress_FINAL.png (final training curves)\")\n",
    "print(f\"  - confusion_matrix_epoch_XX.png (confusion matrix for each epoch)\")\n",
    "print(\"\\n✅ ALL DONE! 🎉\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "skripsi",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
