Transfer Learning for Computer Vision

1import torch
2import torch.nn as nn
3import torchvision.models as models
4from torchvision import transforms
5from torch.utils.data import DataLoader
6import numpy as np
7
8class FeatureExtractor:
9    """
10    Extract features from pretrained models for downstream tasks.
11    """
12
13    def __init__(self, model_name: str = 'resnet50', device: str = 'cuda'):
14        self.device = device if torch.cuda.is_available() else 'cpu'
15
16        # Load pretrained model
17        self.model = self._load_model(model_name)
18        self.model.to(self.device)
19        self.model.eval()
20
21        # Freeze all parameters
22        for param in self.model.parameters():
23            param.requires_grad = False
24
25        # Standard ImageNet preprocessing
26        self.preprocess = transforms.Compose([
27            transforms.Resize(256),
28            transforms.CenterCrop(224),
29            transforms.ToTensor(),
30            transforms.Normalize(
31                mean=[0.485, 0.456, 0.406],
32                std=[0.229, 0.224, 0.225]
33            )
34        ])
35
36    def _load_model(self, model_name: str) -> nn.Module:
37        """Load pretrained model and remove classification head."""
38        if model_name == 'resnet50':
39            model = models.resnet50(pretrained=True)
40            # Remove the final FC layer
41            model = nn.Sequential(*list(model.children())[:-1])
42            self.feature_dim = 2048
43
44        elif model_name == 'efficientnet_b0':
45            model = models.efficientnet_b0(pretrained=True)
46            model = nn.Sequential(*list(model.children())[:-1])
47            self.feature_dim = 1280
48
49        elif model_name == 'vit_b_16':
50            model = models.vit_b_16(pretrained=True)
51            # ViT returns class token, we need to extract it differently
52            self.feature_dim = 768
53            return model
54
55        else:
56            raise ValueError(f"Unknown model: {model_name}")
57
58        return model
59
60    @torch.no_grad()
61    def extract_features(self, images: torch.Tensor) -> np.ndarray:
62        """
63        Extract features from a batch of images.
64
65        Args:
66            images: Preprocessed images [B, C, H, W]
67
68        Returns:
69            Feature vectors [B, feature_dim]
70        """
71        images = images.to(self.device)
72        features = self.model(images)
73
74        # Flatten if needed (e.g., from [B, 2048, 1, 1] to [B, 2048])
75        features = features.squeeze()
76
77        return features.cpu().numpy()
78
79    def extract_from_dataloader(self, dataloader: DataLoader) -> tuple:
80        """
81        Extract features from entire dataset.
82
83        Returns:
84            features: [N, feature_dim]
85            labels: [N]
86        """
87        all_features = []
88        all_labels = []
89
90        for images, labels in dataloader:
91            features = self.extract_features(images)
92            all_features.append(features)
93            all_labels.extend(labels.numpy())
94
95        return np.vstack(all_features), np.array(all_labels)
96
97
98class LinearClassifier(nn.Module):
99    """
100    Simple linear classifier on top of extracted features.
101    """
102
103    def __init__(self, feature_dim: int, num_classes: int):
104        super().__init__()
105        self.classifier = nn.Sequential(
106            nn.Linear(feature_dim, 512),
107            nn.ReLU(),
108            nn.Dropout(0.5),
109            nn.Linear(512, num_classes)
110        )
111
112    def forward(self, x: torch.Tensor) -> torch.Tensor:
113        return self.classifier(x)
114
115
116# Example: Using features with sklearn
117def train_with_sklearn(features: np.ndarray, labels: np.ndarray):
118    """Train simple classifiers on extracted features."""
119    from sklearn.linear_model import LogisticRegression
120    from sklearn.svm import SVC
121    from sklearn.model_selection import train_test_split
122    from sklearn.metrics import accuracy_score
123
124    X_train, X_test, y_train, y_test = train_test_split(
125        features, labels, test_size=0.2, random_state=42
126    )
127
128    # Logistic Regression
129    lr = LogisticRegression(max_iter=1000)
130    lr.fit(X_train, y_train)
131    lr_acc = accuracy_score(y_test, lr.predict(X_test))
132
133    # SVM
134    svm = SVC(kernel='rbf')
135    svm.fit(X_train, y_train)
136    svm_acc = accuracy_score(y_test, svm.predict(X_test))
137
138    return {'logistic_regression': lr_acc, 'svm': svm_acc}
139
140
141# Usage example
142extractor = FeatureExtractor(model_name='resnet50')
143print(f"Feature dimension: {extractor.feature_dim}")
144print(f"Model on device: {extractor.device}")
145
146# Demo with random data
147demo_images = torch.randn(4, 3, 224, 224)
148features = extractor.extract_features(demo_images)
149print(f"Extracted features shape: {features.shape}")

