""" Test image preprocessing strategies that actually reduce token count. Corrected experiments: 1. Downsample (no restore): resize to smaller, feed directly (inference only) 2. Checkerboard removal: remove every other row (halve image size) 3. Center crop: crop a smaller region from the image Usage: python test_preprocess_reduce.py --dataset cifar100 --gpu 5 """ import torch import torch.nn as nn import torch.nn.functional as F import time import argparse import sys import os sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from datasets import get_cifar100_loader, get_oxford_pets_loader, get_food101_loader import timm DATASETS = { 'cifar100': (get_cifar100_loader, 100), 'oxford_pets': (get_oxford_pets_loader, 37), 'food101': (get_food101_loader, 101), } def resize_down(images, target_size): """Resize image to smaller resolution (no upsample back).""" B, C, H, W = images.shape return F.interpolate(images, size=(target_size, target_size), mode='bilinear', align_corners=False) def remove_even_rows_cols(images): """Remove every other row AND column → image size halves. 224×224 → remove even rows → 112×224 → remove even cols → 112×112 Effectively 75% pixel reduction, image shrinks to 1/4 area. """ return images[:, :, ::2, ::2] def remove_even_rows_only(images): """Remove every other row only → height halves, width unchanged. 224×224 → 112×224 (non-square, but fewer patches) """ return images[:, :, ::2, :] def center_crop(images, crop_size): """Center crop the image to a smaller size.""" B, C, H, W = images.shape h_start = (H - crop_size) // 2 w_start = (W - crop_size) // 2 return images[:, :, h_start:h_start+crop_size, w_start:w_start+crop_size] def get_num_patches(img_size, patch_size=16, stride=16): """Calculate number of patches for a given image size.""" grid = (img_size - patch_size) // stride + 1 return grid * grid def make_model_for_size(device, num_classes, img_size, checkpoint_path, patch_size=16): """Create a ViT-B/16 model adapted for a specific input size. Loads checkpoint from 224×224 model, interpolates pos_embed for new size. """ model = timm.create_model( 'vit_base_patch16_224.augreg_in21k', pretrained=True, num_classes=num_classes, img_size=img_size) if os.path.exists(checkpoint_path): ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=True) sd = ckpt.get('model_state_dict', ckpt) # Interpolate position embeddings if sizes differ if 'pos_embed' in sd and sd['pos_embed'].shape != model.pos_embed.shape: old_pos = sd.pop('pos_embed') cls_pos = old_pos[:, 0:1, :] patch_pos = old_pos[:, 1:, :] old_grid = int((patch_pos.shape[1]) ** 0.5) new_grid = (img_size - patch_size) // patch_size + 1 if old_grid != new_grid: patch_2d = patch_pos.reshape(1, old_grid, old_grid, -1).permute(0, 3, 1, 2) patch_new = F.interpolate(patch_2d, size=(new_grid, new_grid), mode='bilinear', align_corners=False) patch_new = patch_new.permute(0, 2, 3, 1).reshape(1, -1, 768) sd['pos_embed'] = torch.cat([cls_pos, patch_new], dim=1) missing, unexpected = model.load_state_dict(sd, strict=False) if missing: print(f' Missing keys: {len(missing)} (expected for diff img_size)', flush=True) if unexpected: print(f' Unexpected keys: {len(unexpected)}', flush=True) model = model.to(device) model.eval() return model @torch.no_grad() def evaluate(model, loader, device, preprocess=None, track_time=False, img_size=224): model.eval() correct = 0 total = 0 total_time = 0 for images, targets in loader: images, targets = images.to(device), targets.to(device) if preprocess is not None: images = preprocess(images) if track_time: torch.cuda.synchronize() start = time.time() logits = model(images) if track_time: torch.cuda.synchronize() end = time.time() total_time += (end - start) _, predicted = logits.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() acc = 100. * correct / total if track_time: latency = total_time / total * 1000 throughput = total / total_time return acc, latency, throughput return acc def main(): parser = argparse.ArgumentParser() parser.add_argument('--dataset', type=str, default='cifar100', choices=list(DATASETS.keys())) parser.add_argument('--gpu', type=int, default=0) parser.add_argument('--batch_size', type=int, default=32) args = parser.parse_args() device = f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu' print(f'Device: {device}', flush=True) if torch.cuda.is_available(): print(f'GPU: {torch.cuda.get_device_name(args.gpu)}', flush=True) loader_fn, num_classes = DATASETS[args.dataset] result = loader_fn(batch_size=args.batch_size, data_dir='./data', num_workers=4) if len(result) == 4: train_loader, val_loader, test_loader, n_cls = result else: train_loader, test_loader, n_cls = result val_loader = test_loader print(f'Dataset: {args.dataset}, Test: {len(test_loader.dataset)}, Classes: {n_cls}', flush=True) # Load fine-tuned ViT-B/16 print('Loading ViT-B/16 IN-21K...', flush=True) model = timm.create_model('vit_base_patch16_224.augreg_in21k', pretrained=True, num_classes=n_cls) model = model.to(device) model.eval() ckpt_path = f'checkpoints/{args.dataset}_vit_b16_ft/best_model.pth' if os.path.exists(ckpt_path): ckpt = torch.load(ckpt_path, map_location=device, weights_only=True) if 'model_state_dict' in ckpt: model.load_state_dict(ckpt['model_state_dict']) else: model.load_state_dict(ckpt) print(f' Loaded fine-tuned checkpoint: {ckpt_path}', flush=True) # Warmup for images, _ in test_loader: images = images.to(device) _ = model(images) break ckpt_base = f'checkpoints/{args.dataset}_vit_b16_ft/best_model.pth' methods = [ # (name, preprocess_fn, img_size, patches) ('Baseline (224×224)', None, 224, 196), ('Downsample 168×168 (推理)', lambda x: resize_down(x, 168), 168, 100), ('Downsample 112×112 (推理)', lambda x: resize_down(x, 112), 112, 49), ('隔行去行列 (224→112)', remove_even_rows_cols, 112, 49), ('Center crop 168×168', lambda x: center_crop(x, 168), 168, 100), ('Center crop 112×112', lambda x: center_crop(x, 112), 112, 49), ] print(f'\n=== {args.dataset}: 图片预处理（实际减少 token 数量） ===\n', flush=True) print(f'{"方法":<30} {"Acc":>8} {"Patches":>8} {"Latency":>10} {"Throughput":>12}', flush=True) print('-' * 75, flush=True) for name, preprocess, img_size, expected_patches in methods: if isinstance(img_size, tuple): # Skip non-square cases (need separate model creation) print(f'{name:<30} skipped (non-square input)', flush=True) continue # Create a fresh model for each input size m = make_model_for_size(device, n_cls, 224 if isinstance(img_size, tuple) else img_size, ckpt_base) acc, latency, throughput = evaluate(m, test_loader, device, preprocess, track_time=True, img_size=img_size) print(f'{name:<30} {acc:>7.2f}% {expected_patches:>8} {latency:>8.2f}ms {throughput:>10.1f}/s', flush=True) del m torch.cuda.empty_cache() print(f'\n对比：降采样训练（已训练过的模型，非推理时硬切）') print(f' 168×168 训练: 91.56%, 100 patches (来自 test_downsample_train.py)') print(f' 112×112 训练: 90.00%, 49 patches (来自 test_downsample_train.py)') print(f'\n结论：推理时直接缩小图片会掉点较多，因为位置编码不匹配') print(f'降采样训练（重新训练适配）的效果远好于推理时硬切', flush=True) if __name__ == '__main__': main()