import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader import numpy as np # -------------------- Data: modulo-balanced batches -------------------- def get_modulo_batches(dataset, batch_size=100, num_classes=10): """Return list of tensors, each containing indices for one balanced batch.""" class_indices = [[] for _ in range(num_classes)] for idx, (_, label) in enumerate(dataset): class_indices[label].append(idx) per_class = batch_size // num_classes num_batches = len(dataset) // batch_size batch_indices = [[] for _ in range(num_batches)] for c in range(num_classes): indices = class_indices[c] np.random.shuffle(indices) for i, idx in enumerate(indices): batch_id = i % num_batches if len(batch_indices[batch_id]) < batch_size: batch_indices[batch_id].append(idx) batch_indices = [torch.tensor(b) for b in batch_indices if len(b) == batch_size] return batch_indices def get_full_cifar10_train(): transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) return torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform) # -------------------- Model -------------------- class HybridNet(nn.Module): def __init__(self): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1) ) self.fc_in = 128 self.classifier = nn.Linear(self.fc_in, 10) def forward(self, x): feat = self.features(x).view(x.size(0), -1) return self.classifier(feat) def get_adaptive_params(self): return list(self.features.parameters()) def get_stable_params(self): return list(self.classifier.parameters()) class DummySamples(nn.Module): def __init__(self, num_classes=10, latent_dim=128): super().__init__() self.dummies = nn.Parameter(torch.randn(num_classes, latent_dim) * 0.1) def forward(self): return self.dummies # -------------------- Training -------------------- def train_hybrid_full(model, dummy_samples, train_dataset, num_epochs=1, batch_size=100, inner_steps=10): batches = get_modulo_batches(train_dataset, batch_size=batch_size) print(f"Total batches: {len(batches)}") opt_adaptive = optim.Adam(model.get_adaptive_params(), lr=1e-3) opt_stable = optim.Adam(model.get_stable_params(), lr=1e-3) opt_dummy = optim.Adam(dummy_samples.parameters(), lr=1e-2) model.train() for epoch in range(num_epochs): perm = torch.randperm(len(batches)) prev_batch_real = None prev_batch_labels = None for batch_idx in perm: indices = batches[batch_idx] real_images = torch.stack([train_dataset[i][0] for i in indices]) real_labels = torch.tensor([train_dataset[i][1] for i in indices]).long() # ---- Create imbalanced version of the same real batch ---- num_real = len(real_images) # Class imbalance probabilities (over‑represent first 5 classes) prob = torch.ones(10) * 0.1 prob[:5] = 0.15 prob[5:] = 0.05 prob = prob / prob.sum() sampled_indices = torch.multinomial(prob[real_labels], num_real, replacement=True) imba_real_images = real_images[sampled_indices] imba_real_labels = real_labels[sampled_indices] # ---- INNER LOOP: update adaptive weights only ---- for _ in range(inner_steps): feat_imba = model.features(imba_real_images).view(imba_real_images.size(0), -1) dummy_latents = dummy_samples() # [10, latent_dim] combined_feat = torch.cat([feat_imba, dummy_latents], dim=0) dummy_targets = torch.arange(10, device=real_labels.device) combined_targets = torch.cat([imba_real_labels, dummy_targets]) logits_combined = model.classifier(combined_feat) loss_adaptive = nn.CrossEntropyLoss()(logits_combined, combined_targets) opt_adaptive.zero_grad() loss_adaptive.backward() opt_adaptive.step() print(loss_adaptive.item()) # ---- STABLE UPDATE: use balanced real batch, detach features ---- # Forward pass after inner loop, but detach so gradients don't go to adaptive with torch.no_grad(): feat_balanced = model.features(real_images).view(real_images.size(0), -1) # We still need gradients for the classifier, so reattach with requires_grad feat_balanced = feat_balanced.detach().requires_grad_(True) logits_balanced = model.classifier(feat_balanced) loss_stable = nn.CrossEntropyLoss()(logits_balanced, real_labels) opt_stable.zero_grad() loss_stable.backward() opt_stable.step() # ---- META UPDATE for dummy samples (improve performance on previous batch) ---- if prev_batch_real is not None: # Compute loss on previous batch using current model feat_prev = model.features(prev_batch_real).view(prev_batch_real.size(0), -1) logits_prev = model.classifier(feat_prev) meta_loss = nn.CrossEntropyLoss()(logits_prev, prev_batch_labels) opt_dummy.zero_grad() meta_loss.backward() opt_dummy.step() # Store current batch for next meta step prev_batch_real = real_images prev_batch_labels = real_labels print(f"Epoch {epoch+1}/{num_epochs} completed") # -------------------- Testing -------------------- def test(model): transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) test_set = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform) test_loader = DataLoader(test_set, batch_size=100, shuffle=False) model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: outputs = model(images) _, predicted = torch.max(outputs, 1) total += labels.size(0) correct += (predicted == labels).sum().item() acc = 100 * correct / total print(f"Test Accuracy: {acc:.2f}%") return acc # -------------------- Main -------------------- if __name__ == "__main__": train_set = get_full_cifar10_train() model = HybridNet() dummy_samples = DummySamples(num_classes=10, latent_dim=model.fc_in) train_hybrid_full(model, dummy_samples, train_set, num_epochs=1, batch_size=100, inner_steps=10) test(model) train_hybrid_full(model, dummy_samples, train_set, num_epochs=1, batch_size=100, inner_steps=10) test(model)