Fine-Tuning Strategies

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

import torch
import torch.nn as nn
from torchvision import models
from typing import Optional, Tuple
 
class FullFineTuningModel(nn.Module):
    """
    A production-ready fine-tuning architecture that:
    1. Loads pre-trained backbone weights
    2. Replaces task-specific head for new output space
    3. Enables full gradient flow through all parameters
    4. Supports optional auxiliary objectives
    """
    
    def __init__(
        self,
        backbone_name: str = "resnet50",
        num_classes: int = 1000,
        pretrained: bool = True,
        dropout_rate: float = 0.2
    ):
        super().__init__()
        
        # Load pre-trained backbone (all layers will be fine-tuned)
        if backbone_name == "resnet50":
            self.backbone = models.resnet50(pretrained=pretrained)
            in_features = self.backbone.fc.in_features
            # Remove original classifier
            self.backbone.fc = nn.Identity()
        elif backbone_name == "efficientnet_b0":
            self.backbone = models.efficientnet_b0(pretrained=pretrained)
            in_features = self.backbone.classifier[1].in_features
            self.backbone.classifier = nn.Identity()
        else:
            raise ValueError(f"Unsupported backbone: {backbone_name}")
        
        # New task-specific head with regularization
        self.classifier = nn.Sequential(
            nn.Dropout(p=dropout_rate),
            nn.Linear(in_features, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(p=dropout_rate),
            nn.Linear(512, num_classes)
        )
        
        # Initialize new layers (backbone retains pre-trained weights)
        self._initialize_classifier()
    
    def _initialize_classifier(self):
        """Xavier/Kaiming initialization for new layers only."""
        for module in self.classifier.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.backbone(x)
        logits = self.classifier(features)
        return logits
    
    def get_parameter_groups(
        self,
        backbone_lr: float = 1e-4,
        head_lr: float = 1e-3
    ) -> list:
        """
        Create parameter groups with different learning rates.
        Backbone uses lower LR to preserve pre-trained knowledge.
        Head uses higher LR for faster adaptation.
        """
        return [
            {"params": self.backbone.parameters(), "lr": backbone_lr},
            {"params": self.classifier.parameters(), "lr": head_lr}
        ]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, LinearLR, SequentialLR
from torch.utils.data import DataLoader
from tqdm import tqdm
 
class FineTuningTrainer:
    """
    Enterprise-grade training pipeline for full fine-tuning.
    
    Key features:
    - Differential learning rates (lower for backbone, higher for head)
    - Warmup period to stabilize early training
    - Cosine annealing for smooth convergence
    - Gradient clipping to prevent exploding gradients
    - Early stopping with patience
    """
    
    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        backbone_lr: float = 1e-4,
        head_lr: float = 1e-3,
        weight_decay: float = 1e-4,
        max_epochs: int = 50,
        warmup_epochs: int = 5,
        grad_clip_norm: float = 1.0,
        patience: int = 10,
        device: str = "cuda"
    ):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.max_epochs = max_epochs
        self.grad_clip_norm = grad_clip_norm
        self.patience = patience
        
        # Setup optimizer with differential learning rates
        param_groups = model.get_parameter_groups(backbone_lr, head_lr)
        self.optimizer = optim.AdamW(param_groups, weight_decay=weight_decay)
        
        # Setup learning rate schedule: warmup -> cosine annealing
        warmup_scheduler = LinearLR(
            self.optimizer, 
            start_factor=0.1, 
            total_iters=warmup_epochs * len(train_loader)
        )
        main_scheduler = CosineAnnealingWarmRestarts(
            self.optimizer,
            T_0=10,  # Restart period
            T_mult=2  # Double period after each restart
        )
        self.scheduler = SequentialLR(
            self.optimizer,
            schedulers=[warmup_scheduler, main_scheduler],
            milestones=[warmup_epochs * len(train_loader)]
        )
        
        self.criterion = nn.CrossEntropyLoss()
        self.best_val_loss = float('inf')
        self.epochs_without_improvement = 0
    
    def train_epoch(self) -> Tuple[float, float]:
        """Execute one training epoch."""
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (inputs, targets) in enumerate(tqdm(self.train_loader)):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            
            self.optimizer.zero_grad()
            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()
            
            # Gradient clipping prevents destabilization
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), 
                self.grad_clip_norm
            )
            
            self.optimizer.step()
            self.scheduler.step()
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        avg_loss = total_loss / len(self.train_loader)
        accuracy = 100.0 * correct / total
        return avg_loss, accuracy
    
    def validate(self) -> Tuple[float, float]:
        """Evaluate on validation set."""
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, targets in self.val_loader:
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                outputs = self.model(inputs)
                loss = self.criterion(outputs, targets)
                
                total_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()
        
        avg_loss = total_loss / len(self.val_loader)
        accuracy = 100.0 * correct / total
        return avg_loss, accuracy
    
    def fit(self) -> dict:
        """Full training loop with early stopping."""
        history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []}
        
        for epoch in range(self.max_epochs):
            train_loss, train_acc = self.train_epoch()
            val_loss, val_acc = self.validate()
            
            history["train_loss"].append(train_loss)
            history["val_loss"].append(val_loss)
            history["train_acc"].append(train_acc)
            history["val_acc"].append(val_acc)
            
            print(f"Epoch {epoch+1}/{self.max_epochs}")
            print(f"  Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%")
            print(f"  Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%")
            
            # Early stopping check
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self.epochs_without_improvement = 0
                # Save best model
                torch.save(self.model.state_dict(), "best_model.pth")
            else:
                self.epochs_without_improvement += 1
                if self.epochs_without_improvement >= self.patience:
                    print(f"Early stopping at epoch {epoch+1}")
                    break
        
        return history

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

