Transfer Learning: Pre-trained Models & Fine-Tuning

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader
import time
import copy

# === 1. PERSIAPAN DATA ===
# Transformasi untuk training (dengan augmentasi)
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],  # ImageNet mean
                         [0.229, 0.224, 0.225])   # ImageNet std
])

# Transformasi untuk validasi (tanpa augmentasi)
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# Load dataset (ganti dengan path dataset Anda)
# train_dataset = datasets.ImageFolder('data/train', train_transform)
# val_dataset = datasets.ImageFolder('data/val', val_transform)
# train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# val_loader = DataLoader(val_dataset, batch_size=32)

# === 2. LOAD PRE-TRAINED MODEL ===
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

print("=== Arsitektur ResNet50 ===")
print(model)

# === 3. FEATURE EXTRACTION (Strategy 1: Freeze semua, train head) ===

# Bekukan semua parameter
for param in model.parameters():
    param.requires_grad = False

# Ganti classifier head untuk 10 kelas
num_classes = 10
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, 512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, num_classes)
)

# Hitung parameter yang trainable
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=== Feature Extraction Mode ===")
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")
print(f"Frozen parameters: {total_params - trainable_params:,}")
print(f"Trainable %: {trainable_params/total_params*100:.1f}%")

# === 4. FINE-TUNING (Strategy 2: Unfreeze beberapa layer) ===

def unfreeze_layers(model, num_layers_to_unfreeze):
    """Unfreeze layer terakhir dari ResNet"""
    # Freeze semua dulu
    for param in model.parameters():
        param.requires_grad = False
    
    # Unfreeze layer4 dan fc
    children = list(model.children())
    layers_to_unfreeze = children[-num_layers_to_unfreeze:]
    
    for layer in layers_to_unfreeze:
        for param in layer.parameters():
            param.requires_grad = True
    
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"\nUnfroze {num_layers_to_unfreeze} layers")
    print(f"Trainable: {trainable:,} / {total:,} ({trainable/total*100:.1f}%)")

# Unfreeze 2 layer terakhir (layer4 + fc)
unfreeze_layers(model, 2)

# === 5. DIFFERENTIAL LEARNING RATES ===
# Layer awal → lr kecil, layer akhir → lr lebih besar
optimizer = optim.Adam([
    {'params': model.layer4.parameters(), 'lr': 1e-4},
    {'params': model.fc.parameters(), 'lr': 1e-3},
], weight_decay=1e-4)

# Scheduler: kurangi lr saat training plateau
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=3, verbose=True
)

criterion = nn.CrossEntropyLoss()

# === 6. TRAINING LOOP ===
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'\nEpoch {epoch+1}/{num_epochs}')
        print('-' * 40)
        
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
                dataloader = train_loader
            else:
                model.eval()
                dataloader = val_loader
            
            running_loss = 0.0
            running_corrects = 0
            
            for inputs, labels in dataloader:
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                optimizer.zero_grad()
                
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    
                    if phase == 'train':
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                        optimizer.step()
                
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            
            epoch_loss = running_loss / len(dataloader.dataset)
            epoch_acc = running_corrects.double() / len(dataloader.dataset)
            
            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
            
            if phase == 'val':
                scheduler.step(epoch_loss)
                if epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model_wts = copy.deepcopy(model.state_dict())
    
    print(f'\nBest val Acc: {best_acc:.4f}')
    model.load_state_dict(best_model_wts)
    return model

# Jalankan training (uncomment jika data tersedia)
# model = train_model(model, criterion, optimizer, scheduler, num_epochs=20)

# === 7. PREDIKSI ===
def predict_image(model, image_path, transform):
    from PIL import Image
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.eval()
    
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(image)
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        confidence = probabilities[predicted_class].item()
    
    return predicted_class, confidence, probabilities

print("\n✓ Transfer Learning Pipeline Siap!")
print("  Uncomment fungsi train_model() dan siapkan dataset untuk mulai training.")