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

import torch
import numpy as np
from typing import Dict, List, Tuple
from dataclasses import dataclass
 
@dataclass
class FewShotEpisode:
    """
    A complete few-shot learning episode.
    
    Attributes:
        n_way: Number of classes in this episode
        k_shot: Number of support examples per class
        q_query: Number of query examples per class
        support_images: Tensor of shape [n_way * k_shot, C, H, W]
        support_labels: Tensor of shape [n_way * k_shot], values in [0, n_way)
        query_images: Tensor of shape [n_way * q_query, C, H, W]
        query_labels: Tensor of shape [n_way * q_query], values in [0, n_way)
        original_classes: Original class IDs before relabeling (for debugging)
    """
    n_way: int
    k_shot: int
    q_query: int
    support_images: torch.Tensor
    support_labels: torch.Tensor
    query_images: torch.Tensor
    query_labels: torch.Tensor
    original_classes: List[int]
    
    def __post_init__(self):
        # Validate dimensions
        assert self.support_images.shape[0] == self.n_way * self.k_shot
        assert self.query_images.shape[0] == self.n_way * self.q_query
        assert len(set(self.support_labels.tolist())) == self.n_way
        assert len(set(self.query_labels.tolist())) == self.n_way
    
    @property
    def support_size(self) -> int:
        return self.n_way * self.k_shot
    
    @property
    def query_size(self) -> int:
        return self.n_way * self.q_query
    
    def to(self, device: torch.device) -> 'FewShotEpisode':
        """Move tensors to specified device."""
        return FewShotEpisode(
            n_way=self.n_way,
            k_shot=self.k_shot,
            q_query=self.q_query,
            support_images=self.support_images.to(device),
            support_labels=self.support_labels.to(device),
            query_images=self.query_images.to(device),
            query_labels=self.query_labels.to(device),
            original_classes=self.original_classes
        )
 
class FewShotEpisodeSampler:
    """
    Samples few-shot episodes from a dataset.
    
    Key principle: Each episode contains novel class combinations,
    with fresh support/query splits. The meta-learner sees many
    such episodes during training, learning to generalize across them.
    """
    
    def __init__(
        self, 
        dataset: Dict[int, List[torch.Tensor]],  # class_id -> list of images
        n_way: int = 5,
        k_shot: int = 5,
        q_query: int = 15
    ):
        self.dataset = dataset  # Pre-organized by class
        self.all_classes = list(dataset.keys())
        self.n_way = n_way
        self.k_shot = k_shot
        self.q_query = q_query
        
        # Validate sufficient examples per class
        min_examples = k_shot + q_query
        for cls, examples in dataset.items():
            if len(examples) < min_examples:
                raise ValueError(
                    f"Class {cls} has {len(examples)} examples, "
                    f"need at least {min_examples} for {k_shot}-shot + {q_query} queries"
                )
    
    def sample_episode(self) -> FewShotEpisode:
        """
        Sample a complete episode.
        
        Steps:
        1. Randomly select n_way classes
        2. For each class, sample k_shot + q_query examples
        3. Split into support and query
        4. Relabel classes to [0, n_way)
        5. Shuffle support and query separately
        """
        # Step 1: Select classes
        selected_classes = np.random.choice(
            self.all_classes, 
            size=self.n_way, 
            replace=False
        )
        
        support_images, support_labels = [], []
        query_images, query_labels = [], []
        
        # Step 2-4: Sample and split for each class
        for new_label, original_class in enumerate(selected_classes):
            class_examples = self.dataset[original_class]
            
            # Random selection without replacement
            indices = np.random.choice(
                len(class_examples),
                size=self.k_shot + self.q_query,
                replace=False
            )
            
            selected_examples = [class_examples[i] for i in indices]
            
            # Split
            support_examples = selected_examples[:self.k_shot]
            query_examples = selected_examples[self.k_shot:]
            
            # Add to lists with relabeled class IDs
            support_images.extend(support_examples)
            support_labels.extend([new_label] * self.k_shot)
            
            query_images.extend(query_examples)
            query_labels.extend([new_label] * self.q_query)
        
        # Step 5: Shuffle to remove class ordering
        support_perm = np.random.permutation(len(support_labels))
        query_perm = np.random.permutation(len(query_labels))
        
        return FewShotEpisode(
            n_way=self.n_way,
            k_shot=self.k_shot,
            q_query=self.q_query,
            support_images=torch.stack([support_images[i] for i in support_perm]),
            support_labels=torch.tensor([support_labels[i] for i in support_perm]),
            query_images=torch.stack([query_images[i] for i in query_perm]),
            query_labels=torch.tensor([query_labels[i] for i in query_perm]),
            original_classes=list(selected_classes)
        )

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

