Feature-Based Transfer

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
158
159
160
161
162
163
164
165

import torch
import torch.nn as nn
from torchvision import models
from typing import Dict, List
 
class MultiLayerFeatureExtractor(nn.Module):
    """
    Extract features from multiple layers of a pre-trained model.
    
    Supports flexible layer selection and feature combination strategies.
    """
    def __init__(
        self, 
        model_name: str = "resnet50",
        layers: List[str] = ["layer3", "layer4", "avgpool"],
        combine_strategy: str = "concat"  # "concat", "mean", "learned"
    ):
        super().__init__()
        
        # Load pre-trained model
        base_model = getattr(models, model_name)(
            weights=getattr(models, f"{model_name.upper()}_Weights").IMAGENET1K_V2
        )
        
        # Freeze all parameters
        for param in base_model.parameters():
            param.requires_grad = False
        
        # Store model components for hook access
        self.model = base_model
        self.layers = layers
        self.combine_strategy = combine_strategy
        
        # Storage for intermediate features
        self.features: Dict[str, torch.Tensor] = {}
        
        # Register forward hooks on specified layers
        for name in layers:
            layer = dict(base_model.named_modules())[name]
            layer.register_forward_hook(self._get_hook(name))
        
        # Compute output dimension for downstream heads
        self._compute_output_dim()
    
    def _get_hook(self, name: str):
        def hook(module, input, output):
            # Flatten spatial dimensions if present
            if output.dim() == 4:  # B x C x H x W
                output = output.flatten(start_dim=2).mean(dim=2)  # Global avg pool
            self.features[name] = output
        return hook
    
    def _compute_output_dim(self):
        """Determine output dimension by doing a dummy forward pass."""
        with torch.no_grad():
            dummy = torch.randn(1, 3, 224, 224)
            self.forward(dummy)
            
            if self.combine_strategy == "concat":
                self.output_dim = sum(
                    self.features[l].shape[1] for l in self.layers
                )
            else:
                # Assumes all layers have same dim (may need projection)
                self.output_dim = self.features[self.layers[0]].shape[1]
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Clear previous features
        self.features = {}
        
        # Forward through full model (hooks capture intermediate)
        _ = self.model(x)
        
        # Combine features according to strategy
        feature_list = [self.features[l] for l in self.layers]
        
        if self.combine_strategy == "concat":
            return torch.cat(feature_list, dim=1)
        elif self.combine_strategy == "mean":
            # Stack and average (assumes same dimension)
            return torch.stack(feature_list, dim=0).mean(dim=0)
        else:
            raise ValueError(f"Unknown strategy: {self.combine_strategy}")
 
 
class FrozenFeatureClassifier(nn.Module):
    """
    Complete frozen feature classifier: extractor + trainable head.
    """
    def __init__(
        self,
        num_classes: int,
        extractor: MultiLayerFeatureExtractor = None,
        head_type: str = "linear",  # "linear", "mlp", "attention"
        hidden_dim: int = 512,
        dropout: float = 0.5
    ):
        super().__init__()
        
        # Create default extractor if none provided
        self.extractor = extractor or MultiLayerFeatureExtractor()
        input_dim = self.extractor.output_dim
        
        # Build classification head (this is what we train)
        if head_type == "linear":
            self.head = nn.Linear(input_dim, num_classes)
        elif head_type == "mlp":
            self.head = nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(hidden_dim, num_classes)
            )
        elif head_type == "attention":
            # Self-attention over feature dimensions
            self.head = nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.MultiheadAttention(hidden_dim, num_heads=4, dropout=dropout),
                nn.Linear(hidden_dim, num_classes)
            )
        
        # Initialize head with good defaults
        self._init_head()
    
    def _init_head(self):
        for m in self.head.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.extractor(x)
        return self.head(features)
    
    def get_trainable_params(self):
        """Return only the parameters that should be trained."""
        return self.head.parameters()
 
 
# Usage example
if __name__ == "__main__":
    # Create classifier
    classifier = FrozenFeatureClassifier(
        num_classes=10,
        head_type="mlp"
    )
    
    # Only train the head!
    optimizer = torch.optim.Adam(
        classifier.get_trainable_params(), 
        lr=0.001
    )
    
    # Dummy batch
    images = torch.randn(8, 3, 224, 224)
    labels = torch.randint(0, 10, (8,))
    
    # Forward pass
    logits = classifier(images)
    loss = nn.CrossEntropyLoss()(logits, labels)
    
    print(f"Feature dimension: {classifier.extractor.output_dim}")
    print(f"Output shape: {logits.shape}")
    print(f"Loss: {loss.item():.4f}")

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
158
159
160
161
162
163
164
165
166
167
168
169
170
171

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional
 