import optuna
from optuna.trial import Trial
 
def objective(trial: Trial) -> float:
    """
    Optuna objective function for hyperparameter search.
    
    Key hyperparameters to tune:
    - Learning rate (backbone and head separately)
    - Weight decay
    - Dropout rate
    - Warmup epochs
    """
    
    # Define search space
    backbone_lr = trial.suggest_float("backbone_lr", 1e-6, 1e-4, log=True)
    head_lr = trial.suggest_float("head_lr", 1e-4, 1e-2, log=True)
    weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-2, log=True)
    dropout_rate = trial.suggest_float("dropout_rate", 0.1, 0.5)
    warmup_epochs = trial.suggest_int("warmup_epochs", 1, 10)
    
    # Ensure head learns faster than backbone
    if head_lr < backbone_lr:
        head_lr = backbone_lr * 10
    
    # Create model with sampled hyperparameters
    model = FullFineTuningModel(
        backbone_name="resnet50",
        num_classes=NUM_CLASSES,
        pretrained=True,
        dropout_rate=dropout_rate
    )
    
    # Create trainer with sampled hyperparameters
    trainer = FineTuningTrainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        backbone_lr=backbone_lr,
        head_lr=head_lr,
        weight_decay=weight_decay,
        warmup_epochs=warmup_epochs,
        max_epochs=20,  # Shorter for search
        patience=5
    )
    
    history = trainer.fit()
    
    # Return best validation accuracy
    return max(history["val_acc"])
 
# Run hyperparameter search
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)
 
