Final Push: Beating State-of-the-Art on CIFAR-10

import torch import torch.nn as nn import torch.nn.functional as F import math class ShakeDropFunction(torch.autograd.Function): """ ShakeDrop: Custom autograd function for different forward/backward behavior. During training: - Forward: output = x + (b + alpha - b*alpha) * F(x) - Backward: gradient scaled by (b + beta - b*beta) Where: - b: Bernoulli random variable (survival probability) - alpha, beta: Uniform random in [-1, 1] """ @staticmethod def forward(ctx, x, residual, survival_prob, alpha_range, training): if training: # Sample Bernoulli (1 = keep, 0 = drop) gate = torch.bernoulli( torch.ones(x.size(0), 1, 1, 1, device=x.device) * survival_prob ) # Sample alpha uniformly from [-1, 1] alpha = torch.empty(x.size(0), 1, 1, 1, device=x.device).uniform_( -alpha_range, alpha_range ) # Forward pass scaling # When gate=1: scaling = 1 + alpha - alpha = 1 (no change) # When gate=0: scaling = 0 + alpha - 0 = alpha (shake) scaling = gate + alpha - gate * alpha # Save for backward ctx.save_for_backward(residual, gate) ctx.survival_prob = survival_prob ctx.alpha_range = alpha_range return x + scaling * residual else: # Inference: use expected value # E[gate + alpha - gate*alpha] = p + 0 - p*0 = p return x + survival_prob * residual @staticmethod def backward(ctx, grad_output): residual, gate = ctx.saved_tensors # Sample beta uniformly for backward pass (different from alpha) beta = torch.empty( grad_output.size(0), 1, 1, 1, device=grad_output.device ).uniform_(-ctx.alpha_range, ctx.alpha_range) # Backward scaling scaling = gate + beta - gate * beta # Gradient for x is unchanged grad_x = grad_output # Gradient for residual is scaled grad_residual = scaling * grad_output return grad_x, grad_residual, None, None, None class ShakeDrop(nn.Module): """ ShakeDrop module wrapper. Args: survival_prob: Probability of keeping the residual (not dropping) alpha_range: Range for uniform sampling (default 1.0 = [-1, 1]) """ def __init__(self, survival_prob=1.0, alpha_range=1.0): super().__init__() self.survival_prob = survival_prob self.alpha_range = alpha_range def forward(self, x, residual): return ShakeDropFunction.apply( x, residual, self.survival_prob, self.alpha_range, self.training )

import torch import torch.nn as nn import torch.nn.functional as F class PyramidBottleneckShakeDrop(nn.Module): """ PyramidNet bottleneck block with ShakeDrop. This is the block that achieved 98.7% on CIFAR-10 (without ensemble). """ expansion = 4 def __init__(self, in_channels, out_channels, stride=1, survival_prob=1.0): super().__init__() bottleneck_channels = out_channels // 4 self.bn1 = nn.BatchNorm2d(in_channels) self.conv1 = nn.Conv2d(in_channels, bottleneck_channels, 1, bias=False) self.bn2 = nn.BatchNorm2d(bottleneck_channels) self.conv2 = nn.Conv2d( bottleneck_channels, bottleneck_channels, 3, stride=stride, padding=1, bias=False ) self.bn3 = nn.BatchNorm2d(bottleneck_channels) self.conv3 = nn.Conv2d(bottleneck_channels, out_channels, 1, bias=False) self.bn4 = nn.BatchNorm2d(out_channels) self.shakedrop = ShakeDrop(survival_prob=survival_prob) self.stride = stride self.in_channels = in_channels self.out_channels = out_channels def forward(self, x): # Residual path out = self.conv1(self.bn1(x)) out = self.conv2(F.relu(self.bn2(out))) out = self.conv3(F.relu(self.bn3(out))) out = self.bn4(out) # Shortcut with zero-padding shortcut = x if self.stride != 1: shortcut = F.avg_pool2d(shortcut, 2) if self.in_channels != self.out_channels: pad_channels = self.out_channels - self.in_channels shortcut = F.pad(shortcut, (0, 0, 0, 0, 0, pad_channels)) # Apply ShakeDrop out = self.shakedrop(shortcut, out) return out class PyramidNetShakeDrop(nn.Module): """ PyramidNet with ShakeDrop regularization. Paper: "ShakeDrop Regularization for Deep Residual Learning" This is the architecture that achieved state-of-the-art results on CIFAR. Args: depth: Network depth (272 for SOTA) alpha: Widening factor (200 for SOTA) num_classes: Number of output classes survival_prob_last: Survival probability for deepest block (0.5 typical) """ def __init__( self, depth=272, alpha=200, num_classes=10, survival_prob_last=0.5 ): super().__init__() # Calculate blocks per group n = (depth - 2) // 9 # Bottleneck has 3 convs per block self.in_channels = 16 self.alpha = alpha self.total_blocks = n * 3 self.block_idx = 0 self.add_rate = alpha / self.total_blocks self.survival_prob_last = survival_prob_last # Initial conv self.conv1 = nn.Conv2d(3, 16, 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(16) # Three groups self.group1 = self._make_group(n, stride=1) self.group2 = self._make_group(n, stride=2) self.group3 = self._make_group(n, stride=2) # Final layers self.final_channels = int(round(16 + alpha)) * 4 self.bn_final = nn.BatchNorm2d(self.final_channels) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(self.final_channels, num_classes) self._initialize_weights() def _get_survival_prob(self): """Linear decay of survival probability with depth.""" self.block_idx += 1 prob = 1.0 - (self.block_idx / self.total_blocks) * (1.0 - self.survival_prob_last) return prob def _make_group(self, num_blocks, stride): layers = [] for i in range(num_blocks): self.block_idx_temp = self.block_idx out_channels = int(round(16 + self.add_rate * (self.block_idx + 1))) * 4 survival_prob = self._get_survival_prob() s = stride if i == 0 else 1 layers.append( PyramidBottleneckShakeDrop( self.in_channels, out_channels, s, survival_prob ) ) self.in_channels = out_channels return nn.Sequential(*layers) def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out') elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) def forward(self, x): out = self.conv1(x) out = self.bn1(out) out = self.group1(out) out = self.group2(out) out = self.group3(out) out = F.relu(self.bn_final(out)) out = self.avgpool(out) out = out.view(out.size(0), -1) out = self.fc(out) return out def pyramidnet272_shakedrop(num_classes=10): """PyramidNet-272 with ShakeDrop - SOTA configuration.""" return PyramidNetShakeDrop( depth=272, alpha=200, num_classes=num_classes, survival_prob_last=0.5 )

