CNN: Convolutional Neural Networks

import torch
import torch.nn as nn
import torchvision.models as models

# =============================================
# TRANSFER LEARNING: Pre-trained ResNet-18
# =============================================

# Load pre-trained ResNet-18
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

print("=== Original ResNet-18 ===")
print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")

# =============================================
# STRATEGY 1: Feature Extraction
# =============================================
# Freeze ALL convolution layers
for param in model.parameters():
    param.requires_grad = False

# Replace classifier (final FC layer)
# ResNet-18 final layer: model.fc (512 → 1000)
model.fc = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 10)  # 10 classes for CIFAR-10
)

trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"\n=== After Feature Extraction ===")
print(f"Total params: {total_params:,}")
print(f"Trainable params: {trainable_params:,}")
print(f"Frozen params: {total_params - trainable_params:,}")

# =============================================
# STRATEGY 2: Fine-Tuning
# =============================================
model_ft = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# Freeze early layers (layer1, layer2)
for name, param in model_ft.named_parameters():
    if 'layer1' in name or 'layer2' in name:
        param.requires_grad = False

# Unfreeze later layers (layer3, layer4) + FC
for name, param in model_ft.named_parameters():
    if 'layer3' in name or 'layer4' in name:
        param.requires_grad = True

# Replace FC
model_ft.fc = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 10)
)

trainable_ft = sum(p.numel() for p in model_ft.parameters() if p.requires_grad)
print(f"\n=== After Fine-Tuning ===")
print(f"Trainable params: {trainable_ft:,}")
print(f"Frozen early layers: layer1, layer2")
print(f"Unfrozen: layer3, layer4, fc")

# =============================================
# DIFFERENT LEARNING RATES (LR Discriminative)
# =============================================
# Berikan learning rate berbeda per layer group
param_groups = [
    {'params': [p for n, p in model_ft.named_parameters()
                if 'layer1' in n or 'layer2' in n],
     'lr': 1e-5},          # Frozen layer: very small LR (if any)
    {'params': [p for n, p in model_ft.named_parameters()
                if 'layer3' in n or 'layer4' in n],
     'lr': 1e-4},          # Middle layers: small LR
    {'params': model_ft.fc.parameters(),
     'lr': 1e-3}           # New layers: larger LR
]

optimizer_ft = torch.optim.Adam(param_groups, weight_decay=1e-4)
print("\n=== Discriminative Learning Rates ===")
print("  Early layers: 1e-5")
print("  Middle layers: 1e-4")
print("  FC layer:     1e-3")