How to do TTA in classification problem ?

[xiaoerlaigeid]

No description provided.

[talhaanwarch]

import ttach as tta
transforms = tta.Compose(
    [
        tta.HorizontalFlip(),
        tta.VerticalFlip(),
        tta.Rotate90(angles=[0, 90,180]),
    ]
)

tta_model = tta.ClassificationTTAWrapper(model, transforms)


model.cuda().eval()
labels,preds=[],[]
with torch.no_grad():
    for batch in loader:
        image,label=batch
        pred=tta_model(image.cuda())
        pred=torch.argmax(pred,dim=1).detach().cpu().numpy()
        labels.append(label.cpu().numpy())
        preds.append(pred)

[puneeshkhanna]

One quick doubt when we do TTA, are we also predicting on original image and the predictions probs are mean of original image + the augmented images based upon the transforms ?