import torch import torch.nn as nn from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR import time import os import timm from timm.data import Mixup from timm.utils import ModelEmaV2 from models import create_model def compute_flops(model, input_size=(1, 3, 224, 224), device='cpu'): """Estimate FLOPs using fvcore.""" try: from fvcore.nn import FlopCountAnalysis dummy = torch.randn(input_size).to(device) model = model.to(device) flops = FlopCountAnalysis(model, dummy) return flops.total() / 1e9 # GFLOPs except Exception: return 0.0 class Trainer: def __init__(self, model, train_loader, test_loader, num_classes, device='cuda', lr=1e-3, weight_decay=0.05, label_smoothing=0.0, mixup_fn=None, ema_decay=0.0, clip_grad=None, warmup_epochs=0, total_epochs=50): self.model = model.to(device) self.train_loader = train_loader self.test_loader = test_loader self.device = device self.num_classes = num_classes self.mixup_fn = mixup_fn self.ema_decay = ema_decay self.clip_grad = clip_grad self.warmup_epochs = warmup_epochs self.criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing) self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) # Warmup + Cosine scheduler (in iterations) self.total_iters = len(train_loader) cosine_iters = max(1, total_epochs - warmup_epochs) * self.total_iters if warmup_epochs > 0: warmup_iters = warmup_epochs * self.total_iters warmup_scheduler = LinearLR(self.optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_iters) cosine_scheduler = CosineAnnealingLR(self.optimizer, T_max=cosine_iters) self.scheduler = SequentialLR(self.optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_iters]) else: self.scheduler = CosineAnnealingLR(self.optimizer, T_max=cosine_iters) # EMA self.ema_model = ModelEmaV2(model, decay=ema_decay) if ema_decay > 0 else None def train_epoch(self, epoch): self.model.train() total_loss = 0 correct = 0 total = 0 for batch_idx, (images, targets) in enumerate(self.train_loader): images, targets = images.to(self.device), targets.to(self.device) # MixUp/CutMix if self.mixup_fn is not None: images, targets = self.mixup_fn(images, targets) self.optimizer.zero_grad() outputs = self.model(images) loss = self.criterion(outputs, targets) loss.backward() # Gradient clipping if self.clip_grad is not None and self.clip_grad > 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad) self.optimizer.step() # EMA update if self.ema_model is not None: self.ema_model.update(self.model) total_loss += loss.item() if self.mixup_fn is None: _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() else: # Approximate accuracy for mixup (count correct if both augmentations agree) _, predicted = outputs.max(1) total += targets.size(0) if targets.dim() > 1: correct += predicted.eq(targets.argmax(dim=1)).sum().item() else: correct += predicted.eq(targets).sum().item() if batch_idx % 50 == 0: print(f' Batch {batch_idx}/{len(self.train_loader)}, Loss: {loss.item():.4f}') avg_loss = total_loss / len(self.train_loader) accuracy = 100. * correct / total return avg_loss, accuracy def evaluate(self, model=None): model = model if model is not None else self.model model.eval() total_loss = 0 correct = 0 total = 0 with torch.no_grad(): for images, targets in self.test_loader: images, targets = images.to(self.device), targets.to(self.device) outputs = model(images) loss = self.criterion(outputs, targets) total_loss += loss.item() _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() avg_loss = total_loss / len(self.test_loader) accuracy = 100. * correct / total return avg_loss, accuracy def train(self, epochs, save_dir='./checkpoints', resume_from=None): os.makedirs(save_dir, exist_ok=True) best_acc = 0 start_epoch = 0 # Resume from checkpoint if resume_from: ckpt = torch.load(resume_from, map_location=self.device) self.model.load_state_dict(ckpt['model_state_dict']) self.optimizer.load_state_dict(ckpt['optimizer_state_dict']) self.scheduler.load_state_dict(ckpt['scheduler_state_dict']) best_acc = ckpt.get('test_acc', 0) start_epoch = ckpt.get('epoch', 0) + 1 print(f'Resumed from epoch {ckpt.get("epoch", 0)} (best acc: {best_acc:.2f}%)') for epoch in range(start_epoch, epochs): start_time = time.time() train_loss, train_acc = self.train_epoch(epoch) # Use raw model for evaluation tracking; EMA is for final inference only test_loss, test_acc = self.evaluate(self.model) self.scheduler.step() epoch_time = time.time() - start_time print(f'Epoch {epoch+1}/{epochs} ({epoch_time:.1f}s)') print(f' Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%') print(f' Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%') print(f' LR: {self.optimizer.param_groups[0]["lr"]:.6f}') print() if test_acc > best_acc: best_acc = test_acc ckpt = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scheduler_state_dict': self.scheduler.state_dict(), 'test_acc': test_acc, } if self.ema_model is not None: ckpt['ema_state_dict'] = self.ema_model.module.state_dict() torch.save(ckpt, os.path.join(save_dir, 'best_model.pth')) print(f' -> Saved best model (Acc: {best_acc:.2f}%)') print() # Periodic checkpoint every 10 epochs if (epoch + 1) % 10 == 0: ckpt = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scheduler_state_dict': self.scheduler.state_dict(), 'test_acc': test_acc, } if self.ema_model is not None: ckpt['ema_state_dict'] = self.ema_model.module.state_dict() torch.save(ckpt, os.path.join(save_dir, f'checkpoint_epoch{epoch+1}.pth')) print(f' -> Saved periodic checkpoint (epoch {epoch+1})') print() print(f'Training complete! Best accuracy: {best_acc:.2f}%') return best_acc def compute_efficiency_metrics(model, test_loader, device='cuda'): model.eval() model = model.to(device) total_time = 0 total_samples = 0 with torch.no_grad(): for images, _ in test_loader: images = images.to(device) batch_size = images.size(0) if device != 'cpu': torch.cuda.synchronize() start_time = time.time() _ = model(images) if device != 'cpu': torch.cuda.synchronize() end_time = time.time() total_time += (end_time - start_time) total_samples += batch_size avg_latency = total_time / total_samples * 1000 throughput = total_samples / total_time return avg_latency, throughput def get_model_summary(model, input_size=(1, 3, 224, 224), device='cpu'): """Print model parameter count and FLOPs.""" model = model.to(device) params = sum(p.numel() for p in model.parameters()) / 1e6 gflops = compute_flops(model, input_size, device) return params, gflops def get_dataset_loader(dataset, batch_size, data_dir='./data', num_workers=4, use_randaugment=False): """Dispatch to the right dataset loader.""" from datasets import get_cifar10_loader, get_cifar100_loader, get_tiny_imagenet_loader # Tiny-ImageNet at 224x224, others default if dataset == 'cifar10': train_loader, test_loader, num_classes = get_cifar10_loader( batch_size=batch_size, data_dir=data_dir, num_workers=num_workers) elif dataset == 'cifar100': train_loader, test_loader, num_classes = get_cifar100_loader( batch_size=batch_size, data_dir=data_dir, num_workers=num_workers, use_randaugment=use_randaugment) elif dataset == 'tiny_imagenet': train_loader, test_loader, num_classes = get_tiny_imagenet_loader( batch_size=batch_size, data_dir=data_dir, image_size=224, num_workers=num_workers) else: raise ValueError(f'Unknown dataset: {dataset}') return train_loader, test_loader, num_classes if __name__ == '__main__': import argparse import sys from datetime import datetime parser = argparse.ArgumentParser(description='Vision Transformer Training with Patch Selection') parser.add_argument('--model', type=str, default='vit_small', choices=['vit_small', 'swin_tiny', 'patch_selection_vit', 'random_prune_vit'], help='Model architecture') parser.add_argument('--dataset', type=str, default='cifar100', choices=['cifar10', 'cifar100', 'tiny_imagenet'], help='Dataset') parser.add_argument('--batch_size', type=int, default=128, help='Batch size') parser.add_argument('--epochs', type=int, default=50, help='Total number of epochs') parser.add_argument('--resume', type=str, default=None, help='Resume from checkpoint path') parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate') parser.add_argument('--weight_decay', type=float, default=0.05, help='Weight decay') parser.add_argument('--keep_ratio', type=float, default=0.5, help='Ratio of patches to keep (for pruning methods)') parser.add_argument('--selection_mode', type=str, default='topk', choices=['topk', 'adaptive'], help='Patch selection strategy: fixed top-k or adaptive threshold') parser.add_argument('--adaptive_alpha', type=float, default=0.5, help='Threshold multiplier for adaptive mode: thresh = mean + alpha*std') parser.add_argument('--patch_size', type=int, default=16, help='Patch size for image splitting (8, 16, 32)') parser.add_argument('--patch_stride', type=int, default=None, help='Patch stride (None = same as patch_size, smaller = overlap)') parser.add_argument('--num_workers', type=int, default=10, help='Data loading workers') parser.add_argument('--use_randaugment', action='store_true', default=False, help='Use RandAugment for data augmentation') parser.add_argument('--label_smoothing', type=float, default=0.0, help='Label smoothing factor (0.0 = disabled, 0.1 recommended)') parser.add_argument('--device', type=str, default='auto', help='Device (auto, cpu, cuda, cuda:7, etc.)') parser.add_argument('--save_dir', type=str, default='./checkpoints', help='Checkpoint save directory') parser.add_argument('--pretrained', action='store_true', default=True, help='Use pretrained weights') parser.add_argument('--load_mae', type=str, default=None, help='Path to MAE pretrained encoder weights for stage 2 finetuning') parser.add_argument('--log_dir', type=str, default='./logs', help='Directory for log files') # Anti-overfitting parser.add_argument('--drop_path', type=float, default=0.0, help='Stochastic depth drop path rate (0.0 = disabled, 0.1 recommended)') parser.add_argument('--mixup', type=float, default=0.0, help='MixUp alpha (0.0 = disabled, 0.8 recommended)') parser.add_argument('--cutmix', type=float, default=0.0, help='CutMix alpha (0.0 = disabled, 1.0 recommended)') parser.add_argument('--ema_decay', type=float, default=0.0, help='EMA decay rate (0.0 = disabled, 0.9999 recommended)') parser.add_argument('--clip_grad', type=float, default=None, help='Gradient clipping max norm (None = disabled, 1.0 recommended)') parser.add_argument('--warmup_epochs', type=int, default=0, help='Number of linear warmup epochs (5-10 recommended)') args = parser.parse_args() # Setup log file os.makedirs(args.log_dir, exist_ok=True) timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') exp_tag = f'{args.dataset}_{args.model}' if 'prune' in args.model or 'selection' in args.model: exp_tag += f'_keep{int(args.keep_ratio*100)}_{args.selection_mode}' if args.model == 'patch_selection_vit' and args.selection_mode == 'adaptive': exp_tag += f'_alpha{args.adaptive_alpha}' if args.patch_size != 16: exp_tag += f'_ps{args.patch_size}' log_file = f'{args.log_dir}/{exp_tag}_{timestamp}.log' log_fh = open(log_file, 'w', buffering=1) # line-buffered def teeprint(*args_tee, **kwargs_tee): """Print to both console and log file.""" kwargs_tee.pop('file', None) original_print(*args_tee, **kwargs_tee) original_print(*args_tee, file=log_fh, **kwargs_tee) log_fh.flush() original_print = print print = teeprint # Device setup if args.device == 'auto': device = 'cuda' if torch.cuda.is_available() else 'cpu' else: device = args.device print(f'Using device: {device}') # Dataset train_loader, test_loader, num_classes = get_dataset_loader( args.dataset, args.batch_size, num_workers=args.num_workers) print(f'Dataset: {args.dataset}, {len(train_loader.dataset)} train, ' f'{len(test_loader.dataset)} test, {num_classes} classes') # Model model = create_model( model_name=args.model, num_classes=num_classes, keep_ratio=args.keep_ratio, pretrained=args.pretrained, selection_mode=args.selection_mode, adaptive_alpha=args.adaptive_alpha, patch_size=args.patch_size, patch_stride=args.patch_stride, drop_path_rate=args.drop_path, ) # Load MAE pretrained weights (for Stage 2) if args.load_mae and hasattr(model, 'load_mae_pretrained'): model.load_mae_pretrained(args.load_mae) # Compute model summary on CPU first (fvcore doesn't always work on CUDA) summary_device = 'cpu' if not device.startswith('cuda') else device params_m, gflops = get_model_summary(model, device=summary_device) print(f'Model: {args.model}') print(f' Parameters: {params_m:.2f}M') print(f' FLOPs: {gflops:.2f}G') if 'prune' in args.model or 'selection' in args.model: print(f' Keep ratio: {args.keep_ratio} ({int(args.keep_ratio * 100)}%)') print(f' Selection mode: {args.selection_mode}') if args.selection_mode == 'adaptive': print(f' Adaptive alpha: {args.adaptive_alpha}') if args.patch_size != 16 or (args.patch_stride and args.patch_stride != args.patch_size): print(f' Patch config: size={args.patch_size}, stride={args.patch_stride or args.patch_size}') print(f' Patches per image: {model.patch_embed.num_patches}') # Anti-overfitting config print(f' RandAugment: {args.use_randaugment}') print(f' Label smoothing: {args.label_smoothing}') print(f' Drop path: {args.drop_path}') print(f' MixUp alpha: {args.mixup}, CutMix alpha: {args.cutmix}') print(f' EMA decay: {args.ema_decay}') print(f' Gradient clipping: {args.clip_grad}') print(f' Warmup epochs: {args.warmup_epochs}') # Experiment tag for checkpoint dir tag = f'{args.dataset}_{args.model}' if 'prune' in args.model or 'selection' in args.model: tag += f'_keep{int(args.keep_ratio*100)}_{args.selection_mode}' if args.model == 'patch_selection_vit' and args.selection_mode == 'adaptive': tag += f'_alpha{args.adaptive_alpha}' if args.patch_size != 16: tag += f'_ps{args.patch_size}' save_dir = f'{args.save_dir}/{tag}' # MixUp / CutMix mixup_fn = None if args.mixup > 0 or args.cutmix > 0: mixup_fn = Mixup( mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, label_smoothing=args.label_smoothing, num_classes=num_classes, ) print(f'Using MixUp/CutMix: alpha={args.mixup}/{args.cutmix}') # Training trainer = Trainer( model, train_loader, test_loader, num_classes, device=device, lr=args.lr, weight_decay=args.weight_decay, label_smoothing=args.label_smoothing, mixup_fn=mixup_fn, ema_decay=args.ema_decay, clip_grad=args.clip_grad, warmup_epochs=args.warmup_epochs, total_epochs=args.epochs, ) best_acc = trainer.train(epochs=args.epochs, save_dir=save_dir, resume_from=args.resume) # Efficiency evaluation print('\n=== Efficiency Metrics ===') latency, throughput = compute_efficiency_metrics(model, test_loader, device=device) print(f' Latency: {latency:.2f} ms per sample') print(f' Throughput: {throughput:.2f} samples/sec') print(f' FLOPs: {gflops:.2f}G') print(f' Best Accuracy: {best_acc:.2f}%') print(f'========================\n')