Kaggle小技巧:TTA(test time augmentation)測試時增強

Pluto0054發表於2020-10-24

1.原理簡述

TTA:test time augmentation
顧名思義是測試時增強的意思,可將準確率提高若干個百分點,它就是測試時增強(test time augmentation, TTA)。這裡會為原始影像造出多個不同版本,包括不同區域裁剪和更改縮放程度等,並將它們輸入到模型中;然後對多個版本進行計算得到平均輸出,作為影像的最終輸出分數。
儘管有作弊的嫌疑。
但這種技術很有效,因為原始影像顯示的區域可能會缺少一些重要特徵,在模型中輸入影像的多個版本並取平均值,能解決上述問題。

2.程式碼實現

這次打kaggle比賽時,是影像目標檢測類的,用到了TTA,提高了LB,實現的主要API類如下:

***
類中函式說明:
augment:為一張圖片製作TTA
batch_augment:為多張圖片批量製作TTA
deaugment_boxes:將TTA預測框返回到影像的原始狀態中
***

class BaseTTA:
    image_size = 512

    def augment(self, image):
        raise NotImplementedError
    
    def batch_augment(self, images):
        raise NotImplementedError
    
    def deaugment_boxes(self, boxes):
        raise NotImplementedError

# 水平翻轉
class TTAHorizontalFlip(BaseTTA):
    def augment(self, image):
        return image.flip(1)
    
    def batch_augment(self, images):
        return images.flip(2)
    
    def deaugment_boxes(self, boxes):
        boxes[:, [1,3]] = self.image_size - boxes[:, [3,1]]
        return boxes

# 垂直翻轉
class TTAVerticalFlip(BaseTTA):
    
    def augment(self, image):
        return image.flip(2)
    
    def batch_augment(self, images):
        return images.flip(3)
    
    def deaugment_boxes(self, boxes):
        boxes[:, [0,2]] = self.image_size - boxes[:, [2,0]]
        return boxes

# 旋轉90°    
class TTARotate90(BaseTTA):
    def augment(self, image):
        return torch.rot90(image, 1, (1, 2))

    def batch_augment(self, images):
        return torch.rot90(images, 1, (2, 3))
    
    def deaugment_boxes(self, boxes):
        res_boxes = boxes.copy()
        res_boxes[:, [0,2]] = self.image_size - boxes[:, [1,3]]
        res_boxes[:, [1,3]] = boxes[:, [2,0]]
        return res_boxes

# 組合使用
class TTACompose(BaseTTA):
    def __init__(self, transforms):
        self.transforms = transforms
        
    def augment(self, image):
        for transform in self.transforms:
            image = transform.augment(image)
        return image
    
    def batch_augment(self, images):
        for transform in self.transforms:
            images = transform.batch_augment(images)
        return images
    
    def prepare_boxes(self, boxes):
        result_boxes = boxes.copy()
        result_boxes[:,0] = np.min(boxes[:, [0,2]], axis=1)
        result_boxes[:,2] = np.max(boxes[:, [0,2]], axis=1)
        result_boxes[:,1] = np.min(boxes[:, [1,3]], axis=1)
        result_boxes[:,3] = np.max(boxes[:, [1,3]], axis=1)
        return result_boxes
    
    def deaugment_boxes(self, boxes):
        for transform in self.transforms[::-1]:
            boxes = transform.deaugment_boxes(boxes)
        return self.prepare_boxes(boxes)

須知少時凌雲志,曾許人間第一流!

相關文章