import numpy as np
from typing import Tuple, List
import torch
from scipy import stats
 
def evaluate_few_shot(
    model,
    episode_sampler,
    n_episodes: int = 600,
    confidence: float = 0.95
) -> Tuple[float, float, List[float]]:
    """
    Standard few-shot evaluation protocol.
    
    Args:
        model: Few-shot classifier
        episode_sampler: Samples (support, query) episodes
        n_episodes: Number of test episodes
        confidence: Confidence level for interval
    
    Returns:
        mean_accuracy: Average accuracy across episodes
        confidence_interval: Half-width of confidence interval
        all_accuracies: Per-episode accuracies for analysis
    """
    model.eval()
    accuracies = []
    
    with torch.no_grad():
        for episode_idx in range(n_episodes):
            episode = episode_sampler.sample_episode()
            
            # Model adapts to support set and predicts on query
            predictions = model(
                support_x=episode.support_images,
                support_y=episode.support_labels,
                query_x=episode.query_images
            )
            
            # Compute accuracy for this episode
            correct = (predictions == episode.query_labels).float()
            episode_accuracy = correct.mean().item()
            accuracies.append(episode_accuracy)
            
            if (episode_idx + 1) % 100 == 0:
                running_mean = np.mean(accuracies)
                running_std = np.std(accuracies)
                print(f"Episode {episode_idx + 1}/{n_episodes}: "
                      f"Accuracy = {running_mean:.2%} ± {running_std:.4f}")
    
    # Compute statistics
    mean_acc = np.mean(accuracies)
    std_acc = np.std(accuracies)
    
    # Confidence interval using t-distribution for small samples
    # For large n_episodes, this approximates the normal distribution
    t_value = stats.t.ppf((1 + confidence) / 2, n_episodes - 1)
    ci_half_width = t_value * (std_acc / np.sqrt(n_episodes))
    
    print(f"\nFinal Results ({n_episodes} episodes):")
    print(f"  Mean Accuracy: {mean_acc:.2%}")
    print(f"  Std Deviation: {std_acc:.4f}")
    print(f"  {int(confidence * 100)}% CI: [{mean_acc - ci_half_width:.2%}, {mean_acc + ci_half_width:.2%}]")
    
    return mean_acc, ci_half_width, accuracies
 
def compare_methods(
    results_a: Tuple[float, float],
    results_b: Tuple[float, float],
    n_episodes: int
) -> bool:
    """
    Test if method A is significantly better than method B.
    
    Uses a two-sample t-test to determine significance.
    Returns True if A is significantly better at p < 0.05.
    """
    mean_a, ci_a = results_a
    mean_b, ci_b = results_b
    
    # If confidence intervals don't overlap, difference is significant
    # This is a conservative test (overlapping CIs can still be significant)
    overlap = (mean_a - ci_a < mean_b + ci_b) and (mean_b - ci_b < mean_a + ci_a)
    
    if not overlap:
        return mean_a > mean_b  
    
    # Proper statistical test would require per-episode accuracies
    # This is a simplification
    return False
 
# Metric variations for specific scenarios
def evaluate_imbalanced_few_shot(
    model,
    episode_sampler,
    n_episodes: int = 600
) -> dict:
    """
    Evaluation for imbalanced few-shot scenarios.
    
    In realistic settings, classes may have different numbers of
    support examples or query examples. Report class-balanced metrics.
    """
    model.eval()
    all_predictions = []
    all_labels = []
    all_n_ways = []
    per_class_accuracies = []
    
    for _ in range(n_episodes):
        episode = episode_sampler.sample_variable_episode()
        
        predictions = model(
            support_x=episode.support_images,
            support_y=episode.support_labels,
            query_x=episode.query_images
        )
        
        # Per-class accuracy (balanced across classes)
        n_way = episode.n_way
        class_accs = []
        for c in range(n_way):
            mask = episode.query_labels == c
            if mask.sum() > 0:
                class_acc = (predictions[mask] == c).float().mean().item()
                class_accs.append(class_acc)
        
        # Macro-averaged (balanced) accuracy
        balanced_acc = np.mean(class_accs) if class_accs else 0.0
        per_class_accuracies.append(balanced_acc)
    
    return {
        'balanced_accuracy': np.mean(per_class_accuracies),
        'balanced_std': np.std(per_class_accuracies),
    }