import torch import torch.nn as nn import torch.optim as optim import math class LongTrainingScheduler: """ Extended cosine annealing for SOTA training. Paper settings: - 1800 epochs total - Initial LR: 0.5 - Batch size: 128 - Weight decay: 1e-4 - Cosine annealing to 0 """ def __init__( self, optimizer, total_epochs=1800, warmup_epochs=10, min_lr=0.0 ): self.optimizer = optimizer self.total_epochs = total_epochs self.warmup_epochs = warmup_epochs self.min_lr = min_lr self.base_lr = optimizer.param_groups[0]['lr'] self.current_epoch = 0 def step(self): if self.current_epoch < self.warmup_epochs: # Linear warmup lr = self.base_lr * (self.current_epoch + 1) / self.warmup_epochs else: # Cosine annealing progress = (self.current_epoch - self.warmup_epochs) / ( self.total_epochs - self.warmup_epochs ) lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * ( 1 + math.cos(math.pi * progress) ) for param_group in self.optimizer.param_groups: param_group['lr'] = lr self.current_epoch += 1 return lr def train_sota(model, train_loader, test_loader, device, epochs=1800): """ SOTA training configuration. Based on PyramidNet+ShakeDrop paper settings. """ # Optimizer optimizer = optim.SGD( model.parameters(), lr=0.5, momentum=0.9, weight_decay=1e-4, nesterov=True ) # Scheduler scheduler = LongTrainingScheduler( optimizer, total_epochs=epochs, warmup_epochs=10 ) # Loss with label smoothing criterion = LabelSmoothingCrossEntropy(epsilon=0.1) # EMA ema = EMA(model, decay=0.9999) # Mixup/CutMix mixup_fn = MixupCutMix(mixup_alpha=0.2, cutmix_alpha=1.0) best_acc = 0 for epoch in range(epochs): model.train() for inputs, targets in train_loader: inputs, targets = inputs.to(device), targets.to(device) # Mixup/CutMix inputs, targets_a, targets_b, lam = mixup_fn(inputs, targets) # Forward outputs = model(inputs) loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam) # Backward optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() # EMA update ema.update() # LR schedule lr = scheduler.step() # Evaluate every 10 epochs if (epoch + 1) % 10 == 0: ema.apply_shadow() acc = evaluate(model, test_loader, device) ema.restore() print(f'Epoch {epoch+1}/{epochs} | LR: {lr:.6f} | Acc: {acc:.2f}%') if acc > best_acc: best_acc = acc ema.apply_shadow() torch.save(model.state_dict(), 'sota_model.pth') ema.restore() return best_acc