print("Best hyperparameters:")
for key, value in study.best_params.items():
    print(f"  {key}: {value}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

import torch
from torch.cuda.amp import GradScaler, autocast
from torch.utils.checkpoint import checkpoint_sequential
 
class MemoryEfficientFineTuning:
    """
    Memory optimization techniques for full fine-tuning of large models.
    
    Techniques demonstrated:
    1. Mixed Precision Training (FP16)
    2. Gradient Checkpointing
    3. Gradient Accumulation
    """
    
    def __init__(
        self,
        model: nn.Module,
        use_mixed_precision: bool = True,
        use_gradient_checkpointing: bool = False,
        gradient_accumulation_steps: int = 1
    ):
        self.model = model
        self.use_mixed_precision = use_mixed_precision
        self.gradient_accumulation_steps = gradient_accumulation_steps
        
        # Mixed precision training
        if use_mixed_precision:
            self.scaler = GradScaler()
        
        # Gradient checkpointing trades compute for memory
        if use_gradient_checkpointing:
            self._enable_gradient_checkpointing()
    
    def _enable_gradient_checkpointing(self):
        """Enable gradient checkpointing for memory savings."""
        # For models with defined checkpointing functions
        if hasattr(self.model.backbone, 'set_grad_checkpointing'):
            self.model.backbone.set_grad_checkpointing(True)
        elif hasattr(self.model.backbone, 'layer1'):  # ResNet-like
            # Wrap layer blocks with checkpointing
            def checkpoint_forward_hook(module, input, output):
                return checkpoint_sequential(module, 2, input)
            # This is a simplified example; production code
            # should use proper layer-wise checkpointing
    
    def training_step(self, batch, optimizer, batch_idx):
        """Single training step with memory optimizations."""
        inputs, targets = batch
        
        # Gradient accumulation: normalize loss by accumulation steps
        loss_scale = 1.0 / self.gradient_accumulation_steps
        
        if self.use_mixed_precision:
            # Mixed precision forward pass
            with autocast():
                outputs = self.model(inputs)
                loss = nn.functional.cross_entropy(outputs, targets)
                loss = loss * loss_scale
            
            # Scaled backward pass
            self.scaler.scale(loss).backward()
            
            # Only update on accumulation boundary
            if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
                self.scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.scaler.step(optimizer)
                self.scaler.update()
                optimizer.zero_grad()
        else:
            outputs = self.model(inputs)
            loss = nn.functional.cross_entropy(outputs, targets) * loss_scale
            loss.backward()
            
            if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad()
        
        return loss.item() / loss_scale
 
 
# Memory savings comparison
def estimate_memory_savings():
    """Compare memory usage across optimization techniques."""
    # Baseline: Full precision, no optimizations
    baseline_memory = "100%"
    
    optimizations = {
        "Mixed Precision (FP16)": "~60% (1.67x batch size)",
        "Gradient Checkpointing": "~70% (saves activation memory)",
        "Gradient Accumulation (4 steps)": "~50% (effective batch = 4x)",
        "All Combined": "~35-40% of baseline"
    }
    
    return optimizations

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

import torch
import numpy as np
from collections import defaultdict
from torch.utils.tensorboard import SummaryWriter
 
class FineTuningMonitor:
    """
    Comprehensive monitoring for full fine-tuning.
    
    Tracks:
    - Loss and accuracy curves
    - Gradient statistics per layer
    - Weight distance from initialization
    - Learning dynamics indicators
    """
    
    def __init__(
        self,
        model: nn.Module,
        log_dir: str = "./runs/fine_tuning"
    ):
        self.model = model
        self.writer = SummaryWriter(log_dir)
        
        # Store initial weights for distance computation
        self.initial_weights = {}
        for name, param in model.named_parameters():
            self.initial_weights[name] = param.data.clone()
        
        # Track gradient history
        self.gradient_history = defaultdict(list)
    
    def log_gradients(self, step: int):
        """Log gradient statistics for each layer."""
        grad_norms = {}
        
        for name, param in self.model.named_parameters():
            if param.grad is not None:
                grad_norm = param.grad.data.norm(2).item()
                grad_norms[name] = grad_norm
                
                # Log to TensorBoard
                self.writer.add_scalar(
                    f"gradients/{name}_norm",
                    grad_norm,
                    step
                )
                
                self.gradient_history[name].append(grad_norm)
        
        # Log summary statistics
        all_norms = list(grad_norms.values())
        self.writer.add_scalar("gradients/mean_norm", np.mean(all_norms), step)
        self.writer.add_scalar("gradients/max_norm", np.max(all_norms), step)
        self.writer.add_scalar("gradients/min_norm", np.min(all_norms), step)
    
    def log_weight_distance(self, step: int):
        """Measure how far weights have moved from initialization."""
        distances = {}
        
        for name, param in self.model.named_parameters():
            if name in self.initial_weights:
                distance = (param.data - self.initial_weights[name]).norm(2).item()
                relative_distance = distance / (self.initial_weights[name].norm(2).item() + 1e-8)
                
                distances[name] = relative_distance
                self.writer.add_scalar(
                    f"weight_distance/{name}",
                    relative_distance,
                    step
                )
        
        # Summary
        all_distances = list(distances.values())
        self.writer.add_scalar(
            "weight_distance/mean_relative",
            np.mean(all_distances),
            step
        )
        
        return distances
    
    def log_layer_learning_rate(self, optimizer, step: int):
        """Log effective learning rate per parameter group."""
        for i, group in enumerate(optimizer.param_groups):
            self.writer.add_scalar(
                f"learning_rate/group_{i}",
                group['lr'],
                step
            )
    
    def detect_training_issues(self, step: int) -> list:
        """Analyze metrics to detect common fine-tuning issues."""
        issues = []
        
        # Check for vanishing gradients
        for name, history in self.gradient_history.items():
            if len(history) > 10:
                recent_mean = np.mean(history[-10:])
                if recent_mean < 1e-7:
                    issues.append(
                        f"Vanishing gradients detected in {name}. "
                        f"Consider higher learning rate or different initialization."
                    )
        
        # Check for exploding gradients
        for name, history in self.gradient_history.items():
            if len(history) > 10:
                recent_max = np.max(history[-10:])
                if recent_max > 100:
                    issues.append(
                        f"Large gradients detected in {name} (max: {recent_max:.2f}). "
                        f"Consider gradient clipping or lower learning rate."
                    )
        
        # Check for catastrophic forgetting indicators
        distances = self.log_weight_distance(step)
        mean_distance = np.mean(list(distances.values()))
        if mean_distance > 1.0:  # Weights have moved more than 100% from initialization
            issues.append(
                f"Weights have moved significantly from pre-trained values "
                f"(mean relative distance: {mean_distance:.2f}). "
                f"Risk of catastrophic forgetting. Consider lower LR or regularization."
            )
        
        return issues
    
    def create_diagnostic_report(self) -> str:
        """Generate a comprehensive diagnostic report."""
        report = []
        report.append("=" * 60)
        report.append("FINE-TUNING DIAGNOSTIC REPORT")
        report.append("=" * 60)
        
        # Gradient statistics
        report.append("\nGRADIENT STATISTICS:")
        for name, history in self.gradient_history.items():
            if history:
                report.append(
                    f"  {name}: mean={np.mean(history):.6f}, "
                    f"std={np.std(history):.6f}, "
                    f"max={np.max(history):.6f}"
                )
        
        # Weight distances
        report.append("\nWEIGHT DISTANCE FROM INITIALIZATION:")
        distances = {}
        for name, param in self.model.named_parameters():
            if name in self.initial_weights:
                distance = (param.data - self.initial_weights[name]).norm(2).item()
                relative = distance / (self.initial_weights[name].norm(2).item() + 1e-8)
                distances[name] = relative
        
        for name, dist in sorted(distances.items(), key=lambda x: -x[1])[:10]:
            report.append(f"  {name}: {dist:.4f}")
        
        return "\n".join(report)