Test-time augmentation that remembers to un-flip

What it does

TTAch wraps PyTorch models to perform Test Time Augmentation on image inputs. Instead of predicting once on a clean test image, it runs inference on multiple augmented variants—flips, rotations, scaling, crops—and aggregates the results. The library handles both the forward augmentation and the reverse transformation of outputs like segmentation masks or keypoint coordinates.

The interesting bit

The real friction in TTA is usually undoing geometric transforms on model outputs so they align spatially before averaging. ttach automates this de-augmentation for standard wrappers, and exposes a manual loop for custom multi-input or multi-output architectures where you need to reverse each tensor yourself.

Key highlights

• Ready-made wrappers for segmentation, classification, and keypoint models
• Pre-baked aliases such as d4_transform (flips plus four rotations) and ten_crop_transform
• Six merge modes including mean, geometric mean, and temperature sharpening
• Compose chains support arbitrary combinations of flips, scales, resizes, and intensity shifts
• Explicit de-augmentation methods for masks, labels, and keypoints

Caveats

• Keypoint models must return coordinates in a strict torch([x1, y1, ..., xn, yn]) format
• The README is silent on inference cost and memory usage when running dozens of augmentations per batch

Verdict

Handy if you are already doing TTA in PyTorch and tired of writing inverse flip logic. Skip it if you need object detection or video inference—the README only covers classification, segmentation, and keypoints.