class LinearProbe(nn.Module):
    """
    Standard linear probe for frozen feature evaluation.
    Includes optional normalization and temperature scaling.
    """
    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        normalize: bool = True,
        temperature: float = 1.0
    ):
        super().__init__()
        self.normalize = normalize
        self.temperature = temperature
        self.classifier = nn.Linear(input_dim, num_classes)
        
        # Initialize with scaling aware of normalization
        nn.init.normal_(self.classifier.weight, std=0.01)
        nn.init.zeros_(self.classifier.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.normalize:
            x = F.normalize(x, p=2, dim=1)
        logits = self.classifier(x) / self.temperature
        return logits
 
 
class MLPHead(nn.Module):
    """
    Flexible MLP head with configurable architecture.
    """
    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        hidden_dims: list = [512],
        activation: str = "relu",
        dropout: float = 0.5,
        batch_norm: bool = True
    ):
        super().__init__()
        
        layers = []
        prev_dim = input_dim
        
        # Activation function selection
        act_fn = {
            "relu": nn.ReLU,
            "gelu": nn.GELU,
            "silu": nn.SiLU
        }[activation]
        
        # Build hidden layers
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            if batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(act_fn())
            layers.append(nn.Dropout(dropout))
            prev_dim = hidden_dim
        
        # Final classification layer
        layers.append(nn.Linear(prev_dim, num_classes))
        
        self.mlp = nn.Sequential(*layers)
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.mlp(x)
 
 
class PrototypicalHead(nn.Module):
    """
    Prototype-based classification head.
    Learns class prototypes and classifies by nearest prototype.
    Particularly effective for few-shot scenarios.
    """
    def __init__(
        self, 
        input_dim: int, 
        num_classes: int,
        metric: str = "euclidean"  # or "cosine"
    ):
        super().__init__()
        self.metric = metric
        
        # Learnable prototypes (one per class)
        self.prototypes = nn.Parameter(
            torch.randn(num_classes, input_dim)
        )
        nn.init.xavier_uniform_(self.prototypes)
        
        # Optional learnable temperature
        self.temperature = nn.Parameter(torch.tensor(1.0))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.metric == "euclidean":
            # Negative squared distance (higher = closer)
            dists = -torch.cdist(x, self.prototypes).pow(2)
        elif self.metric == "cosine":
            # Cosine similarity
            x_norm = F.normalize(x, dim=1)
            p_norm = F.normalize(self.prototypes, dim=1)
            dists = x_norm @ p_norm.T
        
        return dists / self.temperature
 
 
def select_head(
    input_dim: int,
    num_classes: int,
    target_data_size: int,
    task_type: str = "classification"
) -> nn.Module:
    """
    Heuristic head selection based on data availability.
    
    Args:
        input_dim: Dimension of frozen features
        num_classes: Number of output classes
        target_data_size: Number of labeled examples available
        task_type: "classification" or "regression"
    
    Returns:
        Appropriately-sized classification head
    """
    samples_per_class = target_data_size / num_classes
    
    if samples_per_class < 10:
        # Very few examples: use prototypical head
        print("Selected: PrototypicalHead (few-shot regime)")
        return PrototypicalHead(input_dim, num_classes, metric="cosine")
    
    elif samples_per_class < 50:
        # Limited data: use linear probe
        print("Selected: LinearProbe (limited data regime)")
        return LinearProbe(input_dim, num_classes, normalize=True)
    
    elif samples_per_class < 200:
        # Moderate data: small MLP
        print("Selected: MLPHead with 1 hidden layer")
        return MLPHead(
            input_dim, num_classes, 
            hidden_dims=[512], 
            dropout=0.5
        )
    
    else:
        # Ample data: larger MLP
        print("Selected: MLPHead with 2 hidden layers")
        return MLPHead(
            input_dim, num_classes,
            hidden_dims=[1024, 512],
            dropout=0.3
        )

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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
from pathlib import Path
from tqdm import tqdm
from typing import Tuple, Optional
import h5py
 
