import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader, Subset from torchvision import datasets, transforms import numpy as np from collections import defaultdict # ------------------------------- # 1. Graph neural network layer # ------------------------------- class GCNLayer(nn.Module): """Simple graph convolution: aggregation + linear transform.""" def __init__(self, in_dim, out_dim): super().__init__() self.linear = nn.Linear(in_dim, out_dim) def forward(self, x, adj_norm): # x: [N, in_dim], adj_norm: [N, N] row-normalised adjacency x = self.linear(x) # [N, out_dim] x = torch.matmul(adj_norm, x) # aggregate neighbours return x # ------------------------------------------- # 2. Per‑class graph network & classifier # ------------------------------------------- class ClassGraphClassifier(nn.Module): def __init__(self, feature_extractor, feat_dim, gcn_hidden=64, gcn_layers=2, k_neighbors=5): """ Args: feature_extractor: a nn.Module that maps images to feature vectors. feat_dim: dimension of feature vectors. gcn_hidden: hidden size of GCN layers. gcn_layers: number of GCN layers. k_neighbors: k for k‑NN graph construction (within each class support set). """ super().__init__() self.feature_extractor = feature_extractor self.feat_dim = feat_dim self.k = k_neighbors # Build shared GCN (could also be separate per class – here we share weights) gcn = [] in_dim = feat_dim for i in range(gcn_layers): out_dim = gcn_hidden if i < gcn_layers-1 else feat_dim # keep final dim = feat_dim gcn.append(GCNLayer(in_dim, out_dim)) in_dim = out_dim self.gcn = nn.Sequential(*gcn) def build_graph(self, features): """ Build a k‑NN graph (with self‑loops) and return row‑normalised adjacency. features: [N, D] returns: adj_norm [N, N] """ N = features.size(0) k = min(self.k, N) # avoid k > N # compute pairwise distances dist = torch.cdist(features, features) # [N, N] _, idx = dist.topk(k, dim=1, largest=False) # [N, k] indices of nearest neighbours adj = torch.zeros(N, N, device=features.device) adj.scatter_(1, idx, 1.0) # fill ones for neighbours adj = adj + torch.eye(N, device=features.device) # self‑loops # row normalisation degree = adj.sum(dim=1, keepdim=True).clamp(min=1) # avoid division by zero adj_norm = adj / degree return adj_norm def forward(self, support_imgs, support_labels, query_imgs): """ Args: support_imgs: [N_support, C, H, W] all support images (mixed classes) support_labels: [N_support] integer labels query_imgs: [N_query, C, H, W] all query images Returns: logits: [N_query, num_classes] (negative distances) """ num_classes = support_labels.max().item() + 1 device = support_imgs.device # 1. Extract features for all support images with torch.no_grad(): # we can also allow gradients through feature extractor support_feats = self.feature_extractor(support_imgs) # [N_s, D] # 2. Per class: build graph, run GCN, compute prototype prototypes = [] for c in range(num_classes): mask = (support_labels == c) if mask.sum() == 0: # fallback: zero prototype prototypes.append(torch.zeros(self.feat_dim, device=device)) continue c_feats = support_feats[mask] # [N_c, D] adj = self.build_graph(c_feats) # [N_c, N_c] refined_feats = c_feats for layer in self.gcn: refined_feats = layer(refined_feats, adj) # [N_c, D] refined_feats = F.relu(refined_feats) # optional non‑linearity # prototype = mean of refined support features proto = refined_feats.mean(dim=0) # [D] prototypes.append(proto) prototypes = torch.stack(prototypes, dim=0) # [C, D] # 3. Extract query features and compute distances query_feats = self.feature_extractor(query_imgs) # [N_q, D] # squared Euclidean distance: [N_q, C] dists = torch.cdist(query_feats, prototypes, p=2) ** 2 # logits = negative distances (higher = more similar) logits = -dists return logits # ------------------------------- # 3. Training loop (episodic style) # ------------------------------- def train_episodic(model, train_loader, num_classes, support_size, query_size, epochs, lr=1e-3, device='cpu'): model.to(device) optimizer = optim.Adam(model.parameters(), lr=lr) # get full training set as tensors (for easy random sampling) all_imgs = [] all_labels = [] for x, y in train_loader: all_imgs.append(x) all_labels.append(y) all_imgs = torch.cat(all_imgs, dim=0) all_labels = torch.cat(all_labels, dim=0) model.train() for epoch in range(epochs): # Sample support and query indices for each class support_idxs = [] query_idxs = [] for c in range(num_classes): c_mask = (all_labels == c).nonzero(as_tuple=False).squeeze() c_indices = c_mask[torch.randperm(len(c_mask))] # Take first support_size for support, rest for query (if enough) n_support = min(support_size, len(c_mask)) n_query = min(query_size, len(c_mask) - n_support) support_idxs.append(c_indices[:n_support]) query_idxs.append(c_indices[n_support:n_support + n_query]) support_idxs = torch.cat(support_idxs) query_idxs = torch.cat(query_idxs) support_imgs = all_imgs[support_idxs].to(device) support_labels = all_labels[support_idxs].to(device) query_imgs = all_imgs[query_idxs].to(device) query_labels = all_labels[query_idxs].to(device) optimizer.zero_grad() logits = model(support_imgs, support_labels, query_imgs) loss = F.cross_entropy(logits, query_labels) loss.backward() optimizer.step() if (epoch+1) % 10 == 0: pred = logits.argmax(dim=1) acc = (pred == query_labels).float().mean().item() print(f"Epoch {epoch+1:3d} | Loss {loss.item():.4f} | Query Acc {acc:.3f}") # ------------------------------- # 4. Example on MNIST # ------------------------------- if __name__ == "__main__": transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # MNIST mean/std ]) train_set = datasets.MNIST('../data', train=True, download=True, transform=transform) train_loader = DataLoader(train_set, batch_size=256, shuffle=True) # used only to collect tensors # Simple CNN feature extractor class MNISTFeatureExtractor(nn.Module): def __init__(self, out_dim=64): super().__init__() self.conv = nn.Sequential( nn.Conv2d(1, 32, 3, 1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, 1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*5*5, out_dim) # MNIST: after two poolings -> 5x5 ) def forward(self, x): return self.conv(x) feat_extractor = MNISTFeatureExtractor(out_dim=64) classifier = ClassGraphClassifier( feature_extractor=feat_extractor, feat_dim=64, gcn_hidden=64, gcn_layers=2, k_neighbors=5 ) train_episodic( model=classifier, train_loader=train_loader, num_classes=10, support_size=20, # 20 support images per class query_size=30, # 30 query images per class epochs=200, lr=1e-3, device='cuda' if torch.cuda.is_available() else 'cpu' )