import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as T from torch.utils.data import DataLoader from tqdm import tqdm # Import all our components from models.pyramidnet_shakedrop import pyramidnet272_shakedrop from augmentations.cutout import Cutout from augmentations.mixup import MixupCutMix, mixup_criterion from training.label_smoothing import LabelSmoothingCrossEntropy from training.weight_averaging import EMA from training.long_training import LongTrainingScheduler def get_sota_transforms(): """SOTA augmentation pipeline.""" MEAN = (0.4914, 0.4822, 0.4465) STD = (0.2470, 0.2435, 0.2616) train_transform = T.Compose([ T.RandomCrop(32, padding=4, padding_mode='reflect'), T.RandomHorizontalFlip(), T.AutoAugment(policy=T.AutoAugmentPolicy.CIFAR10), T.ToTensor(), T.Normalize(MEAN, STD), Cutout(n_holes=1, length=16), ]) test_transform = T.Compose([ T.ToTensor(), T.Normalize(MEAN, STD), ]) return train_transform, test_transform def main(): device = torch.device('cuda') print(f'Using: {torch.cuda.get_device_name(0)}') # Configuration config = { 'epochs': 1800, 'batch_size': 128, 'lr': 0.5, 'weight_decay': 1e-4, 'label_smoothing': 0.1, 'ema_decay': 0.9999, 'mixup_alpha': 0.2, 'cutmix_alpha': 1.0, } # Data train_transform, test_transform = get_sota_transforms() train_set = torchvision.datasets.CIFAR10( './data', train=True, download=True, transform=train_transform ) test_set = torchvision.datasets.CIFAR10( './data', train=False, download=True, transform=test_transform ) train_loader = DataLoader( train_set, batch_size=config['batch_size'], shuffle=True, num_workers=8, pin_memory=True ) test_loader = DataLoader( test_set, batch_size=100, shuffle=False, num_workers=8, pin_memory=True ) # Model model = pyramidnet272_shakedrop(num_classes=10).to(device) num_params = sum(p.numel() for p in model.parameters()) print(f'Parameters: {num_params:,}') # EMA ema = EMA(model, decay=config['ema_decay']) # Loss criterion = LabelSmoothingCrossEntropy(epsilon=config['label_smoothing']) # Optimizer optimizer = optim.SGD( model.parameters(), lr=config['lr'], momentum=0.9, weight_decay=config['weight_decay'], nesterov=True ) # Scheduler scheduler = LongTrainingScheduler(optimizer, total_epochs=config['epochs']) # Mixup/CutMix mixup_fn = MixupCutMix( mixup_alpha=config['mixup_alpha'], cutmix_alpha=config['cutmix_alpha'] ) # Training best_acc = 0 for epoch in range(config['epochs']): model.train() pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}') for inputs, targets in pbar: inputs, targets = inputs.to(device), targets.to(device) # Mixup/CutMix inputs, targets_a, targets_b, lam = mixup_fn(inputs, targets) # Forward outputs = model(inputs) loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam) # Backward optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() # EMA ema.update() pbar.set_postfix({'loss': f'{loss.item():.4f}'}) # LR lr = scheduler.step() # Evaluate if (epoch + 1) % 10 == 0: ema.apply_shadow() model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, targets in test_loader: inputs, targets = inputs.to(device), targets.to(device) outputs = model(inputs) _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() acc = 100 * correct / total ema.restore() print(f'Epoch {epoch+1} | LR: {lr:.6f} | Acc: {acc:.2f}%') if acc > best_acc: best_acc = acc ema.apply_shadow() torch.save(model.state_dict(), 'pyramid272_sota.pth') ema.restore() print(f'New best: {best_acc:.2f}%') print(f'\nFinal Best: {best_acc:.2f}%') if __name__ == '__main__': main()

def create_sota_ensemble(device): """ Create the ultimate SOTA ensemble. Combines: 1. PyramidNet-272 + ShakeDrop (x3 seeds) 2. WideResNet-40-10 + ShakeDrop (x2 seeds) 3. Test-Time Augmentation Target: 99%+ accuracy """ models_config = {} # PyramidNet-272 models (3 different seeds) for seed in [1, 2, 3]: model = pyramidnet272_shakedrop(num_classes=10) model.load_state_dict(torch.load(f'checkpoints/pyramid272_seed{seed}.pth')) model.to(device).eval() models_config[f'pyramid_{seed}'] = (model, 1.2) # WideResNet-40-10 models (2 different seeds) for seed in [1, 2]: model = wrn_40_10(num_classes=10) model.load_state_dict(torch.load(f'checkpoints/wrn40_seed{seed}.pth')) model.to(device).eval() models_config[f'wrn_{seed}'] = (model, 1.0) # Create ensemble base_ensemble = DiverseEnsemble(models_config) # Add TTA tta_transforms = [ lambda x: x, lambda x: torch.flip(x, [-1]), ] final_ensemble = FullEnsembleWithTTA(base_ensemble, tta_transforms) return final_ensemble