class FeatureCacher:
    """
    Efficient feature extraction and caching for frozen transfer.
    
    Supports multiple storage backends and streaming extraction.
    """
    def __init__(
        self,
        encoder: nn.Module,
        cache_dir: str = "./feature_cache",
        device: str = "cuda",
        dtype: torch.dtype = torch.float16  # Half precision saves 50% storage
    ):
        self.encoder = encoder.to(device).eval()
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.device = device
        self.dtype = dtype
        
        # Freeze encoder
        for param in self.encoder.parameters():
            param.requires_grad = False
    
    @torch.no_grad()
    def extract_and_cache(
        self,
        dataloader: DataLoader,
        cache_name: str,
        backend: str = "hdf5"  # "hdf5", "numpy", "torch"
    ) -> Path:
        """
        Extract features from all data and save to disk.
        
        Args:
            dataloader: DataLoader yielding (images, labels)
            cache_name: Name for the cached features file
            backend: Storage format
        
        Returns:
            Path to cached features
        """
        all_features = []
        all_labels = []
        
        print(f"Extracting features from {len(dataloader.dataset)} samples...")
        
        for images, labels in tqdm(dataloader, desc="Extracting"):
            images = images.to(self.device)
            
            # Extract features
            features = self.encoder(images)
            
            # Flatten if needed (e.g., from spatial features)
            if features.dim() > 2:
                features = features.flatten(start_dim=1)
            
            # Convert to target dtype for storage efficiency
            features = features.to(self.dtype).cpu()
            
            all_features.append(features)
            all_labels.append(labels)
        
        # Concatenate all batches
        features_tensor = torch.cat(all_features, dim=0)
        labels_tensor = torch.cat(all_labels, dim=0)
        
        # Save based on backend
        cache_path = self.cache_dir / f"{cache_name}.{backend}"
        
        if backend == "hdf5":
            with h5py.File(cache_path, 'w') as f:
                f.create_dataset('features', data=features_tensor.numpy(),
                               compression='gzip', compression_opts=4)
                f.create_dataset('labels', data=labels_tensor.numpy())
        
        elif backend == "numpy":
            np.savez_compressed(
                cache_path.with_suffix('.npz'),
                features=features_tensor.numpy(),
                labels=labels_tensor.numpy()
            )
            cache_path = cache_path.with_suffix('.npz')
        
        elif backend == "torch":
            torch.save({
                'features': features_tensor,
                'labels': labels_tensor
            }, cache_path.with_suffix('.pt'))
            cache_path = cache_path.with_suffix('.pt')
        
        print(f"Cached {len(features_tensor)} samples to {cache_path}")
        print(f"Feature shape: {features_tensor.shape}")
        print(f"Storage size: {cache_path.stat().st_size / 1e6:.1f} MB")
        
        return cache_path
    
    @staticmethod
    def load_cached_features(
        cache_path: Path,
        device: str = "cpu"
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Load cached features back into memory."""
        
        suffix = cache_path.suffix
        
        if suffix == '.hdf5' or suffix == '.h5':
            with h5py.File(cache_path, 'r') as f:
                features = torch.tensor(f['features'][:], dtype=torch.float32)
                labels = torch.tensor(f['labels'][:], dtype=torch.long)
        
        elif suffix == '.npz':
            data = np.load(cache_path)
            features = torch.tensor(data['features'], dtype=torch.float32)
            labels = torch.tensor(data['labels'], dtype=torch.long)
        
        elif suffix == '.pt':
            data = torch.load(cache_path)
            features = data['features'].float()
            labels = data['labels']
        
        return features.to(device), labels.to(device)
 
 
class CachedFeatureDataset(torch.utils.data.Dataset):
    """
    Dataset that loads features from disk lazily.
    Useful for very large datasets that don't fit in memory.
    """
    def __init__(self, cache_path: Path):
        self.cache_path = cache_path
        
        # Open file handle for lazy loading
        self.h5_file = h5py.File(cache_path, 'r')
        self.features = self.h5_file['features']
        self.labels = self.h5_file['labels']
    
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        feature = torch.tensor(self.features[idx], dtype=torch.float32)
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        return feature, label
    
    def __del__(self):
        self.h5_file.close()
 
 
def train_on_cached_features(
    features: torch.Tensor,
    labels: torch.Tensor,
    head: nn.Module,
    num_epochs: int = 100,
    batch_size: int = 256,
    lr: float = 0.01,
    weight_decay: float = 1e-4,
    device: str = "cuda"
) -> nn.Module:
    """
    Train a classification head on pre-cached features.
    
    This is extremely fast since there's no encoder forward pass.
    """
    head = head.to(device)
    
    # Create simple dataset from tensors
    dataset = TensorDataset(features, labels)
    loader = DataLoader(
        dataset, 
        batch_size=batch_size, 
        shuffle=True,
        pin_memory=True
    )
    
    # Optimizer
    optimizer = torch.optim.AdamW(
        head.parameters(), 
        lr=lr, 
        weight_decay=weight_decay
    )
    
    # Cosine schedule
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=num_epochs
    )
    
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(num_epochs):
        head.train()
        total_loss = 0
        correct = 0
        total = 0
        
        for batch_features, batch_labels in loader:
            batch_features = batch_features.to(device)
            batch_labels = batch_labels.to(device)
            
            optimizer.zero_grad()
            logits = head(batch_features)
            loss = criterion(logits, batch_labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            correct += (logits.argmax(1) == batch_labels).sum().item()
            total += len(batch_labels)
        
        scheduler.step()
        
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}: Loss={total_loss/len(loader):.4f}, "
                  f"Acc={100*correct/total:.2f}%")
    
    return head