1import torch
2import torch.nn as nn
3import torchvision.models as models
4from typing import List, Optional
5import math
6
7class FineTuner:
8    """
9    Comprehensive fine-tuning strategies for pretrained models.
10    """
11
12    def __init__(self, model_name: str = 'resnet50', num_classes: int = 10,
13                 device: str = 'cuda'):
14        self.device = device if torch.cuda.is_available() else 'cpu'
15        self.num_classes = num_classes
16
17        # Load base model
18        self.model = self._create_model(model_name)
19        self.model.to(self.device)
20
21    def _create_model(self, model_name: str) -> nn.Module:
22        """Create model with custom classification head."""
23        if model_name == 'resnet50':
24            model = models.resnet50(pretrained=True)
25            num_features = model.fc.in_features
26            model.fc = nn.Linear(num_features, self.num_classes)
27
28        elif model_name == 'efficientnet_b0':
29            model = models.efficientnet_b0(pretrained=True)
30            num_features = model.classifier[1].in_features
31            model.classifier = nn.Sequential(
32                nn.Dropout(0.2),
33                nn.Linear(num_features, self.num_classes)
34            )
35
36        return model
37
38    def freeze_layers(self, num_layers_to_freeze: int = None,
39                     freeze_bn: bool = True):
40        """
41        Freeze early layers of the network.
42
43        Args:
44            num_layers_to_freeze: Number of layers to freeze (from start)
45            freeze_bn: Whether to freeze batch normalization layers
46        """
47        if num_layers_to_freeze is None:
48            # Default: freeze all except final layer
49            for name, param in self.model.named_parameters():
50                if 'fc' not in name and 'classifier' not in name:
51                    param.requires_grad = False
52        else:
53            # Freeze specific number of layers
54            children = list(self.model.children())
55            for i, child in enumerate(children[:num_layers_to_freeze]):
56                for param in child.parameters():
57                    param.requires_grad = False
58
59        # Optionally freeze batch norm statistics
60        if freeze_bn:
61            for module in self.model.modules():
62                if isinstance(module, nn.BatchNorm2d):
63                    module.eval()
64
65    def gradual_unfreeze(self, epoch: int, unfreeze_schedule: List[int]):
66        """
67        Gradually unfreeze layers during training.
68
69        Args:
70            epoch: Current training epoch
71            unfreeze_schedule: List of epochs at which to unfreeze next block
72        """
73        children = list(self.model.children())
74        num_blocks = len(children)
75
76        # Determine how many blocks to unfreeze
77        blocks_to_unfreeze = sum(1 for e in unfreeze_schedule if epoch >= e)
78
79        # Unfreeze from the end (closest to output)
80        for i, child in enumerate(reversed(children)):
81            if i < blocks_to_unfreeze:
82                for param in child.parameters():
83                    param.requires_grad = True
84
85    def get_layer_groups(self) -> List[List[nn.Parameter]]:
86        """
87        Get parameter groups for discriminative learning rates.
88
89        Returns groups from earliest to latest layers.
90        """
91        # For ResNet: group by residual blocks
92        groups = []
93
94        # Early layers (stem)
95        stem_params = []
96        for name, param in self.model.named_parameters():
97            if 'layer1' not in name and 'layer2' not in name and                'layer3' not in name and 'layer4' not in name and                'fc' not in name:
98                stem_params.append(param)
99        groups.append(stem_params)
100
101        # Residual blocks
102        for layer_name in ['layer1', 'layer2', 'layer3', 'layer4']:
103            layer_params = []
104            for name, param in self.model.named_parameters():
105                if layer_name in name:
106                    layer_params.append(param)
107            if layer_params:
108                groups.append(layer_params)
109
110        # Classification head
111        head_params = []
112        for name, param in self.model.named_parameters():
113            if 'fc' in name or 'classifier' in name:
114                head_params.append(param)
115        groups.append(head_params)
116
117        return groups
118
119    def get_discriminative_lr_optimizer(self, base_lr: float = 1e-3,
120                                        lr_decay: float = 0.1):
121        """
122        Create optimizer with discriminative learning rates.
123
124        Earlier layers get smaller learning rates.
125        """
126        groups = self.get_layer_groups()
127
128        param_groups = []
129        for i, group_params in enumerate(groups):
130            # Earlier groups get smaller LR
131            lr = base_lr * (lr_decay ** (len(groups) - 1 - i))
132            param_groups.append({
133                'params': group_params,
134                'lr': lr
135            })
136
137        return torch.optim.Adam(param_groups)
138
139
140class WarmupCosineScheduler:
141    """
142    Learning rate scheduler with warmup and cosine annealing.
143    Recommended for fine-tuning.
144    """
145
146    def __init__(self, optimizer, warmup_epochs: int, total_epochs: int,
147                 warmup_lr: float = 1e-6, min_lr: float = 1e-6):
148        self.optimizer = optimizer
149        self.warmup_epochs = warmup_epochs
150        self.total_epochs = total_epochs
151        self.warmup_lr = warmup_lr
152        self.min_lr = min_lr
153
154        # Store base learning rates
155        self.base_lrs = [group['lr'] for group in optimizer.param_groups]
156
157    def step(self, epoch: int):
158        """Update learning rate based on epoch."""
159        if epoch < self.warmup_epochs:
160            # Linear warmup
161            alpha = epoch / self.warmup_epochs
162            for i, group in enumerate(self.optimizer.param_groups):
163                group['lr'] = self.warmup_lr + alpha * (self.base_lrs[i] - self.warmup_lr)
164        else:
165            # Cosine annealing
166            progress = (epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)
167            for i, group in enumerate(self.optimizer.param_groups):
168                group['lr'] = self.min_lr + 0.5 * (self.base_lrs[i] - self.min_lr) *                              (1 + math.cos(math.pi * progress))
169
170
171# Example usage
172finetuner = FineTuner(model_name='resnet50', num_classes=10)
173print(f"Total parameters: {sum(p.numel() for p in finetuner.model.parameters()):,}")
174
175# Freeze all but last layer
176finetuner.freeze_layers()
177trainable = sum(p.numel() for p in finetuner.model.parameters() if p.requires_grad)
178print(f"Trainable parameters after freezing: {trainable:,}")
179
180# Get discriminative LR optimizer
181optimizer = finetuner.get_discriminative_lr_optimizer(base_lr=1e-3, lr_decay=0.1)
182print(f"\nLearning rates per group:")
183for i, group in enumerate(optimizer.param_groups):
184    print(f"  Group {i}: {group['lr']:.6f}")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from typing import Tuple
5
6class DomainAdversarialNetwork(nn.Module):
7    """
8    Domain-Adversarial Neural Network (DANN).
9
10    Uses gradient reversal to learn domain-invariant features.
11    """
12
13    def __init__(self, feature_extractor: nn.Module, feature_dim: int,
14                 num_classes: int):
15        super().__init__()
16
17        self.feature_extractor = feature_extractor
18
19        # Task classifier (label predictor)
20        self.classifier = nn.Sequential(
21            nn.Linear(feature_dim, 256),
22            nn.ReLU(),
23            nn.Dropout(0.5),
24            nn.Linear(256, num_classes)
25        )
26
27        # Domain discriminator
28        self.domain_discriminator = nn.Sequential(
29            nn.Linear(feature_dim, 256),
30            nn.ReLU(),
31            nn.Dropout(0.5),
32            nn.Linear(256, 2)  # source vs target
33        )
34
35    def forward(self, x: torch.Tensor, alpha: float = 1.0
36               ) -> Tuple[torch.Tensor, torch.Tensor]:
37        """
38        Forward pass with gradient reversal for domain discriminator.
39
40        Args:
41            x: Input images
42            alpha: Gradient reversal strength
43
44        Returns:
45            class_output: Class predictions
46            domain_output: Domain predictions
47        """
48        features = self.feature_extractor(x)
49        features = features.view(features.size(0), -1)
50
51        # Class prediction (normal gradient flow)
52        class_output = self.classifier(features)
53
54        # Domain prediction (reversed gradient)
55        reversed_features = GradientReversal.apply(features, alpha)
56        domain_output = self.domain_discriminator(reversed_features)
57
58        return class_output, domain_output
59
60
61class GradientReversal(torch.autograd.Function):
62    """
63    Gradient Reversal Layer for domain adversarial training.
64    """
65
66    @staticmethod
67    def forward(ctx, x: torch.Tensor, alpha: float) -> torch.Tensor:
68        ctx.alpha = alpha
69        return x.view_as(x)
70
71    @staticmethod
72    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
73        return -ctx.alpha * grad_output, None
74
75
76class MMDLoss(nn.Module):
77    """
78    Maximum Mean Discrepancy loss for domain adaptation.
79
80    Measures distance between source and target distributions.
81    """
82
83    def __init__(self, kernel_type: str = 'rbf', kernel_mul: float = 2.0,
84                 kernel_num: int = 5):
85        super().__init__()
86        self.kernel_type = kernel_type
87        self.kernel_mul = kernel_mul
88        self.kernel_num = kernel_num
89
90    def gaussian_kernel(self, source: torch.Tensor, target: torch.Tensor,
91                       kernel_mul: float, kernel_num: int,
92                       fix_sigma: float = None) -> torch.Tensor:
93        """Compute Gaussian kernel matrix."""
94        n_samples = source.size(0) + target.size(0)
95        total = torch.cat([source, target], dim=0)
96
97        # Compute pairwise distances
98        total0 = total.unsqueeze(0).expand(
99            total.size(0), total.size(0), total.size(1)
100        )
101        total1 = total.unsqueeze(1).expand(
102            total.size(0), total.size(0), total.size(1)
103        )
104        L2_distance = ((total0 - total1) ** 2).sum(2)
105
106        # Bandwidth selection
107        if fix_sigma:
108            bandwidth = fix_sigma
109        else:
110            bandwidth = torch.sum(L2_distance) / (n_samples ** 2 - n_samples)
111
112        bandwidth /= kernel_mul ** (kernel_num // 2)
113        bandwidth_list = [bandwidth * (kernel_mul ** i) for i in range(kernel_num)]
114
115        # Sum of Gaussian kernels
116        kernel_val = [torch.exp(-L2_distance / bw) for bw in bandwidth_list]
117        return sum(kernel_val)
118
119    def forward(self, source: torch.Tensor,
120                target: torch.Tensor) -> torch.Tensor:
121        """
122        Compute MMD loss between source and target features.
123        """
124        batch_size = source.size(0)
125        kernels = self.gaussian_kernel(
126            source, target, self.kernel_mul, self.kernel_num
127        )
128
129        # Split kernel matrix
130        XX = kernels[:batch_size, :batch_size]
131        YY = kernels[batch_size:, batch_size:]
132        XY = kernels[:batch_size, batch_size:]
133
134        # MMD^2 = E[k(X,X')] + E[k(Y,Y')] - 2*E[k(X,Y)]
135        mmd = XX.mean() + YY.mean() - 2 * XY.mean()
136        return mmd
137
138
139class DomainAdaptationTrainer:
140    """Training pipeline for domain adaptation."""
141
142    def __init__(self, model: DomainAdversarialNetwork, device: str = 'cuda'):
143        self.model = model.to(device)
144        self.device = device
145
146        self.optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
147        self.class_criterion = nn.CrossEntropyLoss()
148        self.domain_criterion = nn.CrossEntropyLoss()
149
150    def train_step(self, source_data: Tuple[torch.Tensor, torch.Tensor],
151                   target_data: torch.Tensor, epoch: int, max_epochs: int):
152        """
153        Single training step with domain adversarial loss.
154        """
155        source_imgs, source_labels = source_data
156        target_imgs = target_data
157
158        source_imgs = source_imgs.to(self.device)
159        source_labels = source_labels.to(self.device)
160        target_imgs = target_imgs.to(self.device)
161
162        # Adaptation parameter (schedule from 0 to 1)
163        p = epoch / max_epochs
164        alpha = 2.0 / (1.0 + math.exp(-10 * p)) - 1
165
166        self.optimizer.zero_grad()
167
168        # Source forward pass
169        class_output, domain_output_source = self.model(source_imgs, alpha)
170
171        # Target forward pass
172        _, domain_output_target = self.model(target_imgs, alpha)
173
174        # Classification loss (only on source)
175        class_loss = self.class_criterion(class_output, source_labels)
176
177        # Domain loss
178        source_domain_labels = torch.zeros(source_imgs.size(0),
179                                          dtype=torch.long, device=self.device)
180        target_domain_labels = torch.ones(target_imgs.size(0),
181                                         dtype=torch.long, device=self.device)
182
183        domain_loss = self.domain_criterion(domain_output_source, source_domain_labels) +                      self.domain_criterion(domain_output_target, target_domain_labels)
184
185        # Total loss
186        loss = class_loss + domain_loss
187        loss.backward()
188        self.optimizer.step()
189
190        return {
191            'total_loss': loss.item(),
192            'class_loss': class_loss.item(),
193            'domain_loss': domain_loss.item(),
194            'alpha': alpha
195        }
196
197
198# Example
199import torchvision.models as models
200
201# Create feature extractor from pretrained ResNet
202backbone = models.resnet18(pretrained=True)
203feature_extractor = nn.Sequential(*list(backbone.children())[:-1])
204
205dann = DomainAdversarialNetwork(
206    feature_extractor=feature_extractor,
207    feature_dim=512,
208    num_classes=10
209)
210
211print(f"DANN parameters: {sum(p.numel() for p in dann.parameters()):,}")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from typing import Optional
5
6class KnowledgeDistillation:
7    """
8    Knowledge Distillation framework.
9
10    Transfer knowledge from a large teacher model to a smaller student.
11    """
12
13    def __init__(self, teacher: nn.Module, student: nn.Module,
14                 temperature: float = 4.0, alpha: float = 0.7,
15                 device: str = 'cuda'):
16        """
17        Args:
18            teacher: Pretrained teacher model
19            student: Student model to train
20            temperature: Softmax temperature (higher = softer distributions)
21            alpha: Weight for distillation loss (1-alpha for hard labels)
22        """
23        self.device = device if torch.cuda.is_available() else 'cpu'
24
25        self.teacher = teacher.to(self.device)
26        self.student = student.to(self.device)
27
28        # Freeze teacher
29        self.teacher.eval()
30        for param in self.teacher.parameters():
31            param.requires_grad = False
32
33        self.temperature = temperature
34        self.alpha = alpha
35
36        self.optimizer = torch.optim.Adam(self.student.parameters(), lr=1e-3)
37        self.hard_loss_fn = nn.CrossEntropyLoss()
38
39    def distillation_loss(self, student_logits: torch.Tensor,
40                         teacher_logits: torch.Tensor,
41                         labels: torch.Tensor) -> torch.Tensor:
42        """
43        Compute combined distillation and hard label loss.
44
45        Distillation loss: KL divergence between soft teacher and student outputs
46        Hard loss: Cross entropy with ground truth labels
47        """
48        T = self.temperature
49
50        # Soft targets from teacher
51        soft_teacher = F.softmax(teacher_logits / T, dim=1)
52
53        # Soft predictions from student
54        soft_student = F.log_softmax(student_logits / T, dim=1)
55
56        # KL divergence loss (scaled by T^2 as per original paper)
57        distill_loss = F.kl_div(soft_student, soft_teacher,
58                               reduction='batchmean') * (T ** 2)
59
60        # Hard label loss
61        hard_loss = self.hard_loss_fn(student_logits, labels)
62
63        # Combined loss
64        loss = self.alpha * distill_loss + (1 - self.alpha) * hard_loss
65
66        return loss, distill_loss, hard_loss
67
68    def train_step(self, images: torch.Tensor,
69                   labels: torch.Tensor) -> dict:
70        """Single training step."""
71        images = images.to(self.device)
72        labels = labels.to(self.device)
73
74        self.optimizer.zero_grad()
75
76        # Get teacher predictions (no gradient)
77        with torch.no_grad():
78            teacher_logits = self.teacher(images)
79
80        # Get student predictions
81        student_logits = self.student(images)
82
83        # Compute loss
84        loss, distill_loss, hard_loss = self.distillation_loss(
85            student_logits, teacher_logits, labels
86        )
87
88        loss.backward()
89        self.optimizer.step()
90
91        # Calculate accuracy
92        _, predicted = student_logits.max(1)
93        accuracy = (predicted == labels).float().mean()
94
95        return {
96            'total_loss': loss.item(),
97            'distill_loss': distill_loss.item(),
98            'hard_loss': hard_loss.item(),
99            'accuracy': accuracy.item()
100        }
101
102
103class FeatureDistillation(nn.Module):
104    """
105    Feature-level knowledge distillation.
106
107    Matches intermediate feature representations, not just outputs.
108    """
109
110    def __init__(self, teacher: nn.Module, student: nn.Module,
111                 hint_layers: list):
112        """
113        Args:
114            teacher: Teacher model
115            student: Student model
116            hint_layers: Names of layers to match features
117        """
118        super().__init__()
119        self.teacher = teacher
120        self.student = student
121        self.hint_layers = hint_layers
122
123        # Hooks to capture intermediate features
124        self.teacher_features = {}
125        self.student_features = {}
126
127        self._register_hooks()
128
129        # Adaptation layers to match dimensions
130        self.adaptation = nn.ModuleDict()
131
132    def _register_hooks(self):
133        """Register forward hooks to capture features."""
134        def get_teacher_hook(name):
135            def hook(module, input, output):
136                self.teacher_features[name] = output
137            return hook
138
139        def get_student_hook(name):
140            def hook(module, input, output):
141                self.student_features[name] = output
142            return hook
143
144        for name in self.hint_layers:
145            # Register hooks for both models
146            for n, m in self.teacher.named_modules():
147                if n == name:
148                    m.register_forward_hook(get_teacher_hook(name))
149
150            for n, m in self.student.named_modules():
151                if n == name:
152                    m.register_forward_hook(get_student_hook(name))
153
154    def feature_loss(self) -> torch.Tensor:
155        """Compute feature matching loss."""
156        total_loss = 0
157
158        for name in self.hint_layers:
159            teacher_feat = self.teacher_features.get(name)
160            student_feat = self.student_features.get(name)
161
162            if teacher_feat is not None and student_feat is not None:
163                # Adapt student features if dimensions don't match
164                if teacher_feat.shape != student_feat.shape:
165                    if name not in self.adaptation:
166                        # Create adaptation layer
167                        self.adaptation[name] = nn.Conv2d(
168                            student_feat.size(1),
169                            teacher_feat.size(1),
170                            kernel_size=1
171                        ).to(student_feat.device)
172                    student_feat = self.adaptation[name](student_feat)
173
174                # L2 loss between features
175                total_loss += F.mse_loss(student_feat, teacher_feat)
176
177        return total_loss
178
179
180class AttentionTransfer(nn.Module):
181    """
182    Attention Transfer for knowledge distillation.
183
184    Transfers attention maps (spatial importance) from teacher to student.
185    """
186
187    @staticmethod
188    def attention_map(features: torch.Tensor) -> torch.Tensor:
189        """
190        Compute attention map from feature maps.
191
192        Attention = sum over channels of squared activations.
193        """
194        # features: [B, C, H, W]
195        return torch.sum(features ** 2, dim=1)  # [B, H, W]
196
197    @staticmethod
198    def attention_loss(student_features: torch.Tensor,
199                      teacher_features: torch.Tensor) -> torch.Tensor:
200        """Compute attention transfer loss."""
201        # Get attention maps
202        student_att = AttentionTransfer.attention_map(student_features)
203        teacher_att = AttentionTransfer.attention_map(teacher_features)
204
205        # Normalize attention maps
206        student_att = student_att / (student_att.sum(dim=(1, 2), keepdim=True) + 1e-8)
207        teacher_att = teacher_att / (teacher_att.sum(dim=(1, 2), keepdim=True) + 1e-8)
208
209        # Resize if needed
210        if student_att.shape != teacher_att.shape:
211            student_att = F.interpolate(
212                student_att.unsqueeze(1),
213                size=teacher_att.shape[1:],
214                mode='bilinear',
215                align_corners=False
216            ).squeeze(1)
217
218        # L2 loss between attention maps
219        return F.mse_loss(student_att, teacher_att)
220
221
222# Example: Create teacher and student models
223import torchvision.models as models
224
225# Large teacher
226teacher = models.resnet50(pretrained=True)
227
228# Smaller student
229student = models.resnet18(pretrained=False)
230
231# Initialize distillation
232distiller = KnowledgeDistillation(
233    teacher=teacher,
234    student=student,
235    temperature=4.0,
236    alpha=0.7
237)
238
239# Demo training step
240images = torch.randn(8, 3, 224, 224)
241labels = torch.randint(0, 1000, (8,))
242
243metrics = distiller.train_step(images, labels)
244print("Distillation metrics:")
245for k, v in metrics.items():
246    print(f"  {k}: {v:.4f}")