This is the official PyTorch Implementation of "SoTTA: Robust Test-Time Adaptation on Noisy Data Streams (NeurIPS '23)" by Taesik Gong*, Yewon Kim*, Taeckyung Lee*, Sorn Chottananurak, and Sung-Ju Lee (* Equal contribution).
args = None

CIFAR10Opt = {
    'name': 'cifar10',
    'batch_size': 128,

    'learning_rate': 0.1,  # initial learning rate
    'weight_decay': 0.0005,
    'momentum': 0.9,
    'img_size': 3072,

    'file_path': './dataset/CIFAR-10-C',
    'classes': ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'],
    'num_class': 10,
    'severity': 5,
    'domains': ["original",

                "test",

                "gaussian_noise-1", "gaussian_noise-2", "gaussian_noise-3", "gaussian_noise-4", "gaussian_noise-5",
                "gaussian_noise-all",

                "shot_noise-1", "shot_noise-2", "shot_noise-3", "shot_noise-4", "shot_noise-5", "shot_noise-all",

                "impulse_noise-1", "impulse_noise-2", "impulse_noise-3", "impulse_noise-4", "impulse_noise-5",
                "impulse_noise-all",

                "defocus_blur-1", "defocus_blur-2", "defocus_blur-3", "defocus_blur-4", "defocus_blur-5",
                "defocus_blur-all",

                "glass_blur-1", "glass_blur-2", "glass_blur-3", "glass_blur-4", "glass_blur-5", "glass_blur-all",

                "motion_blur-1", "motion_blur-2", "motion_blur-3", "motion_blur-4", "motion_blur-5", "motion_blur-all",

                "zoom_blur-1", "zoom_blur-2", "zoom_blur-3", "zoom_blur-4", "zoom_blur-5", "zoom_blur-all",

                "snow-1", "snow-2", "snow-3", "snow-4", "snow-5", "snow-all",

                "frost-1", "frost-2", "frost-3", "frost-4", "frost-5", "frost-all",

                "fog-1", "fog-2", "fog-3", "fog-4", "fog-5", "fog-all",

                "brightness-1", "brightness-2", "brightness-3", "brightness-4", "brightness-5", "brightness-all",

                "contrast-1", "contrast-2", "contrast-3", "contrast-4", "contrast-5", "contrast-all",

                "elastic_transform-1", "elastic_transform-2", "elastic_transform-3", "elastic_transform-4",
                "elastic_transform-5", "elastic_transform-all",

                "pixelate-1", "pixelate-2", "pixelate-3", "pixelate-4", "pixelate-5", "pixelate-all",

                "jpeg_compression-1", "jpeg_compression-2", "jpeg_compression-3", "jpeg_compression-4",
                "jpeg_compression-5", "jpeg_compression-all",
                ],
    'src_domains': ["original"],
    'tgt_domains': [
        "gaussian_noise-5",
        "shot_noise-5",
        "impulse_noise-5",
        "defocus_blur-5",
        "glass_blur-5",
        "motion_blur-5",
        "zoom_blur-5",
        "snow-5",
        "frost-5",
        "fog-5",
        "brightness-5",
        "contrast-5",
        "elastic_transform-5",
        "pixelate-5",
        "jpeg_compression-5",

    ],
}

CIFAR100Opt = {
    'name': 'cifar100',
    'batch_size': 128,

    'learning_rate': 0.1,  # initial learning rate
    'weight_decay': 0.0005,
    'momentum': 0.9,
    'img_size': 3072,

    'file_path': './dataset/CIFAR-100-C',
    'classes': ['beaver', 'dolphin', 'otter', 'seal', 'whale',
                'aquarium fish', 'flatfish', 'ray', 'shark', 'trout',
                'orchids', 'poppies', 'roses', 'sunflowers', 'tulips',
                'bottles', 'bowls', 'cans', 'cups', 'plates',
                'apples', 'mushrooms', 'oranges', 'pears', 'sweet peppers',
                'clock', 'computer keyboard', 'lamp', 'telephone', 'television',
                'bed', 'chair', 'couch', 'table', 'wardrobe',
                'bee', 'beetle', 'butterfly', 'caterpillar', 'cockroach',
                'bear', 'leopard', 'lion', 'tiger', 'wolf',
                'bridge', 'castle', 'house', 'road', 'skyscraper',
                'cloud', 'forest', 'mountain', 'plain', 'sea',
                'camel', 'cattle', 'chimpanzee', 'elephant', 'kangaroo',
                'fox', 'porcupine', 'possum', 'raccoon', 'skunk',
                'crab', 'lobster', 'snail', 'spider', 'worm',
                'baby', 'boy', 'girl', 'man', 'woman',
                'crocodile', 'dinosaur', 'lizard', 'snake', 'turtle',
                'hamster', 'mouse', 'rabbit', 'shrew', 'squirrel',
                'maple', 'oak', 'palm', 'pine', 'willow',
                'bicycle', 'bus', 'motorcycle', 'pickup truck', 'train',
                'lawn-mower', 'rocket', 'streetcar', 'tank', 'tractor'],
    'num_class': 100,
    'severity': 5,
    # 'corruptions': ["shot_noise", "motion_blur", "snow", "pixelate", "gaussian_noise", "defocus_blur",
    #                 "brightness", "fog", "zoom_blur", "frost", "glass_blur", "impulse_noise", "contrast",
    #                 "jpeg_compression", "elastic_transform"],
    'domains': ["original",

                "test",

                "gaussian_noise-1", "gaussian_noise-2", "gaussian_noise-3", "gaussian_noise-4", "gaussian_noise-5",
                "gaussian_noise-all",

                "shot_noise-1", "shot_noise-2", "shot_noise-3", "shot_noise-4", "shot_noise-5", "shot_noise-all",

                "impulse_noise-1", "impulse_noise-2", "impulse_noise-3", "impulse_noise-4", "impulse_noise-5",
                "impulse_noise-all",

                "defocus_blur-1", "defocus_blur-2", "defocus_blur-3", "defocus_blur-4", "defocus_blur-5",
                "defocus_blur-all",

                "glass_blur-1", "glass_blur-2", "glass_blur-3", "glass_blur-4", "glass_blur-5", "glass_blur-all",

                "motion_blur-1", "motion_blur-2", "motion_blur-3", "motion_blur-4", "motion_blur-5", "motion_blur-all",

                "zoom_blur-1", "zoom_blur-2", "zoom_blur-3", "zoom_blur-4", "zoom_blur-5", "zoom_blur-all",

                "snow-1", "snow-2", "snow-3", "snow-4", "snow-5", "snow-all",

                "frost-1", "frost-2", "frost-3", "frost-4", "frost-5", "frost-all",

                "fog-1", "fog-2", "fog-3", "fog-4", "fog-5", "fog-all",

                "brightness-1", "brightness-2", "brightness-3", "brightness-4", "brightness-5", "brightness-all",

                "contrast-1", "contrast-2", "contrast-3", "contrast-4", "contrast-5", "contrast-all",

                "elastic_transform-1", "elastic_transform-2", "elastic_transform-3", "elastic_transform-4",
                "elastic_transform-5", "elastic_transform-all",

                "pixelate-1", "pixelate-2", "pixelate-3", "pixelate-4", "pixelate-5", "pixelate-all",

                "jpeg_compression-1", "jpeg_compression-2", "jpeg_compression-3", "jpeg_compression-4",
                "jpeg_compression-5", "jpeg_compression-all",
                ],
    'src_domains': ["original"],
    'tgt_domains': [
        "gaussian_noise-5",
        "shot_noise-5",
        "impulse_noise-5",
        "defocus_blur-5",
        "glass_blur-5",
        "motion_blur-5",
        "zoom_blur-5",
        "snow-5",
        "frost-5",
        "fog-5",
        "brightness-5",
        "contrast-5",
        "elastic_transform-5",
        "pixelate-5",
        "jpeg_compression-5",

    ],
}

IMAGENET_C = {
    # referred to for hyperparams: https://github.com/Lornatang/ResNet-PyTorch/blob/9e529757ce0607aafeae2ddd97142201b3d4cadd/examples/imagenet/main.py
    'name': 'imagenet',
    'batch_size': 256,
    'learning_rate': 0.001,
    'weight_decay': 0,
    'momentum': 0.9,
    'img_size': 3072,

    'file_path': './dataset/ImageNet-C',
    'num_class': 1000,
    'severity': 5,
    'domains': ["gaussian_noise-1", "gaussian_noise-2", "gaussian_noise-3", "gaussian_noise-4", "gaussian_noise-5",
                "gaussian_noise-all",

                "shot_noise-1", "shot_noise-2", "shot_noise-3", "shot_noise-4", "shot_noise-5", "shot_noise-all",

                "impulse_noise-1", "impulse_noise-2", "impulse_noise-3", "impulse_noise-4", "impulse_noise-5",
                "impulse_noise-all",

                "defocus_blur-1", "defocus_blur-2", "defocus_blur-3", "defocus_blur-4", "defocus_blur-5",
                "defocus_blur-all",

                "glass_blur-1", "glass_blur-2", "glass_blur-3", "glass_blur-4", "glass_blur-5", "glass_blur-all",

                "motion_blur-1", "motion_blur-2", "motion_blur-3", "motion_blur-4", "motion_blur-5", "motion_blur-all",

                "zoom_blur-1", "zoom_blur-2", "zoom_blur-3", "zoom_blur-4", "zoom_blur-5", "zoom_blur-all",

                "snow-1", "snow-2", "snow-3", "snow-4", "snow-5", "snow-all",

                "frost-1", "frost-2", "frost-3", "frost-4", "frost-5", "frost-all",

                "fog-1", "fog-2", "fog-3", "fog-4", "fog-5", "fog-all",

                "brightness-1", "brightness-2", "brightness-3", "brightness-4", "brightness-5", "brightness-all",

                "contrast-1", "contrast-2", "contrast-3", "contrast-4", "contrast-5", "contrast-all",

                "elastic_transform-1", "elastic_transform-2", "elastic_transform-3", "elastic_transform-4",
                "elastic_transform-5", "elastic_transform-all",

                "pixelate-1", "pixelate-2", "pixelate-3", "pixelate-4", "pixelate-5", "pixelate-all",

                "jpeg_compression-1", "jpeg_compression-2", "jpeg_compression-3", "jpeg_compression-4",
                "jpeg_compression-5", "jpeg_compression-all",
                ],

    'src_domains': ["original"],
    'tgt_domains': ["gaussian_noise-5",
                    "shot_noise-5",
                    "impulse_noise-5",
                    "defocus_blur-5",
                    "glass_blur-5",
                    "motion_blur-5",
                    "zoom_blur-5",
                    "snow-5",
                    "frost-5",
                    "fog-5",
                    "brightness-5",
                    "contrast-5",
                    "elastic_transform-5",
                    "pixelate-5",
                    "jpeg_compression-5",

                    ],
}

MNISTOpt = {
    'name': 'mnist',
    'batch_size': 128,

    'learning_rate': 0.1,  # initial learning rate
    'weight_decay': 0.0005,
    'momentum': 0.9,
    'img_size': 3072,

    'file_path': './dataset/MNIST-C',
    'classes': ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'],
    'num_class': 10,
    'domains': ["original",

                "test",

                "shot_noise",
                "impulse_noise",
                "glass_blur",
                "motion_blur",
                "shear",
                "scale",
                "rotate",
                "brightness",
                "translate",
                "stripe",
                "fog",
                "spatter",
                "dotted_line",
                "zigzag",
                "canny_edges",

                ],
    'src_domains': ["original"],
    'tgt_domains': [
        "shot_noise",
        "impulse_noise",
        "glass_blur",
        "motion_blur",
        "shear",
        "scale",
        "rotate",
        "brightness",
        "translate",
        "stripe",
        "fog",
        "spatter",
        "dotted_line",
        "zigzag",
        "canny_edges",
    ],

}

conf.py
taeckyung/SoTTA

import torch.utils.data
import torchvision.transforms as transforms

import time
import numpy as np
import conf

opt = conf.CIFAR100Opt


class CIFAR100Dataset(torch.utils.data.Dataset):

    def __init__(self, domains=None, activities=None, transform='none'):
        self.domains = domains
        self.activity = activities

        self.domains = domains

        self.img_shape = opt['img_size']
        self.features = None
        self.class_labels = None
        self.domain_labels = None
        self.file_path = opt['file_path']

        assert (len(domains) > 0)
        if domains[0].startswith('original'):
            self.sub_path1 = 'origin'
            self.sub_path2 = ''
            self.data_filename = 'original.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].startswith('test'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-1'  # all data are same in 1~5
            self.data_filename = 'test.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-1'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-1'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-2'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-2'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-3'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-3'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-4'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-4'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-5'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-5'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-all'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-all'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'

        if transform == 'src':
            self.transform = transforms.Compose(
                [
                    transforms.RandomCrop(32, padding=4),
                    transforms.RandomHorizontalFlip(),
                ])

        elif transform == 'val':
            self.transform = None
        else:
            raise NotImplementedError

        self.preprocessing()

    def preprocessing(self):

        path = f'{self.file_path}/{self.sub_path1}/{self.sub_path2}/'

        data = np.load(path + self.data_filename)
        # change NHWC to NCHW format
        data = np.transpose(data, (0, 3, 1, 2))
        # make it compatible with our models (normalize)
        data = data.astype(np.float32) / 255.0

        self.features = data
        self.class_labels = np.load(path + self.label_filename)
        # assume that single domain is passed as List
        self.domain_labels = np.array([0 for _ in range(len(self.features))])

        self.dataset = torch.utils.data.TensorDataset(
            torch.from_numpy(self.features),  # resize for resnet
            torch.from_numpy(self.class_labels),
            torch.from_numpy(self.domain_labels))

    def __len__(self):
        return len(self.dataset)

    def get_num_domains(self):
        return len(self.domains)

    def __getitem__(self, idx):
        if isinstance(idx, torch.Tensor):
            idx = idx.item()
        img, cl, dl = self.dataset[idx]
        if self.transform:
            img = self.transform(img)
        return img, cl, dl


if __name__ == '__main__':
    pass

data_loader/CIFAR100Dataset.py
taeckyung/SoTTA
import torch.utils.data
import torchvision.transforms as transforms

import time
import numpy as np
import conf

opt = conf.CIFAR10Opt


class CIFAR10Dataset(torch.utils.data.Dataset):

    def __init__(self, domains=None, activities=None, transform='none'):
        self.domains = domains
        self.activity = activities

        self.domains = domains

        self.img_shape = opt['img_size']
        self.features = None
        self.class_labels = None
        self.domain_labels = None
        self.file_path = opt['file_path']

        assert (len(domains) > 0)
        if domains[0].startswith('original'):
            self.sub_path1 = 'origin'
            self.sub_path2 = ''
            self.data_filename = 'original.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].startswith('test'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-1'  # all data are same in 1~5
            self.data_filename = 'test.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-1'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-1'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-2'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-2'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-3'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-3'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-4'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-4'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('-5'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity-5'
            self.data_filename = domains[0].split('-')[0] + '.npy'
            self.label_filename = 'labels.npy'
        elif domains[0].endswith('_all'):
            self.sub_path1 = 'corrupted'
            self.sub_path2 = 'severity_all'
            self.data_filename = str(domains[0][:-4]) + '.npy'
            self.label_filename = 'labels.npy'

        if transform == 'src':
            self.transform = transforms.Compose(
                [
                    transforms.RandomCrop(32, padding=4),
                    transforms.RandomHorizontalFlip(),
                ])

        elif transform == 'val':
            self.transform = None
        else:
            raise NotImplementedError

        self.preprocessing()

    def preprocessing(self):

        path = f'{self.file_path}/{self.sub_path1}/{self.sub_path2}/'

        data = np.load(path + self.data_filename)
        # change NHWC to NCHW format
        data = np.transpose(data, (0, 3, 1, 2))
        # make it compatible with our models (normalize)
        data = data.astype(np.float32) / 255.0
        self.features = data
        self.class_labels = np.load(path + self.label_filename)
        # assume that single domain is passed as List
        self.domain_labels = np.array([0 for _ in range(len(self.features))])

        self.dataset = torch.utils.data.TensorDataset(
            torch.from_numpy(self.features),
            torch.from_numpy(self.class_labels),
            torch.from_numpy(self.domain_labels))

    def __len__(self):
        return len(self.dataset)

    def get_num_domains(self):
        return len(self.domains)

    def __getitem__(self, idx):
        if isinstance(idx, torch.Tensor):
            idx = idx.item()
        img, cl, dl = self.dataset[idx]
        if self.transform:
            img = self.transform(img)
        return img, cl, dl


if __name__ == '__main__':
    pass

data_loader/CIFAR10Dataset.py
taeckyung/SoTTA
import torch.utils.data
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder

import time
import numpy as np
import conf

opt = conf.IMAGENET_C


class ImageNetDataset(torch.utils.data.Dataset):

    def __init__(self, domain=None, activities=None, transform='none'):
        self.domain = domain
        self.activity = activities

        self.domain = domain
        self.features = None
        self.class_labels = None
        self.domain_labels = None
        self.file_path = opt['file_path']
        self.transform_type = transform

        assert (len(domain) > 0)
        if domain.startswith('original'):
            self.path = 'origin/Data/CLS-LOC/train/'
        elif domain.startswith('test'):
            self.path = 'origin/Data/CLS-LOC/val/'
        else:
            self.path = 'corrupted/'
            corruption, severity = domain.split('-')
            self.path += corruption + '/' + severity + '/'

        if transform == 'src':
            self.transform = transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor()
            ])
        elif transform == 'val':
            if domain.startswith('original') or domain.startswith('test'):
                self.transform = transforms.Compose([
                    transforms.Resize(256),
                    transforms.CenterCrop(224),
                    transforms.ToTensor()
                ])
            else:
                self.transform = transforms.Compose([
                    transforms.CenterCrop(224),
                    transforms.ToTensor()
                ])
        else:
            raise NotImplementedError

        self.preprocessing()

    def preprocessing(self):

        path = self.file_path + '/' + self.path
        self.features = []
        self.class_labels = []
        self.domain_labels = []
        print('preprocessing images..')
        self.dataset = ImageFolder(path)

    def load_features(self):
        path = self.file_path + '/' + self.path
        dataset = ImageFolder(path, transform=self.transform)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False, pin_memory=False,
                                                 drop_last=False)
        # transformed_dataset = []
        for b_i, data in enumerate(dataloader):  # must be loaded from dataloader, due to transform in the __getitem__()
            feat, cl = data
            # convert a batch of tensors to list, and then append to our list one by one
            feats = torch.unbind(feat, dim=0)
            cls = torch.unbind(cl, dim=0)
            for i in range(len(feats)):
                # transformed_dataset.append((feats[i], cls[i]))
                self.features.append(feats[i])
                self.class_labels.append(cls[i])
                self.domain_labels.append(0)
        self.features = np.stack(self.features)
        self.class_labels = np.stack(self.class_labels)
        self.domain_labels = np.stack(self.domain_labels)

    def __len__(self):
        return len(self.dataset)

    def get_num_domains(self):
        return 1

    def __getitem__(self, idx):
        if isinstance(idx, torch.Tensor):
            idx = idx.item()
        img, cl = self.dataset[idx]
        if self.transform:
            img = self.transform(img)
        return img, cl, torch.tensor(0)


if __name__ == '__main__':
    ### code for making imagenet validation data compatiable with ImageFolder!
    '''
    import os
    root = '/mnt/sting/tsgong/WWW/dataset/ImageNet-C/origin/Data/CLS-LOC/val/'
    f = open(root+'LOC_val_solution.csv', 'r')
    i=0
    for l in f:
        if i ==0: # ignore header
            i += 1
            continue
        filename=l.split(',')[0]
        label=l.split(',')[1].split(' ')[0]
        dir = root+label
        ### 1. make dir
        # if not os.path.exists(dir):
        #     os.makedirs(dir)
        # print(os.path.join(root,filename,'.JPEG'))
        ### 2. move files to dir
        print(label)
        if os.path.isfile(os.path.join(root, filename + '.JPEG')):
            os.rename(os.path.join(root, filename + '.JPEG'), os.path.join(dir, filename + '.JPEG'))
        i += 1
    print(i)
    '''

data_loader/IMAGENETDataset.py
taeckyung/SoTTA
import random
import torch.utils.data
import torchvision.transforms as transforms

import time
import numpy as np
from torchvision.datasets import ImageFolder

import conf

from data_loader.CIFAR100Dataset import CIFAR100Dataset
from data_loader.CIFAR10Dataset import CIFAR10Dataset
from data_loader.IMAGENETDataset import ImageNetDataset
from utils.normalize_layer import *

opt10 = conf.CIFAR10Opt
opt100 = conf.CIFAR100Opt
imagenet = conf.IMAGENET_C
mnist = conf.MNISTOpt

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(conf.args.gpu_idx)  #

NOISY_CLASS_IDX = 10000


class NoisyDataset(torch.utils.data.Dataset):
    """
    Noisy data stream
    Currently supports cifar10noisy, imagenetnoisy, and cifar100noisy
    with noise data of original, cifar100, imagenet, mnist, uniform, and repeat.
    """
    noisy_data_types = ["original", "divide", "repeat", "oneclassrepeat", "cifar100", "cifar100c",
                        "gaussian", "uniform", "mnist", "cifar10", "imagenet"]

    def __init__(self, base="cifar10noisy",
                 domains=None, activities=None, transform='none',
                 noisy_data=None, noisy_data_size=None, noisy_data_class=None):
        self.domains = domains
        self.activity = activities
        self.noisy_data = noisy_data
        self.noisy_data_size = noisy_data_size
        self.noisy_data_class = noisy_data_class

        self.domain = domains[0]

        self.img_shape = opt10['img_size']
        self.features = None
        self.class_labels = None
        self.domain_labels = None

        self.file_path10 = opt10['file_path']
        self.file_path100 = opt100['file_path']
        self.file_path_imagenet = imagenet['file_path']
        self.file_path_mnist = mnist['file_path']

        assert (base in ["cifar10noisy", "imagenetnoisy", "cifar100noisy"])
        assert (self.noisy_data in self.noisy_data_types)
        assert not (self.noisy_data in ['oneclassrepeat', 'gaussian', 'uniform'] and self.noisy_data_size is None)
        assert not (self.noisy_data in ['divide', 'oneclassrepeat'] and self.noisy_data_class is None)

        if transform == 'src':
            self.is_src = True
        else:
            self.is_src = False

        if base == 'cifar10noisy':
            if self.domain == 'none':
                self.base_dataset = CIFAR10Dataset(domains=['original'], transform=transform)
            else:
                self.base_dataset = CIFAR10Dataset(domains=self.domains, transform=transform)
            self.img_size = 32
            if transform == 'src':
                self.transform = transforms.Compose(
                    [
                        transforms.RandomCrop(32, padding=4),
                        transforms.RandomHorizontalFlip(),
                    ])

            elif transform == 'val':
                self.transform = None
            else:
                raise NotImplementedError
        elif base == 'cifar100noisy':
            if self.domain == 'none':
                self.base_dataset = CIFAR100Dataset(domains=['original'], transform=transform)
            else:
                self.base_dataset = CIFAR100Dataset(domains=self.domains, transform=transform)
            self.img_size = 32
            if transform == 'src':
                self.transform = transforms.Compose(
                    [
                        transforms.RandomCrop(32, padding=4),
                        transforms.RandomHorizontalFlip(),
                    ])

            elif transform == 'val':
                self.transform = None
            else:
                raise NotImplementedError
        elif base == 'imagenetnoisy':
            self.base_dataset = ImageNetDataset(domain=self.domain, transform=transform)
            self.img_size = 224

            if transform == 'src':
                self.transform = transforms.Compose([
                    transforms.RandomResizedCrop(224),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor()
                ])
            elif transform == 'val':
                if self.domain.startswith('original') or self.domain.startswith('test'):
                    self.transform = transforms.Compose([
                        transforms.Resize(256),
                        transforms.CenterCrop(224),
                        transforms.ToTensor()
                    ])
                else:
                    self.transform = transforms.Compose([
                        transforms.CenterCrop(224),
                        transforms.ToTensor()
                    ])

                if self.noisy_data in ["oneclassrepeat", "divide", "repeat"]:
                    self.base_dataset.load_features()  # imagenet does not originally load features
            else:
                raise NotImplementedError

        else:
            raise NotImplementedError

        self.preprocessing()

    def resample(self, data: np.ndarray, size: int):
        """
        Resample the input data to corresponding size.
        """
        if size == len(data):
            return data
        elif size < len(data):
            return data[random.sample(range(len(data)), size)]
        else:
            remainder = data[random.sample(range(len(data)), size % len(data))]
            repeat = np.tile(data, (int(size / len(data)), 1, 1, 1))
            return np.concatenate((repeat, remainder), axis=0)

    def preprocessing(self):
        if self.noisy_data == "cifar100":  # cifar100
            path100 = f'{self.file_path100}/origin/'
            data100 = np.load(path100 + 'original.npy')
            # change NHWC to NCHW format
            data100 = np.transpose(data100, (0, 3, 1, 2))
            # make it compatible with our models (normalize)
            data100 = data100.astype(np.float32) / 255.0

            if self.noisy_data_size:
                data100 = self.resample(data100, self.noisy_data_size)

            tr = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
            ])
            data100 = np.array([tr(torch.Tensor(img)).numpy() for img in data100])

            noisy_stream = data100

        elif self.noisy_data == "cifar100c":  # cifar100c
            path100c = f'{self.file_path100}/corrupted/severity-{self.domains[0][-1]}/'
            data_filename = self.domains[0].split('-')[0] + '.npy'
            data100c = np.load(path100c + data_filename)
            # change NHWC to NCHW format
            data100c = np.transpose(data100c, (0, 3, 1, 2))
            # make it compatible with our models (normalize)
            data100c = data100c.astype(np.float32) / 255.0

            tr = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
            ])
            data100c = np.array([tr(torch.Tensor(img)).numpy() for img in data100c])

            noisy_stream = data100c

        elif self.noisy_data == "cifar10":  # cifar10
            path10 = f'{self.file_path10}/origin/'
            data10 = np.load(path10 + 'original.npy')
            # change NHWC to NCHW format
            data10 = np.transpose(data10, (0, 3, 1, 2))
            # make it compatible with our models (normalize)
            data10 = data10.astype(np.float32) / 255.0

            if self.noisy_data_size:
                data10 = data10[random.sample(range(len(data10)), self.noisy_data_size)]

            tr = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
            ])
            data10 = np.array([tr(torch.Tensor(img)).numpy() for img in data10])

            noisy_stream = data10

        elif self.noisy_data == "imagenet":  # imagenet validation set
            path = f'{self.file_path_imagenet}/origin/Data/CLS-LOC/val/'

            tr = transforms.Compose([
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
            ])

            features = []

            dataset = ImageFolder(path, transform=tr)
            dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False, pin_memory=False,
                                                     drop_last=False)
            
            for b_i, data in enumerate(dataloader):  # must be loaded from dataloader, due to transform in the __getitem__()
                feat, _ = data
                # convert a batch of tensors to list, and then append to our list one by one
                feats = torch.unbind(feat, dim=0)
                for i in range(len(feats)):
                    features.append(feats[i])

            noisy_stream = np.stack(features)
            if self.noisy_data_size:
                noisy_stream = noisy_stream[random.sample(range(len(noisy_stream)), self.noisy_data_size)]

        elif self.noisy_data == "gaussian":  # gaussian noise
            gaussian_img = np.random.normal(0, 1, size=(self.noisy_data_size, 3, self.img_size, self.img_size))
            norm_gaussian_img = np.float32((gaussian_img - np.min(gaussian_img)) / np.ptp(gaussian_img))
            noisy_stream = norm_gaussian_img

        elif self.noisy_data == "repeat":  # repeat base dataset
            noisy_stream = self.base_dataset.features

        elif self.noisy_data == "uniform":  # uniform50
            noisy_stream = np.float32(np.random.random(size=(self.noisy_data_size, 3, self.img_size, self.img_size)))

        elif self.noisy_data == "mnist":  # mnist test set
            path_mnist = f'{self.file_path_mnist}/identity/'
            data_mnist = np.load(path_mnist + 'test_images.npy')

            # change NHWC to NCHW format
            data_mnist = np.transpose(data_mnist, (0, 3, 1, 2))
            # make it compatible with our models (normalize)
            data_mnist = data_mnist.astype(np.float32) / 255.0

            if self.noisy_data_size:
                data_mnist = self.resample(data_mnist, self.noisy_data_size)

            mnist_transform = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Grayscale(3),
                transforms.Resize((self.img_size, self.img_size)),
                transforms.ToTensor(),
            ])
            data_mnist = np.array([mnist_transform(torch.Tensor(img)).numpy() for img in data_mnist])

            noisy_stream = data_mnist

        elif self.noisy_data == "original":
            noisy_stream = np.array([])
        else:
            raise NotImplementedError

        if self.noisy_data == "original" or self.is_src is True:
            self.dataset = self.base_dataset.dataset
        else:
            noisy_dataset = torch.utils.data.TensorDataset(
                torch.from_numpy(noisy_stream),  
                torch.from_numpy(NOISY_CLASS_IDX * np.ones((len(noisy_stream)))),
                torch.from_numpy(np.array(np.zeros((len(noisy_stream))))))

            if self.domain == "none":
                self.dataset = noisy_dataset
            else:
                self.dataset = torch.utils.data.ConcatDataset([noisy_dataset, self.base_dataset.dataset])

    def __len__(self):
        return len(self.dataset)

    def get_num_domains(self):
        return len(self.domains)

    def __getitem__(self, idx):
        if isinstance(idx, torch.Tensor):
            idx = idx.item()

        data = self.dataset[idx]
        if len(data) == 3:
            img, cl, dl = data
        elif len(data) == 2:
            img, cl, dl = data[0], torch.tensor(data[1]), torch.tensor(0)  # avoid errors on torch.stack()
        else:
            raise NotImplementedError

        if int(cl) != int(NOISY_CLASS_IDX):  # in-dist data
            if self.transform:
                img = self.transform(img)

        return img, cl, dl


if __name__ == '__main__':
    pass

data_loader/NoisyDataset.py
taeckyung/SoTTA
from torch.utils.data import DataLoader
import torch
import numpy as np
import time
import math

from .IMAGENETDataset import ImageNetDataset
from .CIFAR10Dataset import CIFAR10Dataset
from .CIFAR100Dataset import CIFAR100Dataset
from .NoisyDataset import NoisyDataset
import os
import re
import random
import conf


def keep_order_split(entire_data, train_size, valid_size, test_size):
    # Split the dataset while keeping the original order of it.
    all_indices = [i for i in range(len(entire_data))]

    valid_indices = random.sample(all_indices, valid_size)

    for i in sorted(valid_indices, reverse=True):  # reverse is required
        all_indices.pop(i)

    test_indices = random.sample(all_indices, test_size)

    for i in sorted(test_indices, reverse=True):
        all_indices.pop(i)

    valid_data = torch.utils.data.Subset(entire_data, valid_indices)
    test_data = torch.utils.data.Subset(entire_data, test_indices)
    train_data = torch.utils.data.Subset(entire_data, all_indices)

    return train_data, valid_data, test_data


def split_data(entire_data, valid_split, test_split, train_max_rows, valid_max_rows, test_max_rows):
    valid_size = math.floor(len(entire_data) * valid_split)
    test_size = math.floor(len(entire_data) * test_split)

    train_size = len(entire_data) - valid_size - test_size

    assert (train_size >= 0 and valid_size >= 0 and test_size >= 0)

    train_data, valid_data, test_data = keep_order_split(entire_data, train_size, valid_size, test_size)

    if len(entire_data) > train_max_rows:
        train_data = torch.utils.data.Subset(train_data, range(train_max_rows))
    if len(valid_data) > valid_max_rows:
        valid_data = torch.utils.data.Subset(valid_data, range(valid_max_rows))
    if len(test_data) > test_max_rows:
        test_data = torch.utils.data.Subset(test_data, range(test_max_rows))

    return train_data, valid_data, test_data


def datasets_to_dataloader(datasets, batch_size, concat=True, shuffle=True, drop_last=False):
    if concat:
        data_loader = None
        if len(datasets):
            if type(datasets) is torch.utils.data.dataset.Subset:
                datasets = [datasets]
            if sum([len(dataset) for dataset in datasets]) > 0:  # at least one dataset has data
                data_loader = DataLoader(torch.utils.data.ConcatDataset(datasets), batch_size=batch_size,
                                         shuffle=shuffle, drop_last=drop_last, pin_memory=False)
        return data_loader
    else:
        data_loaders = []
        for dataset in datasets:
            if len(dataset) == 0:
                continue
            else:
                data_loaders.append(DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                                               drop_last=drop_last, pin_memory=False))

        return data_loaders


def load_cache(dataset, cond, data_file_path, transform=None):
    try:
        root = './cached_data/'
        dir = root + str(dataset) + '/'
        filename = re.sub("[^a-zA-Z0-9 \n]", '_', str(cond) + '_' + str(data_file_path))
        if transform:
            filename += '_' + transform + '.pkl'
        else:
            filename += '.pkl'
        cache_path = dir + filename

        if os.path.isfile(cache_path):
            print(f'Cache hit:{cache_path}')
            return torch.load(cache_path)
        else:
            print(f'Cache miss:{cache_path}')
            return None
    except:  # RuntimeError, EOFError
        return None


def save_cache(loaded_data, dataset, cond, data_file_path, transform=None):
    root = './cached_data/'
    dir = root + str(dataset) + '/'
    filename = re.sub("[^a-zA-Z0-9 \n]", '_', str(cond) + '_' + str(data_file_path))
    if transform:
        filename += '_' + transform + '.pkl'
    else:
        filename += '.pkl'
    cache_path = dir + filename
    try:
        if not os.path.exists(dir):
            os.makedirs(dir)
    except:
        pass
    return torch.save(loaded_data, cache_path, pickle_protocol=4)


def domain_data_loader(dataset, domains, file_path, batch_size, train_max_rows=np.inf, valid_max_rows=np.inf,
                       test_max_rows=np.inf, valid_split=0, test_split=0, is_src=True):
    entire_datasets = []
    train_datasets = []

    valid_datasets = []
    test_datasets = []
    st = time.time()

    if domains is not None:
        if domains == 'src':
            processed_domains = conf.args.opt['src_domains']
        elif isinstance(domains, (list,)):
            processed_domains = domains
        else:
            processed_domains = [domains]
    elif is_src:

        if conf.args.validation:
            processed_domains = sorted(
                list(set(conf.args.opt['src_domains']) - set([conf.args.tgt])))  # Leave-one-user-out
        else:
            processed_domains = conf.args.opt['src_domains']
    else:
        if conf.args.validation:
            processed_domains = conf.args.opt['src_domains']
        else:
            processed_domains = conf.args.opt['tgt_domains']

    ##-- load dataset per each domain
    print('Domains:{}'.format(processed_domains))

    if dataset in ['imagenetnoisy', 'cifar10noisy', 'cifar100noisy']:

        cond = processed_domains
        transform = 'src' if is_src else 'val'
        loaded_data = NoisyDataset(base=dataset, domains=cond, transform=transform,
                                   noisy_data=conf.args.noisy_type, noisy_data_size=conf.args.noisy_size,
                                   noisy_data_class=conf.args.noisy_class)
        train_data = loaded_data
        entire_datasets.append(train_data)

    elif dataset in ['cifar10']:

        cond = processed_domains

        transform = 'src' if is_src else 'val'
        loaded_data = load_cache(dataset, processed_domains, file_path, transform=transform)

        if not loaded_data:
            loaded_data = CIFAR10Dataset(domains=cond, transform=transform)
            save_cache(loaded_data, dataset, processed_domains, file_path, transform=transform)

        train_data = loaded_data
        entire_datasets.append(train_data)

    elif dataset in ['cifar100']:

        cond = processed_domains

        transform = 'src' if is_src else 'val'
        loaded_data = load_cache(dataset, processed_domains, file_path, transform=transform)

        if not loaded_data:
            loaded_data = CIFAR100Dataset(domains=cond, transform=transform)
            save_cache(loaded_data, dataset, processed_domains, file_path, transform=transform)

        train_data = loaded_data
        entire_datasets.append(train_data)

    elif dataset in ['imagenet']:

        cond = processed_domains
        transform = 'src' if is_src else 'val'

        file_path = os.path.join(file_path, 'imagenet', cond[0], str(5))
        loaded_data = load_cache(dataset, processed_domains, file_path, transform=transform)

        if not loaded_data:
            loaded_data = ImageNetDataset(domain=cond[0], transform=transform)
            save_cache(loaded_data, dataset, processed_domains, file_path, transform=transform)

        train_data = loaded_data
        entire_datasets.append(train_data)

    else:
        raise NotImplementedError

    ##-- split each dataset into train, valid, and test
    for train_data in entire_datasets:
        total_len = len(train_data)

        train_data, valid_data, test_data = split_data(train_data, valid_split, test_split, train_max_rows,
                                                       valid_max_rows, test_max_rows)

        train_datasets.append(train_data)
        valid_datasets.append(valid_data)
        test_datasets.append(test_data)

        print('#Multi?:{:d} data_loader len:{:d} Train: {:d}\t# Valid: {:d}\t# Test: {:d}'.format(
            True if domains == ['rest'] else False, total_len, len(train_data), len(valid_data),
            len(test_data)))

    print('# Time: {:f} secs'.format(time.time() - st))

    if is_src:
        train_data_loader = datasets_to_dataloader(train_datasets, batch_size=batch_size, concat=True,
                                                   drop_last=True,
                                                   shuffle=True)  # Drop_last for avoiding one-sized minibatches for batchnorm layers
    else:
        train_data_loader = datasets_to_dataloader(train_datasets, batch_size=1, concat=True,
                                                   drop_last=False,
                                                   shuffle=False)
    valid_data_loader = datasets_to_dataloader(valid_datasets, batch_size=batch_size, concat=True,
                                               shuffle=False)
    test_data_loader = datasets_to_dataloader(test_datasets, batch_size=batch_size, concat=True, shuffle=False)

    data_loader = {
        'train': train_data_loader,
        'valid': valid_data_loader,
        'test': test_data_loader,
        'num_domains': sum([dataset.dataset.get_num_domains() for dataset in train_datasets]),
    }
    print('num_domains:' + str(data_loader['num_domains']))
    return data_loader


if __name__ == '__main__':
    pass

data_loader/data_loader.py
taeckyung/SoTTA
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

import conf
from utils import memory
from .dnn import DNN

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class BN_Stats(DNN):

    def __init__(self, *args, **kwargs):
        super(BN_Stats, self).__init__(*args, **kwargs)

        # turn on grad for BN params only
        for module in self.net.modules():
            # print(module)
            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html

                if conf.args.use_learned_stats:  # use learned stats + update with target via momentum
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum

                else:  # Default: use the target batch stats
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None
        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """
        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)
                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        assert (conf.args.update_every_x == conf.args.memory_size)

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data
                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if evaluation:
            self.evaluation_online(current_num_sample, self.mem.get_memory())

        self.batch_instances = []

        if conf.args.use_learned_stats:  # update batch stats with momentum
            self.net.train()

            feats, cls, dls = self.mem.get_memory()
            feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)
            dataset = torch.utils.data.TensorDataset(feats)
            data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                     shuffle=True,
                                     drop_last=False, pin_memory=False)

            with torch.no_grad():

                for batch_idx, (feats,) in enumerate(data_loader):
                    feats = feats.to(device)
                    _ = self.net(feats)  # update batch stats

        ################################ TT_Single_stats: Mazankiewicz, A., Böhm, K., & Bergés, M. (2020). Incremental Real-Time Personalization In Human Activity Recognition Using Domain Adaptive Batch Normalization. ArXiv, 4(4). https://doi.org/10.1145/3432230

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

learner/bn_stats.py
taeckyung/SoTTA
import copy

import conf
from .dnn import DNN
from torch.utils.data import DataLoader
from utils import memory

from utils.loss_functions import *
import PIL
from utils import cotta_utils
import torchvision.transforms as transforms

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


@torch.jit.script
def softmax_entropy(x, x_ema):  # -> torch.Tensor:
    """Entropy of softmax distribution from logits."""
    return -(x_ema.softmax(1) * x.log_softmax(1)).sum(1)


def update_ema_variables(ema_model, model, alpha_teacher):
    for ema_param, param in zip(ema_model.parameters(), model.parameters()):
        ema_param.data[:] = alpha_teacher * ema_param[:].data[:] + (1 - alpha_teacher) * param[:].data[:]
    return ema_model


def get_tta_transforms(gaussian_std: float = 0.005, soft=False, clip_inputs=False):
    if conf.args.dataset in ['cifar10', 'cifar100', 'cifar10noisy', 'cifar100noisy']:
        img_shape = (32, 32, 3)
    else:
        img_shape = (224, 224, 3)
    n_pixels = img_shape[0]

    clip_min, clip_max = 0.0, 1.0

    p_hflip = 0.5

    tta_transforms = transforms.Compose([
        cotta_utils.Clip(0.0, 1.0),
        cotta_utils.ColorJitterPro(
            brightness=[0.8, 1.2] if soft else [0.6, 1.4],
            contrast=[0.85, 1.15] if soft else [0.7, 1.3],
            saturation=[0.75, 1.25] if soft else [0.5, 1.5],
            hue=[-0.03, 0.03] if soft else [-0.06, 0.06],
            gamma=[0.85, 1.15] if soft else [0.7, 1.3]
        ),
        transforms.Pad(padding=int(n_pixels / 2), padding_mode='edge'),
        transforms.RandomAffine(
            degrees=[-8, 8] if soft else [-15, 15],
            translate=(1 / 16, 1 / 16),
            scale=(0.95, 1.05) if soft else (0.9, 1.1),
            shear=None,
            resample=PIL.Image.BILINEAR,
            fillcolor=None
        ),
        transforms.GaussianBlur(kernel_size=5, sigma=[0.001, 0.25] if soft else [0.001, 0.5]),
        transforms.CenterCrop(size=n_pixels),
        transforms.RandomHorizontalFlip(p=p_hflip),
        cotta_utils.GaussianNoise(0, gaussian_std),
        cotta_utils.Clip(clip_min, clip_max)
    ])
    return tta_transforms


class CoTTA(DNN):
    def __init__(self, *args, **kwargs):
        super(CoTTA, self).__init__(*args, **kwargs)

        for param in self.net.parameters():  # turn on requires_grad for all
            param.requires_grad = True

        for module in self.net.modules():

            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # use of batch stats in train and eval modes: https://github.com/qinenergy/cotta/blob/main/cifar/cotta.py
                if conf.args.use_learned_stats:
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum
                else:
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None

                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

        self.mt = conf.args.ema_factor  # 0.999 for every dataset
        self.rst = conf.args.restoration_factor  # 0.01 for all dataset
        self.ap = conf.args.aug_threshold  # 0.92 for CIFAR10, 0.72 for CIFAR100
        self.episodic = False

        self.net_state = copy.deepcopy(self.net.state_dict())
        self.net_anchor = copy.deepcopy(self.net)
        self.net_ema = copy.deepcopy(self.net)
        self.transform = get_tta_transforms()

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        if conf.args.use_learned_stats and evaluation:  # batch-free inference
            self.evaluation_online(current_num_sample, [[current_sample[0]], [current_sample[1]], [current_sample[2]]])

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if not conf.args.use_learned_stats and evaluation:  # batch-based inference
            self.evaluation_online(current_num_sample, self.mem.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        feats, cls, dls = self.mem.get_memory()
        feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)

        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        for e in range(conf.args.epoch):

            for batch_idx, (x,) in enumerate(data_loader):
                torch.cuda.empty_cache()

                if conf.args.tta_attack_type:
                    x = x.clone().detach()

                x = x.to(device)
                outputs = self.net(x)

                torch.cuda.empty_cache()

                with torch.no_grad():
                    anchor_prob = torch.nn.functional.softmax(self.net_anchor(x), dim=1).max(1)[0]

                torch.cuda.empty_cache()

                standard_ema = self.net_ema(x)

                torch.cuda.empty_cache()

                N = 32
                outputs_emas = []

                # Threshold choice discussed in supplementary
                # enable data augmentation for vision datasets
                if anchor_prob.mean(0) < self.ap:
                    for i in range(N):
                        with torch.no_grad():
                            outputs_ = self.net_ema(self.transform(x)).detach()
                        outputs_emas.append(outputs_)
                    outputs_ema = torch.stack(outputs_emas).mean(0)
                    torch.cuda.empty_cache()
                else:
                    outputs_ema = standard_ema

                del x
                torch.cuda.empty_cache()

                # Student update
                self.optimizer.zero_grad()
                loss = (softmax_entropy(outputs, outputs_ema)).mean(0)
                loss.backward()
                self.optimizer.step()

                # Teacher update
                self.net_ema = update_ema_variables(ema_model=self.net_ema, model=self.net, alpha_teacher=self.mt)
                # Stochastic restore
                if True:
                    for nm, m in self.net.named_modules():
                        for npp, p in m.named_parameters():
                            if npp in ['weight', 'bias'] and p.requires_grad:
                                mask = (torch.rand(p.shape) < self.rst).float().cuda()
                                with torch.no_grad():
                                    p.data = self.net_state[f"{nm}.{npp}"] * mask + p * (1. - mask)

        if evaluation and add_memory:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

learner/cotta.py
taeckyung/SoTTA
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from sklearn.metrics import confusion_matrix
from sklearn.metrics import f1_score
import conf
import torch.nn.functional as F

from data_loader.NoisyDataset import NOISY_CLASS_IDX
from data_loader.data_loader import load_cache, save_cache
from utils import memory

from utils.logging import *
from utils.loss_functions import softmax_entropy, calc_energy
from utils.normalize_layer import *
from utils.sam_optimizer import SAM, sam_collect_params
from utils import memory_rotta
from copy import deepcopy

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(
    conf.args.gpu_idx)  # this prevents unnecessary gpu memory allocation to cuda:0 when using estimator


class DNN:
    def __init__(self, model, source_dataloader, target_dataloader, write_path):
        self.device = device

        # init dataloader
        self.source_dataloader = source_dataloader
        self.target_dataloader = target_dataloader

        if conf.args.dataset in ['cifar10', 'cifar100'] and conf.args.tgt_train_dist == 0:
            self.tgt_train_dist = 4  # Dirichlet is default for non-real-distribution data
        else:
            self.tgt_train_dist = conf.args.tgt_train_dist

        self.target_data_cache()

        self.write_path = write_path

        ################## Init & prepare model###################
        self.conf_list = []

        # Load model
        if 'resnet' in conf.args.model:
            if conf.args.dataset not in ['imagenet', 'imagenetnoisy']:
                num_feats = model.fc.in_features
                if conf.args.noisy_type == "divide":
                    num_class = conf.args.opt['num_class'] - len(conf.args.noisy_class)
                else:
                    num_class = conf.args.opt['num_class']
                model.fc = nn.Linear(num_feats, num_class)  # match class number
            self.net = model
        else:
            self.net = model.Net()

        if conf.args.load_checkpoint_path:  # false if conf.args.load_checkpoint_path==''
            self.load_checkpoint(conf.args.load_checkpoint_path)

        # Add normalization layers
        norm_layer = get_normalize_layer(conf.args.dataset)
        if norm_layer:
            self.net = torch.nn.Sequential(norm_layer, self.net)

        if conf.args.parallel and torch.cuda.device_count() > 1:
            self.net = nn.DataParallel(self.net)

        self.net.to(device)

        ##########################################################

        # init criterions, optimizers, scheduler
        if conf.args.method == 'Src':
            if conf.args.dataset in ['cifar10', 'cifar100', 'cifar10noisy', 'cifar100noisy']:
                self.optimizer = torch.optim.SGD(
                    self.net.parameters(),
                    conf.args.opt['learning_rate'],
                    momentum=conf.args.opt['momentum'],
                    weight_decay=conf.args.opt['weight_decay'],
                    nesterov=True)

                self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=conf.args.epoch * len(
                    self.source_dataloader['train']))
            elif conf.args.dataset in ['tinyimagenet', 'imagenet', 'imagenetnoisy']:
                self.optimizer = torch.optim.SGD(
                    self.net.parameters(),
                    conf.args.opt['learning_rate'],
                    momentum=conf.args.opt['momentum'],
                    weight_decay=conf.args.opt['weight_decay'],
                    nesterov=True)
            else:
                self.optimizer = optim.Adam(self.net.parameters(), lr=conf.args.opt['learning_rate'],
                                            weight_decay=conf.args.opt['weight_decay'])
        elif conf.args.method == 'MEMO':
            if conf.args.dataset in ['cifar10', 'cifar100', 'cifar10noisy', 'imagenetnoisy', 'cifar100noisy']:
                self.optimizer = torch.optim.SGD(
                    self.net.parameters(),
                    conf.args.opt['learning_rate'])
            else:
                raise NotImplementedError

        elif conf.args.method == 'SAR':
            # TODO: set base optimizer dynamically (current: SGD)
            params, _ = sam_collect_params(self.net, freeze_top=True)
            self.optimizer = SAM(params, torch.optim.SGD, lr=conf.args.opt['learning_rate'],
                                 momentum=conf.args.opt['momentum'])
        elif conf.args.method in ['SoTTA', 'LOG'] and conf.args.esm:
            params, _ = sam_collect_params(self.net, freeze_top=True)
            self.optimizer = SAM(params, torch.optim.Adam, rho=0.05, lr=conf.args.opt['learning_rate'],
                                 weight_decay=conf.args.opt['weight_decay'])
        elif conf.args.method == "RoTTA":
            # self.optimizer = optim.Adam(self.net.parameters(),lr=1e-3, betas=(0.9, 0.999), weight_decay=0.0)
            self.optimizer = None
        elif conf.args.method == "EATA":
            self.optimizer = optim.SGD(self.net.parameters(),
                                       lr=conf.args.opt['learning_rate'],
                                       momentum=conf.args.opt['momentum'],
                                       weight_decay=conf.args.opt['weight_decay'])
        else:
            if conf.args.method == 'TENT' and conf.args.dataset == 'imagenetnoisy':  # TENT use SGD for imagenet
                self.optimizer = optim.SGD(self.net.parameters(), lr=conf.args.opt['learning_rate'],
                                           weight_decay=conf.args.opt['weight_decay'])
            else:
                self.optimizer = optim.Adam(self.net.parameters(), lr=conf.args.opt['learning_rate'],
                                            weight_decay=conf.args.opt['weight_decay'])

        self.class_criterion = nn.CrossEntropyLoss()
        self.attack_loss = nn.CrossEntropyLoss()

        # online learning
        if conf.args.memory_type == 'FIFO':
            self.mem = memory.FIFO(capacity=conf.args.memory_size)
        elif conf.args.memory_type == 'HUS':
            self.mem = memory.HUS(capacity=conf.args.memory_size, threshold=conf.args.high_threshold)
        elif conf.args.memory_type == 'CSTU':
            self.mem = memory_rotta.CSTU(capacity=conf.args.memory_size, num_class=conf.args.opt['num_class'],
                                         lambda_t=1, lambda_u=1)  # replace memory with original RoTTA
        elif conf.args.memory_type == 'ConfFIFO':
            self.mem = memory.ConfFIFO(capacity=conf.args.memory_size, threshold=conf.args.high_threshold)

        self.json = {}
        self.l2_distance = []
        self.occurred_class = [0 for i in range(conf.args.opt['num_class'])]

        if conf.args.tta_attack_type:
            self.purturb = None
            self.index_attack = None

        self.fifo = None

    def target_data_cache(self):
        dataset = conf.args.dataset
        cond = conf.args.tgt

        filename = f"{dataset}_{conf.args.noisy_type}_{conf.args.noisy_size}_{conf.args.noisy_class}_" \
                   f"{conf.args.seed}_{conf.args.tgt_train_dist}"

        file_path = conf.args.opt['file_path'] + "_target_train_set"

        self.target_train_set = load_cache(filename, cond, file_path, transform=None)

        if not self.target_train_set:
            self.target_data_processing()
            save_cache(self.target_train_set, filename, cond, file_path, transform=None)

    def target_data_processing(self):

        features = []
        cl_labels = []
        do_labels = []

        for b_i, (feat, cl, dl) in enumerate(self.target_dataloader['train']):
            # must be loaded from dataloader, due to transform in the __getitem__()
            features.append(feat.squeeze(0))
            cl_labels.append(cl.squeeze())
            do_labels.append(dl.squeeze())

        tmp = list(zip(features, cl_labels, do_labels))

        features, cl_labels, do_labels = zip(*tmp)
        features, cl_labels, do_labels = list(features), list(cl_labels), list(do_labels)

        num_class = conf.args.opt['num_class']

        result_feats = []
        result_cl_labels = []
        result_do_labels = []

        # real distribution
        if self.tgt_train_dist == 0:
            num_samples = conf.args.nsample if conf.args.nsample < len(features) else len(features)
            for _ in range(num_samples):
                tgt_idx = 0
                result_feats.append(features.pop(tgt_idx))
                result_cl_labels.append(cl_labels.pop(tgt_idx))
                result_do_labels.append(do_labels.pop(tgt_idx))

        # random distribution
        elif self.tgt_train_dist == 1:
            num_samples = conf.args.nsample if conf.args.nsample < len(features) else len(features)
            for _ in range(num_samples):
                tgt_idx = np.random.randint(len(features))
                result_feats.append(features.pop(tgt_idx))
                result_cl_labels.append(cl_labels.pop(tgt_idx))
                result_do_labels.append(do_labels.pop(tgt_idx))

        # dirichlet distribution
        elif self.tgt_train_dist == 4:
            dirichlet_numchunks = conf.args.opt['num_class']

            tgt_class_list = list(range(num_class))
            if conf.args.noisy_type:
                tgt_class_list += [NOISY_CLASS_IDX]

            # https://github.com/IBM/probabilistic-federated-neural-matching/blob/f44cf4281944fae46cdce1b8bc7cde3e7c44bd70/experiment.py
            min_size = -1
            N = len(features)
            min_size_thresh = 10  # if conf.args.dataset in ['tinyimagenet'] else 10
            while min_size < min_size_thresh:  # prevent any chunk having too less data
                idx_batch = [[] for _ in range(dirichlet_numchunks)]
                idx_batch_cls = [[] for _ in range(dirichlet_numchunks)]  # contains data per each class
                for k in tgt_class_list:
                    cl_labels_np = torch.Tensor(cl_labels).numpy()
                    idx_k = np.where(cl_labels_np == k)[0]
                    np.random.shuffle(idx_k)
                    proportions = np.random.dirichlet(
                        np.repeat(conf.args.dirichlet_beta, dirichlet_numchunks))

                    # balance
                    proportions = np.array([p * (len(idx_j) < N / dirichlet_numchunks) for p, idx_j in
                                            zip(proportions, idx_batch)])
                    proportions = proportions / proportions.sum()
                    proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
                    idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]
                    min_size = min([len(idx_j) for idx_j in idx_batch])

                    # store class-wise data
                    for idx_j, idx in zip(idx_batch_cls, np.split(idx_k, proportions)):
                        idx_j.append(idx)

            sequence_stats = []

            # create temporally correlated toy dataset by shuffling classes
            for chunk in idx_batch_cls:
                cls_seq = list(range(len(tgt_class_list)))
                np.random.shuffle(cls_seq)
                for cls in cls_seq:
                    idx = chunk[cls]
                    result_feats.extend([features[i] for i in idx])
                    result_cl_labels.extend([cl_labels[i] for i in idx])
                    result_do_labels.extend([do_labels[i] for i in idx])
                    sequence_stats.extend(list(np.repeat(cls, len(idx))))

            # trim data if num_sample is smaller than the original data size
            num_samples = conf.args.nsample if conf.args.nsample < len(result_feats) else len(result_feats)
            result_feats = result_feats[:num_samples]
            result_cl_labels = result_cl_labels[:num_samples]
            result_do_labels = result_do_labels[:num_samples]

        elif self.tgt_train_dist == 5:
            num_samples = conf.args.nsample if conf.args.nsample < len(features) else len(features)
            for _ in range(num_samples):
                tgt_idx = 0
                result_feats.append(features.pop(tgt_idx))
                result_cl_labels.append(cl_labels.pop(tgt_idx))
                result_do_labels.append(do_labels.pop(tgt_idx))

        elif self.tgt_train_dist == 6:  # for SAR online imbalanced label shifts
            assert conf.args.dataset in ['imagenet']

            ir = 500000  # just use fixed `ir` for now. Depends on the dataset.

            if conf.args.seed == 2021:
                indices_path = './data_loader/imagenetc_imbalanced/total_{}_ir_{}_class_order_shuffle_yes.npy'.format(
                    100000, ir)
            else:
                indices_path = './data_loader/imagenetc_imbalanced/seed{}_total_{}_ir_{}_class_order_shuffle_yes.npy'.format(
                    conf.args.seed, 100000, ir)

            indices = np.load(indices_path).astype(int).tolist()
            result_feats = [features[i] for i in indices]
            result_cl_labels = [cl_labels[i] for i in indices]
            result_do_labels = [do_labels[i] for i in indices]

        remainder = len(result_feats) % conf.args.update_every_x  # drop leftover samples
        if remainder == 0:
            pass
        else:
            result_feats = result_feats[:-remainder]
            result_cl_labels = result_cl_labels[:-remainder]
            result_do_labels = result_do_labels[:-remainder]

        try:
            self.target_train_set = (torch.stack(result_feats),
                                     torch.stack(result_cl_labels),
                                     torch.stack(result_do_labels))
        except:
            try:
                self.target_train_set = (torch.stack(result_feats),
                                         result_cl_labels,
                                         torch.stack(result_do_labels))
            except:  # for dataset which each image has different shape
                self.target_train_set = (result_feats,
                                         result_cl_labels,
                                         torch.stack(result_do_labels))

    def save_checkpoint(self, epoch, epoch_acc, best_acc, checkpoint_path):
        if isinstance(self.net, nn.Sequential):
            if isinstance(self.net[0], NormalizeLayer):
                cp = self.net[1]
        else:
            cp = self.net

        if isinstance(self.net, nn.DataParallel):
            cp = self.net.module

        torch.save(cp.state_dict(), checkpoint_path)

    def load_checkpoint(self, checkpoint_path):
        self.checkpoint = torch.load(checkpoint_path, map_location=f'cuda:{conf.args.gpu_idx}')
        self.net.load_state_dict(self.checkpoint, strict=True)
        self.net.to(device)

    def get_loss_and_confusion_matrix(self, classifier, criterion, data, label):
        preds_of_data = classifier(data)

        labels = [i for i in range(len(conf.args.opt['classes']))]

        loss_of_data = criterion(preds_of_data, label)
        pred_label = preds_of_data.max(1, keepdim=False)[1]
        cm = confusion_matrix(label.cpu(), pred_label.cpu(), labels=labels)
        return loss_of_data, cm, preds_of_data

    def get_loss_cm_error(self, classifier, criterion, data, label):
        preds_of_data = classifier(data)
        labels = [i for i in range(len(conf.args.opt['classes']))]

        loss_of_data = criterion(preds_of_data, label)
        pred_label = preds_of_data.max(1, keepdim=False)[1]
        assert (len(label) == len(pred_label))
        cm = confusion_matrix(label.cpu(), pred_label.cpu(), labels=labels)
        errors = [0 if label[i] == pred_label[i] else 1 for i in range(len(label))]
        return loss_of_data, cm, errors

    def log_loss_results(self, condition, epoch, loss_avg):

        if condition == 'train_online':
            # print loss
            print('{:s}: [current_sample: {:d}]'.format(
                condition, epoch
            ))
        else:
            # print loss
            print('{:s}: [epoch: {:d}]\tLoss: {:.6f} \t'.format(
                condition, epoch, loss_avg
            ))

        return loss_avg

    def log_accuracy_results(self, condition, suffix, epoch, cm_class):

        assert (condition in ['valid', 'test'])
        # assert (suffix in ['labeled', 'unlabeled', 'test'])

        class_accuracy = 100.0 * np.sum(np.diagonal(cm_class)) / np.sum(cm_class)

        print('[epoch:{:d}] {:s} {:s} class acc: {:.3f}'.format(epoch, condition, suffix, class_accuracy))

        return class_accuracy

    def train(self, epoch):
        """
        Train the model
        """

        # setup models

        self.net.train()

        class_loss_sum = 0.0

        total_iter = 0

        if conf.args.method in ['Src', 'Src_Tgt']:
            num_iter = len(self.source_dataloader['train'])
            total_iter += num_iter

            for batch_idx, labeled_data in tqdm(enumerate(self.source_dataloader['train']), total=num_iter):
                feats, cls, _ = labeled_data
                feats, cls = feats.to(device), cls.to(device)

                if torch.isnan(feats).any() or torch.isinf(feats).any():  # For reallifehar debugging
                    print('invalid input detected at iteration ', batch_idx)
                    exit(1)
                # compute the feature
                preds = self.net(feats)
                if torch.isnan(preds).any() or torch.isinf(preds).any():  # For reallifehar debugging
                    print('invalid input detected at iteration ', batch_idx)
                    exit(1)
                class_loss = self.class_criterion(preds, cls)
                class_loss_sum += float(class_loss * feats.size(0))

                if torch.isnan(class_loss).any() or torch.isinf(class_loss).any():  # For reallifehar debugging
                    print('invalid input detected at iteration ', batch_idx)
                    exit(1)

                self.optimizer.zero_grad()
                class_loss.backward()
                self.optimizer.step()
                if conf.args.dataset in ['cifar10', 'cifar100', 'cifar10noisy', 'cifar100noisy']:
                    self.scheduler.step()

        ######################## LOGGING #######################

        self.log_loss_results('train', epoch=epoch, loss_avg=class_loss_sum / total_iter)
        avg_loss = class_loss_sum / total_iter
        return avg_loss

    def logger(self, name, value, epoch, condition):

        if not hasattr(self, name + '_log'):
            exec(f'self.{name}_log = []')
            exec(f'self.{name}_file = open(self.write_path + name + ".txt", "w")')

        exec(f'self.{name}_log.append(value)')

        if isinstance(value, torch.Tensor):
            value = value.item()
        write_string = f'{epoch}\t{value}\n'
        exec(f'self.{name}_file.write(write_string)')

    def evaluation(self, epoch, condition):
        # Evaluate with a batch of samples, which is a typical way of evaluation. Used for pre-training or offline eval.

        self.net.eval()

        with torch.no_grad():
            inputs, cls, dls = self.target_train_set
            tgt_inputs = inputs.to(device)
            tgt_cls = cls.to(device)

            preds = self.net(tgt_inputs)

            labels = [i for i in range(len(conf.args.opt['classes']))]

            class_loss_of_test_data = self.class_criterion(preds, tgt_cls)
            y_pred = preds.max(1, keepdim=False)[1]
            class_cm_test_data = confusion_matrix(tgt_cls.cpu(), y_pred.cpu(), labels=labels)

        print('{:s}: [epoch : {:d}]\tLoss: {:.6f} \t'.format(
            condition, epoch, class_loss_of_test_data
        ))
        class_accuracy = 100.0 * np.sum(np.diagonal(class_cm_test_data)) / np.sum(class_cm_test_data)
        print('[epoch:{:d}] {:s} {:s} class acc: {:.3f}'.format(epoch, condition, 'test', class_accuracy))

        self.logger('accuracy', class_accuracy, epoch, condition)
        self.logger('loss', class_loss_of_test_data, epoch, condition)

        return class_accuracy, class_loss_of_test_data, class_cm_test_data

    def copy_model_and_optimizer(self, model, optimizer):
        """Copy the model and optimizer states for resetting after adaptation."""
        model_state = deepcopy(model.state_dict())
        optimizer_state = deepcopy(optimizer.state_dict())
        return model_state, optimizer_state

    def load_model_and_optimizer(self, model, optimizer, model_state, optimizer_state):
        """Restore the model and optimizer states from copies."""
        model.load_state_dict(model_state, strict=True)
        optimizer.load_state_dict(optimizer_state)

    def evaluation_online_body(self, epoch, current_samples, feats, cls, dls):

        if conf.args.log_grad:
            model_state, optimizer_state = self.copy_model_and_optimizer(self.net, self.optimizer)

        cls = cls.to(torch.int32)  # bugfix when comparing noisy_index
        feats, cls, dls = feats.to(device), cls.to(device), dls.to(device)
        grad_norms = 0.0

        if conf.args.method == 'LAME':
            y_pred = self.batch_evaluation(feats)  # already softmax
            y_conf = y_pred.max(-1)[0].cpu().numpy()
            y_pred = y_pred.argmax(-1).cpu().numpy()

            entropy = np.zeros_like(y_pred)  # we don't support this now
            y_energy = np.zeros_like(entropy)  # we don't support this now

        elif conf.args.method == 'CoTTA':
            x = feats
            anchor_prob = torch.nn.functional.softmax(self.net_anchor(x), dim=1).max(1)[0]
            standard_ema = self.net_ema(x)

            N = 32
            outputs_emas = []

            # Threshold choice discussed in supplementary
            # enable data augmentation for vision datasets
            if anchor_prob.mean(0) < self.ap:
                for i in range(N):
                    outputs_ = self.net_ema(self.transform(x)).detach()
                    outputs_emas.append(outputs_)
                outputs_ema = torch.stack(outputs_emas).mean(0)
            else:
                outputs_ema = standard_ema
            y_pred = outputs_ema
            entropy = softmax_entropy(y_pred)
            y_conf = F.softmax(y_pred, dim=1).max(1, keepdim=False)[0]
            y_energy = calc_energy(y_pred).cpu()
            y_pred = y_pred.max(1, keepdim=False)[1]

        elif conf.args.method == 'TPT':
            with torch.cuda.amp.autocast():
                y_pred = self.net(feats)
                entropy = softmax_entropy(y_pred)
                y_conf = F.softmax(y_pred, dim=1).max(1, keepdim=False)[0]
                y_energy = calc_energy(y_pred).cpu()
                y_pred = y_pred.max(1, keepdim=False)[1]

        else:
            y_pred = self.net(feats)
            entropy = softmax_entropy(y_pred)
            y_conf = F.softmax(y_pred, dim=1).max(1, keepdim=False)[0]
            y_energy = calc_energy(y_pred).cpu()
            y_pred = y_pred.max(1, keepdim=False)[1]

            if conf.args.log_grad:
                if isinstance(self.optimizer, SAM):
                    self.optimizer.zero_grad()
                    # filtering reliable samples/gradients for further adaptation; first time forward
                    loss = entropy.mean(0)
                    loss.backward()

                    # compute \hat{\epsilon(\Theta)} for first order approximation, Eqn. (4)
                    self.optimizer.first_step(zero_grad=True)

                    # second time backward, update model weights using gradients at \Theta+\hat{\epsilon(\Theta)}
                    entropys2 = softmax_entropy(self.net(feats))
                    loss_second = entropys2.mean(0)
                    loss_second.backward()

                    grad_norms = self.optimizer._grad_norm().item()
                else:
                    self.optimizer.zero_grad()
                    loss = entropy.mean(0)
                    loss.backward()
                    grad_norms = torch.norm(
                        torch.stack([
                            (1.0 * p.grad).norm(p=2)
                            for p in self.net.parameters()
                            if p.grad is not None
                        ]),
                        p=2
                    ).item()

        ###################### SAVE RESULT
        # get lists from json

        try:
            true_cls_list = self.json['gt']
            pred_cls_list = self.json['pred']
            accuracy_list = self.json['accuracy']
            f1_macro_list = self.json['f1_macro']
            distance_l2_list = self.json['distance_l2']
            noisy_pred_cls_list = self.json['pred_noisy']
            conf_list = self.json['confidence']
            noisy_conf_list = self.json['confidence_noisy']
            entropy_list = self.json['entropy']
            noisy_entropy_list = self.json['entropy_noisy']
            ood_pred = self.json['ood_pred']
            ood_gt = self.json['ood_gt']
            gt_noisy = self.json['gt_noisy']
            energy_list = self.json['energy']
            noisy_energy_list = self.json['energy_noisy']
            grad_list = self.json['grad']
            noisy_grad_list = self.json['grad_noisy']
            total_pred_list = self.json['pred_total']
            total_conf_list = self.json['confidence_total']
            total_entropy_list = self.json['entropy_total']
            total_grad_list = self.json['grad_total']
        except KeyError:
            true_cls_list = []
            pred_cls_list = []
            accuracy_list = []
            f1_macro_list = []
            distance_l2_list = []
            noisy_pred_cls_list = []
            conf_list = []
            noisy_conf_list = []
            entropy_list = []
            noisy_entropy_list = []
            ood_pred = []
            ood_gt = []
            gt_noisy = []
            energy_list = []
            noisy_energy_list = []
            grad_list = []
            noisy_grad_list = []
            total_pred_list = []
            total_conf_list = []
            total_entropy_list = []
            total_grad_list = []

        if conf.args.noisy_type:
            cls = cls.cpu().to(torch.int32)
            noisy_y_pred = y_pred[cls == NOISY_CLASS_IDX]
            y_pred = y_pred[cls != NOISY_CLASS_IDX]

            noisy_pred_cls_list += [int(c) for c in noisy_y_pred.tolist()]
            noisy_conf_list += [float(c) for c in y_conf[cls == NOISY_CLASS_IDX].tolist()]
            noisy_entropy_list += [float(c) for c in entropy[cls == NOISY_CLASS_IDX].tolist()]
            noisy_energy_list += [float(c) for c in y_energy[cls == NOISY_CLASS_IDX].tolist()]
            if conf.args.log_grad:
                if cls[0] != NOISY_CLASS_IDX:
                    grad_list += [float(grad_norms)]
                else:
                    noisy_grad_list += [float(grad_norms)]

        # append values to lists
        true_cls_list += [int(c) for c in cls if c != NOISY_CLASS_IDX]
        pred_cls_list += [int(c) for c in y_pred.tolist() if c != NOISY_CLASS_IDX]
        conf_list += [float(c) for c in y_conf[cls != NOISY_CLASS_IDX].tolist()]
        entropy_list += [float(c) for c in entropy[cls != NOISY_CLASS_IDX].tolist()]
        energy_list += [float(c) for c in y_energy[cls != NOISY_CLASS_IDX].tolist()]

        total_pred_list += [int(c) for c in y_pred.tolist()]
        total_conf_list += [float(c) for c in y_conf.tolist()]
        total_entropy_list += [float(c) for c in entropy.tolist()]
        total_grad_list += [float(grad_norms)]

        if len(true_cls_list) > 0:
            cumul_accuracy = sum(1 for gt, pred in zip(true_cls_list, pred_cls_list) if gt == pred) / float(
                len(true_cls_list)) * 100
            accuracy_list.append(cumul_accuracy)
            f1_macro_list.append(f1_score(true_cls_list, pred_cls_list,
                                          average='macro'))

            self.occurred_class = [0 for i in range(conf.args.opt['num_class'])]

            # epoch: 1~len(self.target_train_set[0])
            progress_checkpoint = [int(i * (len(self.target_train_set[0]) / 100.0)) for i in range(1, 101)]
            for i in range(epoch + 1 - len(current_samples[0]), epoch + 1):  # consider a batch input
                if i in progress_checkpoint:
                    print(
                        f'[Online Eval][NumSample:{i}][Epoch:{progress_checkpoint.index(i) + 1}][Accuracy:{cumul_accuracy}]')

        # update self.json file
        self.json = {
            'gt': true_cls_list,
            'pred': pred_cls_list,
            'accuracy': accuracy_list,
            'f1_macro': f1_macro_list,
            'distance_l2': distance_l2_list,
            'pred_noisy': noisy_pred_cls_list,
            'confidence': conf_list,
            'confidence_noisy': noisy_conf_list,
            'entropy': entropy_list,
            'entropy_noisy': noisy_entropy_list,
            'ood_pred': ood_pred,
            'ood_gt': ood_gt,
            'gt_noisy': gt_noisy,
            'energy': energy_list,
            'energy_noisy': noisy_energy_list,
            'grad': grad_list,
            'grad_noisy': noisy_grad_list,
            'pred_total': total_pred_list,
            'confidence_total': total_conf_list,
            'entropy_total': total_entropy_list,
            'grad_total': total_grad_list
        }

        if conf.args.log_grad:
            self.load_model_and_optimizer(self.net, self.optimizer, model_state, optimizer_state)

    def evaluation_online(self, epoch, current_samples):
        # Evaluate with online samples that come one by one while keeping the order.
        self.net.eval()

        if conf.args.log_grad:  # run for each sample to log per-sample gradient
            for feature, cl_label, do_label in zip(*current_samples):
                feats, cls, dls = (torch.stack([feature]), torch.stack([cl_label]), torch.stack([do_label]))
                self.evaluation_online_body(epoch, current_samples, feats, cls, dls)
        else:  # batch-based
            with torch.no_grad():  # we don't log grad here
                # extract each from list of current_sample
                features, cl_labels, do_labels = current_samples
                feats, cls, dls = (torch.stack(features), torch.stack(cl_labels), torch.stack(do_labels))
                self.evaluation_online_body(epoch, current_samples, feats, cls, dls)

    def dump_eval_online_result(self, is_train_offline=False):

        if is_train_offline:

            feats, cls, dls = self.target_train_set

            batchsize = conf.args.opt['batch_size']

            for num_sample in range(0, len(feats), batchsize):
                current_sample = feats[num_sample:num_sample + batchsize], cls[num_sample:num_sample + batchsize], dls[
                                                                                                                   num_sample:num_sample + batchsize]
                self.evaluation_online(num_sample + batchsize,
                                       [list(current_sample[0]), list(current_sample[1]), list(current_sample[2])])

        # logging json files
        json_file = open(self.write_path + 'online_eval.json', 'w')
        json_subsample = {key: self.json[key] for key in self.json.keys() - {'extracted_feat'}}
        json_file.write(to_json(json_subsample))
        json_file.close()

    def validation(self, epoch):
        """
        Validate the performance of the model
        """
        class_accuracy_of_test_data, loss, _ = self.evaluation(epoch, 'valid')

        return class_accuracy_of_test_data, loss

    def test(self, epoch):
        """
        Test the performance of the model
        """

        #### for test data
        class_accuracy_of_test_data, loss, cm_class = self.evaluation(epoch, 'test')

        return class_accuracy_of_test_data, loss

    def evaluation_online_num_sample(self, current_num_sample):
        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]
        self.evaluation_online(current_num_sample, [[current_sample[0]], [current_sample[1]], [current_sample[2]]])
        self.log_loss_results('detect ood : skip tta', epoch=current_num_sample, loss_avg=0)
        return SKIPPED

    # implementation of the paper on 3 attack scenarios on TTA
    # implement version assume that params_{t} are almost similar with params_{t+1}
    def reset_purturb(self):
        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        # get data from evaluation part
        if self.fifo:
            feats, cls, dls = self.fifo.get_memory()
        else:
            raise NotImplementedError

        if conf.args.noisy_type == "repeat":
            self.attack_index = []
            self.benign_index = []
            for i in range(len(cls)):
                if int(cls[i]) == int(NOISY_CLASS_IDX):
                    self.attack_index.append(i)
                else:
                    self.benign_index.append(i)
            num_noisy = len(self.attack_index)

        else:
            num_noisy = conf.args.tta_attack_num_samples
            shuffled_index = torch.randperm(len(cls))
            attack_index = torch.sort(shuffled_index[:num_noisy]).values
            benign_index = torch.sort(shuffled_index[num_noisy:]).values
            self.attack_index = attack_index.tolist()
            self.benign_index = benign_index.tolist()

        self.benign_cls = torch.tensor([cls[i] for i in self.benign_index]).long().to(device)

        from torch.autograd import Variable

        self.purturb = []
        for _ in range(num_noisy):
            currnet_purturb = Variable(torch.zeros(feats[0].shape))
            currnet_purturb.requires_grad = True
            currnet_purturb = currnet_purturb.to(device)
            self.purturb.append(currnet_purturb)

        return TRAINED

    def tta_attack_train(self, current_num_sample, step):

        if len(self.purturb) == 0:  # No noisy samples in the batch
            return

        if self.fifo:
            feats, cls, dls = self.fifo.get_memory()
        else:
            raise NotImplementedError

        ls_feats_updated = []
        index_puturb = 0

        self.mem.set_memory(self.mem_state)  # reset mem for adding new generated attack samples

        for i in range(len(cls)):

            if i in self.attack_index:
                ls_feats_updated.append(torch.clamp(feats[i].to(device) + self.purturb[index_puturb], min=0, max=1))
                index_puturb += 1
            else:
                ls_feats_updated.append(feats[i].to(device))

            current_sample = (ls_feats_updated[-1], cls[i], dls[i])
            # add attacked samples in mem for adaptation
            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS', 'ConfFIFO']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)
                    logit = self.net(f.unsqueeze(0))
                    pseudo_cls = logit.max(1, keepdim=False)[1][0].cpu().numpy()
                    pseudo_conf = F.softmax(logit, dim=1).max(1, keepdim=False)[0][0].cpu().numpy()
                    self.mem.add_instance([f, pseudo_cls, d, pseudo_conf])

                elif conf.args.memory_type in ['CSTU']:
                    # RoTTA; Category-balanced Sampling with Timeliness and Uncertainty
                    with torch.no_grad():
                        f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(
                            device)
                        self.net.eval()
                        if conf.args.method in ['RoTTA']:
                            self.net_not_ema.eval()
                        ema_out = self.net(f.unsqueeze(0))
                        predict = torch.softmax(ema_out, dim=1)
                        pseudo_label = torch.argmax(predict, dim=1)
                        entropy = torch.sum(- predict * torch.log(predict + 1e-6), dim=1)

                    # add into memory
                    for i, data in enumerate(f.unsqueeze(0)):
                        p_l = pseudo_label[i].item()
                        uncertainty = entropy[i].item()
                        current_instance = (data, p_l, uncertainty)
                        self.mem.add_instance(current_instance)

                else:
                    raise NotImplementedError

        self.train_online(current_num_sample, add_memory=False, evaluation=False)

        feats_updated = torch.stack(ls_feats_updated).to(device)

        if conf.args.method == "SRC_FOR_TTA_ATTACK":
            self.tta_attack_mode_turn_on()

        self.net.eval()

        preds = self.net(feats_updated)

        # loss calculation
        if conf.args.tta_attack_type == "targeted":
            raise NotImplementedError

        elif conf.args.tta_attack_type == "indiscriminate":
            loss = -self.attack_loss(preds[self.benign_index], self.benign_cls)

            acc = torch.sum(torch.argmax(preds[self.benign_index], axis=1) == self.benign_cls) / len(self.benign_cls)
            print('{:s}: [epoch: {:d}]\tLoss: {:.6f} \tAcc: {:.6f}'.format("tta_attack_train", step, loss, acc))

        elif conf.args.tta_attack_type == "stealthy":
            raise NotImplementedError

        else:
            raise NotImplementedError

        # update perturb
        alpha = 1.0 / 255

        grad = torch.autograd.grad(loss, self.purturb, retain_graph=False, create_graph=False)

        for i in range(len(self.purturb)):
            self.purturb[i] -= alpha * torch.sign(grad[i])
            self.purturb[i] = torch.clamp(self.purturb[i], min=-conf.args.tta_attack_eps, max=conf.args.tta_attack_eps)

    def tta_attack_update(self, current_num_sample):  # update memory with final perturbation
        # get data from evaluation part
        if self.fifo:
            feats, cls, dls = self.fifo.get_memory()
        else:
            raise NotImplementedError

        feats_updated = []
        index_puturb = 0

        self.mem.set_memory(self.mem_state)  # reset mem for adding new generated attack samples

        for i in range(len(cls)):

            if i in self.attack_index:
                feats_updated.append(torch.clamp(feats[i] + self.purturb[index_puturb].detach().cpu(), min=0, max=1))
                index_puturb += 1
            else:
                feats_updated.append(feats[i])

            # add attacked samples in mem for adaptation
            current_sample = (feats_updated[-1], cls[i], dls[i])
            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS', 'ConfFIFO']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)
                    logit = self.net(f.unsqueeze(0))
                    pseudo_cls = logit.max(1, keepdim=False)[1][0].cpu().numpy()
                    pseudo_conf = F.softmax(logit, dim=1).max(1, keepdim=False)[0][0].cpu().numpy()
                    self.mem.add_instance([f, pseudo_cls, d, pseudo_conf])

                elif conf.args.memory_type in [
                    'CSTU']:  # RoTTA; Category-balanced Sampling with Timeliness and Uncertainty
                    with torch.no_grad():
                        f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(
                            device)
                        self.net.eval()
                        if conf.args.method in ['RoTTA']:
                            self.net_not_ema.eval()
                        ema_out = self.net(f.unsqueeze(0))
                        predict = torch.softmax(ema_out, dim=1)
                        pseudo_label = torch.argmax(predict, dim=1)
                        entropy = torch.sum(- predict * torch.log(predict + 1e-6), dim=1)

                    # add into memory
                    for i, data in enumerate(f.unsqueeze(0)):
                        p_l = pseudo_label[i].item()
                        uncertainty = entropy[i].item()
                        current_instance = (data, p_l, uncertainty)
                        self.mem.add_instance(current_instance)

                else:
                    raise NotImplementedError

        # update fifo memory
        dic = {'data': [feats_updated, cls, dls]}
        self.fifo.set_memory(dic)
        self.mem_state = self.mem.save_state_dict()  # update final mem_state

        if conf.args.method == "SRC_FOR_TTA_ATTACK":
            self.tta_attack_mode_turn_off()
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        else:
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        return

learner/dnn.py
taeckyung/SoTTA
import conf
from data_loader.NoisyDataset import NoisyDataset
from data_loader.data_loader import datasets_to_dataloader
from .dnn import DNN
from torch.utils.data import DataLoader, random_split

from utils.loss_functions import *
from utils import memory

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class EATA(DNN):
    def __init__(self, *args, **kwargs):
        super(EATA, self).__init__(*args, **kwargs)

        fisher_dataset = NoisyDataset(base=conf.args.dataset, domains=['test'], transform='val',
                                      noisy_data=conf.args.noisy_type, noisy_data_size=conf.args.noisy_size,
                                      noisy_data_class=conf.args.noisy_class)
        fisher_dataset = \
        random_split(fisher_dataset, [conf.args.fisher_size, len(fisher_dataset) - conf.args.fisher_size])[0]
        fisher_loader = datasets_to_dataloader([fisher_dataset], batch_size=64, concat=True, shuffle=True)

        subnet = configure_model(self.net)
        params, param_names = collect_params(subnet)
        ewc_optimizer = torch.optim.SGD(params, 0.001)
        fishers = {}
        train_loss_fn = nn.CrossEntropyLoss().cuda()
        for iter_, (images, targets, domains) in enumerate(fisher_loader, start=1):
            if conf.args.gpu_idx is not None:
                images = images.cuda(conf.args.gpu_idx, non_blocking=True)
            # if torch.cuda.is_available():
            #     targets = targets.cuda(conf.args.gpu_idx, non_blocking=True)
            outputs = subnet(images)
            _, targets = outputs.max(1)
            loss = train_loss_fn(outputs, targets)
            loss.backward()
            for name, param in subnet.named_parameters():
                if param.grad is not None:
                    if iter_ > 1:
                        fisher = param.grad.data.clone().detach() ** 2 + fishers[name][0]
                    else:
                        fisher = param.grad.data.clone().detach() ** 2
                    if iter_ == len(fisher_loader):
                        fisher = fisher / iter_
                    fishers.update({name: [fisher, param.data.clone().detach()]})
            ewc_optimizer.zero_grad()
        # logger.info("compute fisher matrices finished")
        del ewc_optimizer

        # self.steps = 1 # SoTTA: replaced to epoch
        # assert self.steps > 0, "EATA requires >= 1 step(s) to forward and update"
        # self.episodic = False  # SoTTA: we don't use episodic

        self.num_samples_update_1 = 0  # number of samples after First filtering, exclude unreliable samples
        self.num_samples_update_2 = 0  # number of samples after Second filtering, exclude both unreliable and redundant samples
        self.e_margin = conf.args.e_margin  # hyper-parameter E_0 (Eqn. 3)
        self.d_margin = conf.args.d_margin  # hyper-parameter \epsilon for consine simlarity thresholding (Eqn. 5)

        self.current_model_probs = None  # the moving average of probability vector (Eqn. 4)

        self.fishers = fishers  # fisher regularizer items for anti-forgetting, need to be calculated pre model adaptation (Eqn. 9)
        self.fisher_alpha = conf.args.fisher_alpha  # trade-off \beta for two losses (Eqn. 8)

        # note: if the model is never reset, like for continual adaptation,
        # then skipping the state copy would save memory
        self.model_state, self.optimizer_state = self.copy_model_and_optimizer(self.net, self.optimizer)

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)
                else:
                    raise NotImplementedError

        if conf.args.use_learned_stats and evaluation:  # batch-free inference
            self.evaluation_online(current_num_sample,
                                   [[current_sample[0]], [current_sample[1]], [current_sample[2]]])

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if not conf.args.use_learned_stats and evaluation:  # batch-based inference
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        feats, cls, dls = self.mem.get_memory()
        feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)

        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        for e in range(conf.args.epoch):
            for batch_idx, (feats,) in enumerate(data_loader):
                feats = feats.to(device)

                if conf.args.tta_attack_type:
                    feats = feats.clone().detach()

                outputs, num_counts_2, num_counts_1, updated_probs = forward_and_adapt_eata(feats, self.net,
                                                                                            self.optimizer,
                                                                                            self.fishers,
                                                                                            self.e_margin,
                                                                                            self.current_model_probs,
                                                                                            fisher_alpha=self.fisher_alpha,
                                                                                            num_samples_update=self.num_samples_update_2,
                                                                                            d_margin=self.d_margin)
                self.num_samples_update_2 += num_counts_2
                self.num_samples_update_1 += num_counts_1
                self.reset_model_probs(updated_probs)

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

    def reset(self):
        if self.model_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer state")
        self.load_model_and_optimizer(self.net, self.optimizer,
                                      self.model_state, self.optimizer_state)

    def reset_model_probs(self, probs):
        self.current_model_probs = probs


@torch.enable_grad()  # ensure grads in possible no grad context for testing
def forward_and_adapt_eata(x, model, optimizer, fishers, e_margin, current_model_probs, fisher_alpha=50.0,
                           d_margin=0.05, scale_factor=2, num_samples_update=0):
    """Forward and adapt model on batch of data.
    Measure entropy of the model prediction, take gradients, and update params.
    Return:
    1. model outputs;
    2. the number of reliable and non-redundant samples;
    3. the number of reliable samples;
    4. the moving average  probability vector over all previous samples
    """
    # forward
    outputs = model(x)
    # adapt
    entropys = softmax_entropy(outputs)
    # filter unreliable samples
    filter_ids_1 = torch.where(entropys < e_margin)
    ids1 = filter_ids_1
    ids2 = torch.where(ids1[0] > -0.1)
    entropys = entropys[filter_ids_1]
    # filter redundant samples
    if current_model_probs is not None:
        cosine_similarities = F.cosine_similarity(current_model_probs.unsqueeze(dim=0),
                                                  outputs[filter_ids_1].softmax(1), dim=1)
        filter_ids_2 = torch.where(torch.abs(cosine_similarities) < d_margin)
        entropys = entropys[filter_ids_2]
        ids2 = filter_ids_2
        updated_probs = update_model_probs(current_model_probs, outputs[filter_ids_1][filter_ids_2].softmax(1))
    else:
        updated_probs = update_model_probs(current_model_probs, outputs[filter_ids_1].softmax(1))
    coeff = 1 / (torch.exp(entropys.clone().detach() - e_margin))
    # implementation version 1, compute loss, all samples backward (some unselected are masked)
    entropys = entropys.mul(coeff)  # reweight entropy losses for diff. samples
    loss = entropys.mean(0)
    """
    # implementation version 2, compute loss, forward all batch, forward and backward selected samples again.
    # if x[ids1][ids2].size(0) != 0:
    #     loss = softmax_entropy(model(x[ids1][ids2])).mul(coeff).mean(0) # reweight entropy losses for diff. samples
    """
    if fishers is not None:
        ewc_loss = 0
        for name, param in model.named_parameters():
            if name in fishers:
                ewc_loss += fisher_alpha * (fishers[name][0] * (param - fishers[name][1]) ** 2).sum()
        loss += ewc_loss
    if x[ids1][ids2].size(0) != 0:
        loss.backward()
        optimizer.step()
    optimizer.zero_grad()
    return outputs, entropys.size(0), filter_ids_1[0].size(0), updated_probs


@torch.jit.script
def softmax_entropy(x: torch.Tensor) -> torch.Tensor:
    """Entropy of softmax distribution from logits."""
    temprature = 1
    x = x / temprature
    x = -(x.softmax(1) * x.log_softmax(1)).sum(1)
    return x


def update_model_probs(current_model_probs, new_probs):
    if current_model_probs is None:
        if new_probs.size(0) == 0:
            return None
        else:
            with torch.no_grad():
                return new_probs.mean(0)
    else:
        if new_probs.size(0) == 0:
            with torch.no_grad():
                return current_model_probs
        else:
            with torch.no_grad():
                return 0.9 * current_model_probs + (1 - 0.9) * new_probs.mean(0)


def collect_params(model):
    """Collect the affine scale + shift parameters from batch norms.
    Walk the model's modules and collect all batch normalization parameters.
    Return the parameters and their names.
    Note: other choices of parameterization are possible!
    """
    params = []
    names = []
    for nm, m in model.named_modules():
        if isinstance(m, nn.BatchNorm2d):
            for np, p in m.named_parameters():
                if np in ['weight', 'bias']:  # weight is scale, bias is shift
                    params.append(p)
                    names.append(f"{nm}.{np}")
    return params, names


def configure_model(model):
    """Configure model for use with eata."""
    # train mode, because eata optimizes the model to minimize entropy
    model.train()
    # disable grad, to (re-)enable only what eata updates
    model.requires_grad_(False)
    # configure norm for eata updates: enable grad + force batch statisics
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.requires_grad_(True)
            # force use of batch stats in train and eval modes
            m.track_running_stats = False
            m.running_mean = None
            m.running_var = None
    return model


def check_model(model):
    """Check model for compatability with eata."""
    is_training = model.training
    assert is_training, "eata needs train mode: call model.train()"
    param_grads = [p.requires_grad for p in model.parameters()]
    has_any_params = any(param_grads)
    has_all_params = all(param_grads)
    assert has_any_params, "eata needs params to update: " \
                           "check which require grad"
    assert not has_all_params, "eata should not update all params: " \
                               "check which require grad"
    has_bn = any([isinstance(m, nn.BatchNorm2d) for m in model.modules()])
    assert has_bn, "eata needs normalization for its optimization"

learner/eata.py
taeckyung/SoTTA
import conf
from .dnn import DNN
from torch.utils.data import DataLoader
from utils.normalize_layer import *
from utils.loss_functions import *
from utils import memory

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class AffinityMatrix:

    def __init__(self, **kwargs):
        pass

    def __call__(X, **kwargs):
        raise NotImplementedError

    def is_psd(self, mat):
        eigenvalues = torch.eig(mat)[0][:, 0].sort(descending=True)[0]
        return eigenvalues, float((mat == mat.t()).all() and (eigenvalues >= 0).all())

    def symmetrize(self, mat):
        return 1 / 2 * (mat + mat.t())


class kNN_affinity(AffinityMatrix):
    def __init__(self, knn: int, **kwargs):
        self.knn = knn

    def __call__(self, X):
        N = X.size(0)
        dist = torch.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, p=2)  # [N, N]
        n_neighbors = min(self.knn + 1, N)

        knn_index = dist.topk(n_neighbors, -1, largest=False).indices[:, 1:]  # [N, knn]

        W = torch.zeros(N, N, device=X.device)
        W.scatter_(dim=-1, index=knn_index, value=1.0)

        return W


class rbf_affinity(AffinityMatrix):
    def __init__(self, sigma: float, **kwargs):
        self.sigma = sigma
        self.k = kwargs['knn']

    def __call__(self, X):
        N = X.size(0)
        dist = torch.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, p=2)  # [N, N]
        n_neighbors = min(self.k, N)
        kth_dist = dist.topk(k=n_neighbors, dim=-1, largest=False).values[:,
                   -1]  # compute k^th distance for each point, [N, knn + 1]
        sigma = kth_dist.mean()
        rbf = torch.exp(- dist ** 2 / (2 * sigma ** 2))
        # mask = torch.eye(X.size(0)).to(X.device)
        # rbf = rbf * (1 - mask)
        return rbf


class linear_affinity(AffinityMatrix):

    def __call__(self, X: torch.Tensor):
        """
        X: [N, d]
        """
        return torch.matmul(X, X.t())


class LAME(DNN):
    def __init__(self, *args, **kwargs):
        super(LAME, self).__init__(*args, **kwargs)

        self.knn = 5
        self.sigma = 1.0  # from overall_best.yaml in LAME github
        self.affinity = kNN_affinity(knn=self.knn)
        if isinstance(self.net, nn.Sequential):
            if isinstance(self.net[0], NormalizeLayer):
                self.featurizer = torch.nn.Sequential(self.net[0], *list(self.net[1].children())[:-1])
        else:
            self.featurizer = torch.nn.Sequential(*list(self.net.children())[:-1])

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Get sample from target
        xs, cls, dls = self.target_train_set
        current_sample = xs[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        # Add sample to memory
        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        # Skipping depend on "batch size"
        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # setup models
        self.net.eval()

        if len(xs) == 1:  # avoid BN error
            self.net.eval()

        xs, cls, dls = self.mem.get_memory()
        xs, cls, dls = torch.stack(xs), torch.stack(cls), torch.stack(dls)

        dataset = torch.utils.data.TensorDataset(xs, cls, dls)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        if evaluation:
            self.evaluation_online(current_num_sample, self.mem.get_memory())
        else:
            self.net.eval()
            with torch.no_grad():

                # extract each from list of current_sample
                features, cl_labels, do_labels = self.mem.get_memory()
                feats, cls, dls = (torch.stack(features), torch.stack(cl_labels), torch.stack(do_labels))
                cls = cls.to(torch.int32)  # bugfix when comparing noisy index
                feats, cls, dls = feats.to(device), cls.to(device), dls.to(device)

                self.batch_evaluation(feats)

        if evaluation and add_memory:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

    def batch_evaluation(self, extracted_feat):
        out = self.net(extracted_feat)
        unary = -torch.log(out.softmax(-1) + 1e-10)  # softmax the output

        if conf.args.model == 'resnet18':
            feats = self.featurizer(extracted_feat)
            feats = F.avg_pool2d(feats, 4)
            feats = feats.view(feats.size(0), -1)
        else:
            # use special forwarding function instead of featurizer in language models
            if conf.args.model in ['bert', 'bert-tiny', 'bert-mini', 'bert-small', 'bert-medium', 'mobilebert',
                                   'distilbert']:
                if conf.args.parallel:
                    extracted_feat = self.net.module.get_feature(extracted_feat)
                else:
                    extracted_feat = self.net.get_feature(extracted_feat)
            else:
                extracted_feat = self.featurizer(extracted_feat)
            feats = torch.nn.functional.normalize(extracted_feat, p=2, dim=-1).squeeze()

        kernel = self.affinity(feats)

        Y = laplacian_optimization(unary, kernel)
        return Y


def laplacian_optimization(unary, kernel, bound_lambda=1, max_steps=100):
    E_list = []
    oldE = float('inf')
    Y = (-unary).softmax(-1)  # [N, K]
    for i in range(max_steps):
        pairwise = bound_lambda * kernel.matmul(Y)  # [N, K]
        exponent = -unary + pairwise
        Y = exponent.softmax(-1)
        E = entropy_energy(Y, unary, pairwise, bound_lambda).item()
        E_list.append(E)

        if i > 1 and (abs(E - oldE) <= 1e-8 * abs(oldE)):
            break
        else:
            oldE = E

    return Y


def entropy_energy(Y, unary, pairwise, bound_lambda):
    E = (unary * Y - bound_lambda * pairwise * Y + Y * torch.log(Y.clip(1e-20))).sum()
    return E

learner/lame.py
taeckyung/SoTTA
import conf
from .dnn import DNN
from torch.utils.data import DataLoader

from utils.loss_functions import *
from utils import memory

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class PseudoLabel(DNN):
    def __init__(self, *args, **kwargs):
        super(PseudoLabel, self).__init__(*args, **kwargs)

        for param in self.net.parameters():  # initially turn off requires_grad for all
            param.requires_grad = False

        for module in self.net.modules():

            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html
                if conf.args.use_learned_stats:
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum
                else:
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None

                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        if evaluation:
            # Evaluate with a batch
            self.evaluation_online(current_num_sample, self.mem.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        feats, cls, dls = self.mem.get_memory()
        feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)

        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        for e in range(conf.args.epoch):

            for batch_idx, (feats,) in enumerate(data_loader):
                feats = feats.to(device)

                if conf.args.tta_attack_type:
                    feats = feats.clone().detach()

                preds_of_data = self.net(feats)
                pseudo_cls = preds_of_data.max(1, keepdim=False)[1]
                loss = self.class_criterion(preds_of_data, pseudo_cls)

                self.optimizer.zero_grad()
                if conf.args.tta_attack_type:
                    loss.backward(retain_graph=True)
                else:
                    loss.backward()
                self.optimizer.step()

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

learner/pseudo_label.py
taeckyung/SoTTA
from copy import deepcopy

import torch.optim as optim

import conf
from utils import memory, bn_layers_rotta
from utils.custom_transforms import get_tta_transforms
from utils.loss_functions import *
from .dnn import DNN

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class RoTTA(DNN):
    def __init__(self, *args, **kwargs):
        super(RoTTA, self).__init__(*args, **kwargs)

        self.net.requires_grad_(False)
        self.alpha = 0.05
        normlayer_names = []

        for name, sub_module in self.net.named_modules():
            if isinstance(sub_module, nn.BatchNorm1d) or isinstance(sub_module, nn.BatchNorm2d):
                normlayer_names.append(name)

        for name in normlayer_names:
            bn_layer = get_named_submodule(self.net, name)
            if isinstance(bn_layer, nn.BatchNorm1d):
                NewBN = bn_layers_rotta.RobustBN1d
            elif isinstance(bn_layer, nn.BatchNorm2d):
                NewBN = bn_layers_rotta.RobustBN2d
            else:
                raise RuntimeError()

            momentum_bn = NewBN(bn_layer, self.alpha)
            momentum_bn.requires_grad_(True)
            set_named_submodule(self.net, name, momentum_bn)

        params, param_names = self.collect_params(self.net)
        self.optimizer = optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), weight_decay=0.0)

        net_ema = deepcopy(self.net)
        for param in net_ema.parameters():
            param.detach_()

        self.net_not_ema = self.net
        self.net = net_ema  # set self.net to self.net_ema
        self.net.to(device)

        self.transform = get_tta_transforms(tuple(self.target_train_set[0][0].shape[1:]))
        self.nu = 0.001

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():

                self.net.eval()

                if conf.args.memory_type in [
                    'CSTU']:  # RoTTA; Category-balanced Sampling with Timeliness and Uncertainty
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)
                    self.net.eval()
                    self.net_not_ema.eval()
                    ema_out = self.net(f.unsqueeze(0))
                    predict = torch.softmax(ema_out, dim=1)
                    pseudo_label = torch.argmax(predict, dim=1)
                    entropy = torch.sum(- predict * torch.log(predict + 1e-6), dim=1)

                    # add into memory
                    for i, data in enumerate(f.unsqueeze(0)):
                        p_l = pseudo_label[i].item()
                        uncertainty = entropy[i].item()
                        current_instance = (data, p_l, uncertainty)
                        self.mem.add_instance(current_instance)
                        # self.current_instance += 1

                        # if self.current_instance % self.update_frequency == 0:
                        #     self.update_model(model, optimizer)
                else:
                    raise NotImplementedError

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if evaluation:  # batch-based inference
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        # setup models
        self.net.train()
        self.net_not_ema.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()
            self.net_not_ema.eval()

        # get memory data
        # sup_data, ages = self.mem.get_memory_rotta_style()
        sup_data, ages = self.mem.get_memory()

        l_sup = None
        if len(sup_data) > 0:
            sup_data = torch.stack(sup_data)
            strong_sup_aug = self.transform(sup_data)
            ema_sup_out = self.net(sup_data)
            stu_sup_out = self.net_not_ema(strong_sup_aug)
            instance_weight = self.timeliness_reweighting(ages)
            l_sup = (softmax_entropy_rotta(stu_sup_out, ema_sup_out) * instance_weight).mean()

        loss = l_sup
        if loss is not None:
            self.optimizer.zero_grad()
            if conf.args.tta_attack_type:
                loss.backward(retain_graph=True)
            else:
                loss.backward()

            self.optimizer.step()

        self.net = self.update_ema_variables(self.net, self.net_not_ema, self.nu)

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

    def timeliness_reweighting(self, ages):
        if isinstance(ages, list):
            ages = torch.tensor(ages).float().cuda()
        return torch.exp(-ages) / (1 + torch.exp(-ages))

    def update_ema_variables(self, ema_model, model, nu):
        for ema_param, param in zip(ema_model.parameters(), model.parameters()):
            ema_param.data[:] = (1 - nu) * ema_param[:].data[:] + nu * param[:].data[:]
        return ema_model

    def collect_params(self, model: nn.Module):
        names = []
        params = []

        for n, p in model.named_parameters():
            if p.requires_grad:
                names.append(n)
                params.append(p)

        return params, names


def get_named_submodule(model, sub_name: str):
    names = sub_name.split(".")
    module = model
    for name in names:
        module = getattr(module, name)

    return module


def set_named_submodule(model, sub_name, value):
    names = sub_name.split(".")
    module = model
    for i in range(len(names)):
        if i != len(names) - 1:
            module = getattr(module, names[i])

        else:
            setattr(module, names[i], value)

learner/rotta.py
taeckyung/SoTTA
from copy import deepcopy

import math

import conf
from utils import memory
from .dnn import DNN
from torch.utils.data import DataLoader

from utils.loss_functions import *

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


def update_ema(ema, new_data):
    if ema is None:
        return new_data
    else:
        with torch.no_grad():
            return 0.9 * ema + (1 - 0.9) * new_data


def copy_model_and_optimizer(model, optimizer):
    """Copy the model and optimizer states for resetting after adaptation."""
    model_state = deepcopy(model.state_dict())
    optimizer_state = deepcopy(optimizer.state_dict())
    return model_state, optimizer_state


def load_model_and_optimizer(model, optimizer, model_state, optimizer_state):
    """Restore the model and optimizer states from copies."""
    model.load_state_dict(model_state, strict=True)
    optimizer.load_state_dict(optimizer_state)


class SAR(DNN):
    def __init__(self, *args, **kwargs):
        super(SAR, self).__init__(*args, **kwargs)

        # turn on grad for BN params only

        self.net.train()

        for param in self.net.parameters():  # initially turn off requires_grad for all
            param.requires_grad = False

        for module in self.net.modules():
            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                if conf.args.use_learned_stats:
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum
                else:
                    # With below, this module always uses the test batch statistics (no momentum)
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None

                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

            # LayerNorm and GroupNorm for ResNet-GN and Vit-LN models
            # TODO: support use_learned_stats
            if isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
                module.requires_grad_(True)

        # SAR-specific hyperparameters

        if conf.args.noisy_type == "divide":
            num_class = conf.args.opt['num_class'] - len(conf.args.noisy_class)
        else:
            num_class = conf.args.opt['num_class']

        self.margin_e0 = 0.4 * math.log(num_class)  # math.log(1000)
        self.reset_constant_em = 0.2
        self.ema = None

        self.net_state, self.optimizer_state = \
            copy_model_and_optimizer(self.net, self.optimizer)

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

    def reset(self):
        if self.net_state is None or self.optimizer_state is None:
            raise Exception("cannot reset without saved model/optimizer")
        load_model_and_optimizer(self.net, self.optimizer, self.net_state, self.optimizer_state)
        self.ema = None

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        if conf.args.use_learned_stats and evaluation:  # batch-free inference
            self.evaluation_online(current_num_sample,
                                   [[current_sample[0]], [current_sample[1]], [current_sample[2]]])

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if not conf.args.use_learned_stats and evaluation:  # batch-based inference
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        feats, cls, dls = self.mem.get_memory()
        feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)

        if conf.args.tta_attack_type:
            feats = feats.clone().detach()

        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        for e in range(conf.args.epoch):
            for batch_idx, (feats,) in enumerate(data_loader):
                feats = feats.to(device)

                self.optimizer.zero_grad()

                preds_of_data = self.net(feats)

                # filtering reliable samples/gradients for further adaptation; first time forward
                entropys = softmax_entropy(preds_of_data)
                filter_ids_1 = torch.where(entropys < self.margin_e0)
                entropys = entropys[filter_ids_1]
                loss = entropys.mean(0)

                if conf.args.tta_attack_type:
                    loss.backward(retain_graph=True)
                else:
                    loss.backward()

                # compute \hat{\epsilon(\Theta)} for first order approximation, Eqn. (4)
                self.optimizer.first_step(zero_grad=True)

                # second time backward, update model weights using gradients at \Theta+\hat{\epsilon(\Theta)}
                entropys2 = softmax_entropy(self.net(feats))
                entropys2 = entropys2[filter_ids_1]
                filter_ids_2 = torch.where(entropys2 < self.margin_e0)
                loss_second = entropys2[filter_ids_2].mean(0)
                if not np.isnan(loss_second.item()):
                    self.ema = update_ema(self.ema, loss_second.item())

                if conf.args.tta_attack_type:
                    loss_second.backward(retain_graph=True)
                else:
                    loss_second.backward()

                self.optimizer.second_step(zero_grad=True)

                if self.ema is not None and self.ema < 0.2:
                    print("ema < 0.2, now reset the model")
                    ema = self.ema
                    self.reset()
                    self.ema = ema

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

learner/sar.py
taeckyung/SoTTA
from torch.utils.data import DataLoader

import conf
from utils import memory
from utils.loss_functions import *
from utils.sam_optimizer import SAM
from .dnn import DNN

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class SoTTA(DNN):

    def __init__(self, *args, **kwargs):
        super(SoTTA, self).__init__(*args, **kwargs)

        # turn on grad for BN params only

        for param in self.net.parameters():  # initially turn off requires_grad for all
            param.requires_grad = False

        for module in self.net.modules():
            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html

                if conf.args.use_learned_stats:
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum
                else:
                    # With below, this module always uses the test batch statistics (no momentum)
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None

                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

            elif isinstance(module, nn.InstanceNorm1d) or isinstance(module, nn.InstanceNorm2d): 
                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

            elif isinstance(module, nn.LayerNorm):  
                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()

        self.ema = None
        self.batchnorm_stats = []

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Get a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]
        
        # Add into memory
        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():

                self.net.eval()

                if conf.args.memory_type in ['FIFO']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS', 'ConfFIFO']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)
                    logit = self.net(f.unsqueeze(0))
                    pseudo_cls = logit.max(1, keepdim=False)[1][0].cpu().numpy()
                    pseudo_conf = F.softmax(logit, dim=1).max(1, keepdim=False)[0][0].cpu().numpy()
                    self.mem.add_instance([f, pseudo_cls, d, pseudo_conf])

                elif conf.args.memory_type in ['CSTU']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)
                    ema_out = self.net(f.unsqueeze(0))
                    predict = torch.softmax(ema_out, dim=1)
                    pseudo_label = torch.argmax(predict, dim=1)
                    entropy = torch.sum(- predict * torch.log(predict + 1e-6), dim=1)

                    for i, data in enumerate(f.unsqueeze(0)):
                        p_l = pseudo_label[i].item()
                        uncertainty = entropy[i].item()
                        current_instance = (data, p_l, uncertainty)
                        self.mem.add_instance(current_instance)

                else:
                    raise NotImplementedError

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        if evaluation:
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        if conf.args.memory_type in ['CSTU']:
            feats, _ = self.mem.get_memory()
        else:
            feats, _, _ = self.mem.get_memory()

        if len(feats) == 0:
            return TRAINED

        feats = torch.stack(feats)
        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True, drop_last=False, pin_memory=False)

        entropy_loss = HLoss(conf.args.temperature)

        for e in range(conf.args.epoch):
            for batch_idx, (feats,) in enumerate(data_loader):
                self.step(loss_fn=entropy_loss, feats=feats)

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

    def step(self, loss_fn, feats=None):
        assert (feats is not None)

        if conf.args.tta_attack_type: # avoid attack error
            feats = feats.clone().detach()

        self.net.train()
        feats = feats.to(device)
        preds_of_data = self.net(feats)

        loss_first = loss_fn(preds_of_data)

        self.optimizer.zero_grad()

        loss_first.backward()

        if not isinstance(self.optimizer, SAM):
            self.optimizer.step()
        else:
            # compute \hat{\epsilon(\Theta)} for first order approximation, Eqn. (4)
            self.optimizer.first_step(zero_grad=True)

            preds_of_data = self.net(feats)

            # second time backward, update model weights using gradients at \Theta+\hat{\epsilon(\Theta)}
            loss_second = loss_fn(preds_of_data)

            loss_second.backward()

            self.optimizer.second_step(zero_grad=True)

learner/sotta.py
taeckyung/SoTTA
from copy import deepcopy

import conf
from utils import memory
from utils.loss_functions import *
from .dnn import DNN

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class SRC_FOR_TTA_ATTACK(DNN):
    def __init__(self, *args, **kwargs):
        super(SRC_FOR_TTA_ATTACK, self).__init__(*args, **kwargs)

        # turn on grad for BN params only

        for param in self.net.parameters():  # initially turn off requires_grad for all
            param.requires_grad = False

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()
        self.grad_norm_list = []
        self.original_model = deepcopy(self.net)

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()
                self.mem.add_instance(current_sample)

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

    def tta_attack_mode_turn_on(self):

        for module in self.net.modules():

            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html
                # TENT: force use of batch stats in train and eval modes: https://github.com/DequanWang/tent/blob/master/tent.py

                module.track_running_stats = False
                module.running_mean = None
                module.running_var = None

    def tta_attack_mode_turn_off(self):
        # reset model before actual inference
        self.net = deepcopy(self.original_model)

learner/src_for_tta_attack.py
taeckyung/SoTTA
import conf
from .dnn import DNN
from torch.utils.data import DataLoader

from utils.loss_functions import *
from utils import memory

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")


class TENT(DNN):
    def __init__(self, *args, **kwargs):
        super(TENT, self).__init__(*args, **kwargs)

        # turn on grad for BN params only

        for param in self.net.parameters():  # initially turn off requires_grad for all
            param.requires_grad = False

        for module in self.net.modules():

            if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
                # https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html
                # TENT: force use of batch stats in train and eval modes: https://github.com/DequanWang/tent/blob/master/tent.py
                if conf.args.use_learned_stats:
                    module.track_running_stats = True
                    module.momentum = conf.args.bn_momentum
                else:
                    module.track_running_stats = False
                    module.running_mean = None
                    module.running_var = None

                module.weight.requires_grad_(True)
                module.bias.requires_grad_(True)

        self.fifo = memory.FIFO(capacity=conf.args.update_every_x)  # required for evaluation
        self.mem_state = self.mem.save_state_dict()
        self.grad_norm_list = []

    def train_online(self, current_num_sample, add_memory=True, evaluation=True):
        """
        Train the model
        """

        TRAINED = 0
        SKIPPED = 1
        FINISHED = 2

        if not hasattr(self, 'previous_train_loss'):
            self.previous_train_loss = 0

        if current_num_sample > len(self.target_train_set[0]):
            return FINISHED

        # Add a sample
        feats, cls, dls = self.target_train_set
        current_sample = feats[current_num_sample - 1], cls[current_num_sample - 1], dls[current_num_sample - 1]

        if add_memory:
            self.fifo.add_instance(current_sample)  # for batch-based inference

            with torch.no_grad():
                self.net.eval()

                if conf.args.memory_type in ['FIFO', 'Reservoir']:
                    self.mem.add_instance(current_sample)

                elif conf.args.memory_type in ['HUS']:
                    f, c, d = current_sample[0].to(device), current_sample[1].to(device), current_sample[2].to(device)

                    logit = self.net(f.unsqueeze(0))
                    psuedo_conf = logit.max(1, keepdim=False)[0][0].cpu()
                    pseudo_cls = logit.max(1, keepdim=False)[1][0]
                    self.mem.add_instance([f, pseudo_cls, d, psuedo_conf])
                else:
                    raise NotImplementedError

        # self.num_in_mem +=1

        if conf.args.use_learned_stats and evaluation:  # batch-free inference
            self.evaluation_online(current_num_sample, [[current_sample[0]], [current_sample[1]], [current_sample[2]]])

        if current_num_sample % conf.args.update_every_x != 0:  # train only when enough samples are collected
            if not (current_num_sample == len(self.target_train_set[0]) and
                    conf.args.update_every_x >= current_num_sample):  # update with entire data

                self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=self.previous_train_loss)
                return SKIPPED

        # Evaluate with a batch
        if not conf.args.use_learned_stats and evaluation:  # batch-based inference
            self.evaluation_online(current_num_sample, self.fifo.get_memory())

        # setup models
        self.net.train()

        if len(feats) == 1:  # avoid BN error
            self.net.eval()

        feats, cls, dls = self.mem.get_memory()
        feats, cls, dls = torch.stack(feats), cls, torch.stack(dls)

        dataset = torch.utils.data.TensorDataset(feats)
        data_loader = DataLoader(dataset, batch_size=conf.args.opt['batch_size'],
                                 shuffle=True,
                                 drop_last=False, pin_memory=False)

        entropy_loss = HLoss()

        for e in range(conf.args.epoch):

            for batch_idx, (feats,) in enumerate(data_loader):
                feats = feats.to(device)

                if conf.args.tta_attack_type:
                    feats = feats.clone().detach()

                preds_of_data = self.net(feats)

                loss = entropy_loss(preds_of_data)

                self.optimizer.zero_grad()

                if conf.args.tta_attack_type:
                    loss.backward(retain_graph=True)
                else:
                    loss.backward()

                self.optimizer.step()

        if add_memory and evaluation:
            self.log_loss_results('train_online', epoch=current_num_sample, loss_avg=0)

        return TRAINED

learner/tent.py
taeckyung/SoTTA
# -*- coding: utf-8 -*-

import sys
import argparse
import random

import math
import numpy as np
import torch
import time
import os
import conf
from copy import deepcopy

os.environ['TOKENIZERS_PARALLELISM'] = 'false'

TRAINED = 0
SKIPPED = 1
FINISHED = 2


def get_path():
    path = 'log/'

    # information about used data type
    path += conf.args.dataset + '/'

    # information about used model type

    if conf.args.noisy_type:
        path += conf.args.method + "_noisy/"
    else:
        path += conf.args.method + '/'

    # information about domain(condition) of training data
    if conf.args.src == ['rest']:
        path += 'src_rest' + '/'
    elif conf.args.src == ['all']:
        path += 'src_all' + '/'
    elif conf.args.src is not None and len(conf.args.src) >= 1:
        path += 'src_' + '_'.join(conf.args.src) + '/'

    if conf.args.tgt:
        if conf.args.noisy_type:
            path += 'tgt_' + conf.args.tgt + '_{}'.format(conf.args.noisy_type) + \
                    '_{}'.format(conf.args.noisy_size) + '_{}/'.format(conf.args.noisy_class)
        else:
            path += 'tgt_' + conf.args.tgt + '/'

    if conf.args.noisy_type:
        if "FGSM" in conf.args.noisy_type:
            path += "{}/".format(conf.args.tta_attack_eps)

    if conf.args.tta_attack_type:
        path += "tta_attack_{}_num{}_step{}_eps{}/".format(conf.args.tta_attack_type, conf.args.tta_attack_num_samples,
                                                           conf.args.tta_attack_step, conf.args.tta_attack_eps)

    path += conf.args.log_prefix + '/'

    checkpoint_path = path + 'cp/'
    log_path = path
    result_path = path + '/'

    print('Path:{}'.format(path))
    return result_path, checkpoint_path, log_path


def attack_online(learner, current_num_sample):
    if conf.args.method == "SRC_FOR_TTA_ATTACK":  # for ablation
        ret_val = learner.train_online(current_num_sample, add_memory=True, evaluation=False)  # add memory FIFO
        if ret_val == TRAINED:  # if current_num_sample % update_every_x == 0
            learner.reset_purturb()  # reset pruturb to zero
            for step in range(conf.args.tta_attack_step):
                learner.tta_attack_train(current_num_sample, step)  # adding generated attack samples

            learner.tta_attack_update(current_num_sample)
            ret_val = learner.train_online(current_num_sample, add_memory=False, evaluation=False)  # train online

    else:  # normal
        if current_num_sample % conf.args.update_every_x == 0:
            temp_model_state = deepcopy(learner.net.state_dict())  # save model state before train

        ret_val = learner.train_online(current_num_sample, add_memory=True, evaluation=False)  # add memory FIFO

        if ret_val == TRAINED:
            assert current_num_sample % conf.args.update_every_x == 0
            learner.reset_purturb()  # reset pruturb to zero
            learner.net.load_state_dict(temp_model_state)  # reset model
            for step in range(conf.args.tta_attack_step):
                learner.tta_attack_train(current_num_sample, step)  # adding generated attack samples
                learner.net.load_state_dict(temp_model_state)  # reset model for next testing step

            learner.tta_attack_update(current_num_sample)
            ret_val = learner.train_online(current_num_sample, add_memory=False, evaluation=False)  # train online

    return ret_val


def main():
    ######################################################################
    device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")
    # Assume that we are on a CUDA machine, then this should print a CUDA device:
    print(device)

    ################### Hyperparameters #################
    if 'cifar100' in conf.args.dataset:
        opt = conf.CIFAR100Opt
    elif 'cifar10' in conf.args.dataset:
        opt = conf.CIFAR10Opt
    elif 'imagenet' in conf.args.dataset:
        opt = conf.IMAGENET_C
    else:
        raise NotImplementedError

    conf.args.opt = opt
    if conf.args.lr:
        opt['learning_rate'] = conf.args.lr
    if conf.args.weight_decay:
        opt['weight_decay'] = conf.args.weight_decay

    model = None

    if conf.args.model == "resnet18":
        from models import ResNet
        model = ResNet.ResNet18()
    elif conf.args.model == "resnet18_pretrained":
        import torchvision
        model = torchvision.models.resnet18(pretrained=True)
    elif conf.args.model == "resnet50":
        from models import ResNet
        model = ResNet.ResNet50()
    elif conf.args.model == "resnet50_pretrained":
        import torchvision
        model = torchvision.models.resnet50(pretrained=True)

    # import modules after setting the seed
    from data_loader import data_loader as data_loader
    from learner.dnn import DNN
    from learner.bn_stats import BN_Stats
    from learner.pseudo_label import PseudoLabel
    from learner.tent import TENT
    from learner.sotta import SoTTA
    from learner.cotta import CoTTA
    from learner.lame import LAME
    from learner.sar import SAR
    from learner.rotta import RoTTA
    from learner.eata import EATA
    from learner.src_for_tta_attack import SRC_FOR_TTA_ATTACK

    result_path, checkpoint_path, log_path = get_path()

    ########## Dataset loading ############################

    if conf.args.method == 'Src':
        learner_method = DNN
    elif conf.args.method == 'BN_Stats':
        learner_method = BN_Stats
    elif conf.args.method == 'PseudoLabel':
        learner_method = PseudoLabel
    elif conf.args.method == 'TENT':
        learner_method = TENT
    elif conf.args.method == 'CoTTA':
        learner_method = CoTTA
    elif conf.args.method == 'LAME':
        learner_method = LAME
    elif conf.args.method == 'SAR':
        learner_method = SAR
    elif conf.args.method == "RoTTA":
        learner_method = RoTTA
    elif conf.args.method == 'SoTTA':
        learner_method = SoTTA
    elif conf.args.method == "SRC_FOR_TTA_ATTACK":
        learner_method = SRC_FOR_TTA_ATTACK
    elif conf.args.method == "EATA":
        learner_method = EATA
    else:
        raise NotImplementedError

    since = time.time()

    print('##############Source Data Loading...##############')
    source_data_loader = data_loader.domain_data_loader(conf.args.dataset, conf.args.src,
                                                        conf.args.opt['file_path'],
                                                        batch_size=conf.args.opt['batch_size'],
                                                        valid_split=0,  # to be used for the validation
                                                        test_split=0, is_src=True)

    print('##############Target Data Loading...##############')
    target_data_loader = data_loader.domain_data_loader(conf.args.dataset, conf.args.tgt,
                                                        conf.args.opt['file_path'],
                                                        batch_size=conf.args.opt['batch_size'],
                                                        valid_split=0,
                                                        test_split=0, is_src=False)

    learner = learner_method(model, source_dataloader=source_data_loader,
                             target_dataloader=target_data_loader, write_path=log_path)

    time_elapsed = time.time() - since
    print('Data Loading Completion time: {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    #################### Training #########################

    since = time.time()

    # make dir if doesn't exist
    if not os.path.exists(result_path):
        oldumask = os.umask(0)
        os.makedirs(result_path, 0o777)
        os.umask(oldumask)
    if not os.path.exists(checkpoint_path):
        oldumask = os.umask(0)
        os.makedirs(checkpoint_path, 0o777)
        os.umask(oldumask)

    if not conf.args.online:

        start_epoch = 1
        best_acc = -9999
        best_epoch = -1

        for epoch in range(start_epoch, conf.args.epoch + 1):
            learner.train(epoch)

        learner.save_checkpoint(epoch=0, epoch_acc=-1, best_acc=best_acc,
                                checkpoint_path=checkpoint_path + 'cp_last.pth.tar')
        learner.dump_eval_online_result(is_train_offline=True)  # eval with final model

        time_elapsed = time.time() - since
        print('Completion time: {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print('Best val Acc: {:4f} at Epoch: {:d}'.format(best_acc, best_epoch))

    else:  # online
        current_num_sample = 1
        num_sample_end = conf.args.nsample
        best_acc = -9999
        best_epoch = -1

        finished = False

        while not finished and current_num_sample < num_sample_end:
            if conf.args.tta_attack_type:
                ret_val = attack_online(learner, current_num_sample)
            else:
                ret_val = learner.train_online(current_num_sample)

            if ret_val == FINISHED:
                break
            elif ret_val == SKIPPED:
                pass
            elif ret_val == TRAINED:
                pass
            current_num_sample += 1

        if not conf.args.remove_cp:
            learner.save_checkpoint(epoch=0, epoch_acc=-1, best_acc=best_acc,
                                    checkpoint_path=checkpoint_path + 'cp_last.pth.tar')
        learner.dump_eval_online_result()

        time_elapsed = time.time() - since
        print('Completion time: {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print('Best val Acc: {:4f} at Epoch: {:d}'.format(best_acc, best_epoch))

    if conf.args.remove_cp:
        last_path = checkpoint_path + 'cp_last.pth.tar'
        best_path = checkpoint_path + 'cp_best.pth.tar'
        try:
            os.remove(last_path)
            os.remove(best_path)
        except Exception:
            pass


def parse_arguments():
    """Command line parse."""

    # Note that 'type=bool' args should be False in default. Any string argument is recognized as "True". Do not give "--bool_arg 0"

    parser = argparse.ArgumentParser()

    # Required
    parser.add_argument('--dataset', type=str, default='',
                        help='Dataset to be used, in [cifar10noisy, cifar100noisy, imagenetnoisy]')
    parser.add_argument('--model', type=str, default='',
                        help='Base model to use')
    parser.add_argument('--method', type=str, default='',
                        help='TTA method')
    parser.add_argument('--src', nargs='*', default=None,
                        help='Specify source domains; not passing an arg load default src domains specified in conf.py')
    parser.add_argument('--tgt', type=str, default=None,
                        help='specific target domain; give "src" if you test under src domain')
    parser.add_argument('--gpu_idx', type=int, default=0, help='which gpu to use')

    # Noisy data settings
    parser.add_argument('--noisy_type', default=None, type=str,
                        help='Noisy test data : divide, repeat, oneclassrepeat, cifar100, cifar100c, gaussian, uniform, attack')
    parser.add_argument('--noisy_size', default=None, type=int,
                        help='Noisy test data size')
    parser.add_argument('--noisy_class', default=None, nargs="*", type=int,
                        help='Noisy test data target class')

    # Additional
    parser.add_argument('--dummy', action='store_true', default=False,
                        help='do nothing')
    parser.add_argument('--log_prefix', type=str, default='',
                        help='Prefix of log file path')
    parser.add_argument('--load_checkpoint_path', type=str, default='',
                        help='Load checkpoint and train from checkpoint in path?')
    parser.add_argument('--remove_cp', action='store_true',
                        help='Remove checkpoints after evaluation')

    # Learning parameters
    parser.add_argument('--lr', type=float, default=None,
                        help='learning rate to overwrite conf.py')
    parser.add_argument('--weight_decay', type=float, default=None,
                        help='weight_decay to overwrite conf.py')
    parser.add_argument('--seed', type=int, default=0,
                        help='random seed')
    parser.add_argument('--use_learned_stats', action='store_true',
                        help='Use learned stats uses exponential moving average with bn_momentum')
    parser.add_argument('--bn_momentum', type=float, default=0.1,
                        help='Momentum of exponential moving average')
    parser.add_argument('--epoch', type=int, default=1,
                        help='How many epochs do you want to use for train')
    parser.add_argument('--online', action='store_true',
                        help='training via online learning?')
    parser.add_argument('--update_every_x', type=int, default=1,
                        help='number of target samples used for every update')
    parser.add_argument('--memory_type', type=str, default='FIFO',
                        help='Memory type. FIFO as default')
    parser.add_argument('--memory_size', type=int, default=1,
                        help='number of previously trained data to be used for training')

    # Data
    parser.add_argument('--train_max_rows', type=int, default=np.inf,
                        help='How many data do you want to use for train')
    parser.add_argument('--valid_max_rows', type=int, default=np.inf,
                        help='How many data do you want to use for valid')
    parser.add_argument('--test_max_rows', type=int, default=np.inf,
                        help='How many data do you want to use for test')
    parser.add_argument('--nsample', type=int, default=20000000,
                        help='How many samples do you want use for train')
    parser.add_argument('--tgt_train_dist', type=int, default=1,
                        help='0: real selection'
                             '1: random selection'
                             '4: dirichlet distribution'
                             '5: no shuffling'
                             '6: online imbalanced label shifts (ImageNet-C only; details in SAR)'
                        )
    parser.add_argument('--dirichlet_beta', type=float, default=0.1,
                        help='the concentration parameter of the Dirichlet distribution for heterogeneous partition.')
    parser.add_argument('--validation', action='store_true',
                        help='Use validation data instead of test data for hyperparameter tuning')

    # LAME
    parser.add_argument('--parallel', type=bool, default=False)

    # CoTTA
    parser.add_argument('--ema_factor', type=float, default=0.999)
    parser.add_argument('--restoration_factor', type=float, default=0.01)
    parser.add_argument('--aug_threshold', type=float, default=0.92)

    # SoTTA
    parser.add_argument('--temperature', type=float, default=1.0,
                        help='temperature for HLoss')
    parser.add_argument('--loss_scaler', type=float, default=0,
                        help='loss_scaler for entropy_loss')
    parser.add_argument('--esm', action='store_true', default=False,
                        help='changes Adam to ESM+Adam')
    parser.add_argument('--high_threshold', default=0.99, type=float,
                        help='High confidence threshold')
    parser.add_argument('--log_grad', action='store_true', default=False,
                        help="log grad-norm")

    # EATA settings
    parser.add_argument('--fisher_size', default=2000, type=int,
                        help='number of samples to compute fisher information matrix.')
    parser.add_argument('--fisher_alpha', type=float, default=2000.,
                        help='the trade-off between entropy and regularization loss, in Eqn. (8)')
    parser.add_argument('--e_margin', type=float, default=math.log(1000) * 0.40,
                        help='entropy margin E_0 in Eqn. (3) for filtering reliable samples')
    parser.add_argument('--d_margin', type=float, default=0.05,
                        help='\\epsilon in Eqn. (5) for filtering redundant samples')

    # TTA Attack setting
    parser.add_argument('--tta_attack_eps', default=8 / 255, type=float,
                        help='eps of attack, ie 8/255 for FGSM')
    parser.add_argument('--tta_attack_type', default=None, type=str,
                        help='adversarial attack : targeted , indiscriminate, stealthy')
    parser.add_argument('--tta_attack_step', default=200, type=int,
                        help='number of iteration for training attack data')
    parser.add_argument('--tta_attack_num_samples', default=20, type=int,
                        help='number of attack samples')
    parser.add_argument('--tta_attack_target', default=None, type=int,
                        help='target attack class')

    return parser.parse_args()


def set_seed():
    torch.manual_seed(conf.args.seed)
    np.random.seed(conf.args.seed)
    random.seed(conf.args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


if __name__ == '__main__':
    print('Command:', end='\t')
    print(" ".join(sys.argv))
    conf.args = parse_arguments()
    print(conf.args)
    set_seed()
    main()

main.py
taeckyung/SoTTA
import json
import os
import re
import argparse

import pandas as pd
from multiprocessing import Pool
from functools import partial

CORRUPTION_LIST = ["gaussian_noise", "shot_noise", "impulse_noise", "defocus_blur", "glass_blur", "motion_blur",
                   "zoom_blur", "snow", "frost", "fog", "brightness", "contrast", "elastic_transform", "pixelate",
                   "jpeg_compression"]

METHOD_LIST = ["Src", "TENT", "PseudoLabel", "BN_Stats", "SAR", "RoTTA", "CoTTA", "LAME", "MEMO", "SoTTA", "EATA"]

BASE_DATASET = "cifar10noisy"

LOG_PREFIX = "eval_results"

SEED = [0, 1, 2]

DIST = 1

NOISE = ["original", "cifar100", "imagenet", "mnist", "uniform", "repeat"]


def get_avg_online_acc(file_path):
    f = open(file_path)
    json_data = json.load(f)
    f.close()
    return json_data['accuracy'][-1]


def process_path(args, path):
    result = {f"{s}_{t}": pd.DataFrame(columns=CORRUPTION_LIST) for s in args.seed for t in args.noisy_type}
    method = path.split("/")[-1].replace("_noisy", "")
    for (path, _, _) in os.walk(path):
        for corr in CORRUPTION_LIST:
            for seed in args.seed:
                for noise in args.noisy_type:
                    if method == 'Src':
                        pattern_of_path = f'.*{corr}.*{noise}.*/.*{args.prefix}_.*{seed}.*'
                    elif 'repeat' in noise:  # attack
                        if args.dataset in ['cifar10noisy', 'cifar100noisy']:
                            attack = "tta_attack_indiscriminate_num20_step10_eps0.1"
                        elif args.dataset == 'imagenetnoisy':
                            attack = "tta_attack_indiscriminate_num20_step1_eps0.2"
                        else:
                            raise NotImplementedError
                        pattern_of_path = f'.*{corr}.*{noise}.*/{attack}/.*{args.prefix}_.*{seed}_dist{args.dist}.*'
                    else:
                        pattern_of_path = f'.*{corr}.*{noise}.*/.*{args.prefix}_.*{seed}_dist{args.dist}.*'

                    pattern_of_path = re.compile(pattern_of_path)
                    if pattern_of_path.match(path):
                        if not path.endswith('/cp'):  # ignore cp/ dir
                            try:
                                acc = get_avg_online_acc(path + '/online_eval.json')
                                key = method + "_" + path.split("/")[-1]
                                result[f"{seed}_{noise}"].loc[key, corr] = float(acc)
                            except:
                                pass
    return result


def main(args):
    root = 'log/' + args.dataset
    paths = [os.path.join(root, f"{method}_noisy") for method in args.method]
    with Pool(processes=len(paths)) as p:
        func = partial(process_path, args)
        results = p.map(func, paths)

    for noisy_type in args.noisy_type:
        for seed in args.seed:
            print(f"SEED:{seed}, NOISY_TYPE: {noisy_type}")
            result = pd.concat([results[i][f"{seed}_{noisy_type}"] for i in range(len(results))])
            print(result.to_csv())


def parse_arguments():
    """Command line parse."""

    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset', type=str, default=BASE_DATASET,
                        help='Base dataset')

    parser.add_argument('--method', nargs="*", type=str, default=METHOD_LIST,
                        help='Method name')

    parser.add_argument('--seed', nargs="*", type=int, default=SEED,
                        help='Seed')

    parser.add_argument('--noisy_type', nargs="*", type=str, default=NOISE,
                        help='Noisy data type')

    parser.add_argument('--prefix', type=str, default=LOG_PREFIX,
                        help='Log prefix')

    parser.add_argument('--dist', type=str, default=DIST,
                        help='Distribution')

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_arguments()

    print(
        f"DATASET: {args.dataset}\n"
        f"LOG_PREFIX: {args.prefix}\n"
        f"METHOD: {args.method}\n"
        f"SEED: {args.seed}\n"
        f"Noise: {args.noisy_type}\n"
        f"Dist: {args.dist}"
    )

    main(args)

print_acc.py
taeckyung/SoTTA
import os
import sys
import numpy as np
import pickle

severities = [1, 2, 3, 4, 5]
corruptions = ["shot_noise", "motion_blur", "snow", "pixelate", "gaussian_noise", "defocus_blur", "brightness", "fog", \
               "zoom_blur", "frost", "glass_blur", "impulse_noise", "contrast", "jpeg_compression", "elastic_transform"]


def ensure_dir(file_path):
    directory = os.path.dirname(file_path)
    if not os.path.exists(directory):
        os.makedirs(directory)


def reshape(data):
    # original data shape: (N, 3072)
    # reshape the vector into height * weight * channel format
    nchw = data.reshape((data.shape[0], 3, 32, 32))
    nhwc = np.transpose(nchw, (0, 2, 3, 1))
    return nhwc


def get_dirnames(dataset):
    if dataset == "cifar-10c":
        cor_dirname = "CIFAR-10-C"
        org_dirname = "cifar-10-batches-py"
    if dataset == "cifar-100c":
        cor_dirname = "CIFAR-100-C"
        org_dirname = "cifar-100-python"
    return cor_dirname, org_dirname


def load_data(filename):
    with open(filename, 'rb') as f:
        d = pickle.load(f, encoding='bytes')
    return d


def process_corrupted_data(cor_dir):
    """
    process the downloaded data and save the results into the dedicated directories.
    for cifar-10-c or cifar-100-c, we parse data according to the severity label and save each in the corresponding directory.
    e.g., data of severity level 1 with corruption type 'fog' is saved as "./dataset/corrupted/severity-1/fog.npy"
    cor_dir: directory of corrupted dataset.
    """
    os.chdir(cor_dir)

    label_all = np.load("labels.npy")

    print("python: processing data...")
    for i, corruption in enumerate(corruptions):
        corruption_file_name = corruption + ".npy"
        data_all = np.load(corruption_file_name)
        for severity in severities:
            data = data_all[(severity - 1) * 10000: severity * 10000]
            label = label_all[(severity - 1) * 10000: severity * 10000]
            new_data_dir = "./corrupted/severity-" + str(severity) + "/" + corruption + ".npy"
            ensure_dir(new_data_dir)
            np.save(new_data_dir, data)

            if i == 0:
                new_label_dir = "./corrupted/severity-" + str(severity) + "/labels.npy"
                np.save(new_label_dir, label)


def process_original_data(org_dir, dataset):
    """
    process the downloaded data and save the results into the dedicated directories.
    for cifar 10,
        we read all five batches of train data, concatenate them into a single dataset and save it to the "origin" directory.
        we read a test batch and save it in each severity level directories in 'corrupted' directory.
        e.g., "./dataset/corrupted/severity-1/test.npy"
    for cifar 100,
        we read a train data and save it to the "origin" directory.
        we read a test data and save it in each severity level directories in 'corrupted' directory.
        e.g., "./dataset/corrupted/severity-1/test.npy"
    org_dir: directory of origin dataset.
    dataset: dataset type, "cifar-10c" or "cifar-100c".
    """
    os.chdir(org_dir)

    if dataset == "cifar-10c":

        # 1. load train data
        files = ["data_batch_1", "data_batch_2", "data_batch_3", "data_batch_4", "data_batch_5"]
        data_list = []
        label_list = []

        for filename in files:
            data_dict = load_data(filename)
            batch_data = data_dict[b'data']
            batch_label = np.asarray(data_dict[b'labels'])
            reshaped_data = reshape(batch_data)  # reshape data

            # append data and label to the predefined list
            data_list.append(reshaped_data)
            label_list.append(batch_label)

        # concat data and labels
        data = np.concatenate(np.asarray(data_list), axis=0)
        label = np.concatenate(np.asarray(label_list), axis=0)

        # 2. load test data
        test_file = "test_batch"
        test_data_dict = load_data(test_file)
        test_data = test_data_dict[b'data']
        test_data = reshape(test_data)  # reshape data

    else:  # dataset == "cifar-100c"
        # 1. load train data
        data_dict = load_data("train")
        data = data_dict[b'data']
        label = np.asarray(data_dict[b'fine_labels'])
        data = reshape(data)  # reshape data

        # 2. load test data
        test_data_dict = load_data("test")
        test_data = test_data_dict[b'data']
        test_data = reshape(test_data)  # reshape data

    # save train data
    org_data_dir = f'{cor_dir}/origin'
    if not os.path.exists(org_data_dir):
        os.makedirs(org_data_dir)

    np.save(f"{org_data_dir}/original.npy", data)
    np.save(f"{org_data_dir}/labels.npy", label)

    # save test data in each severity directories in 'corrupted' directory
    for severity in severities:
        save_data_dir = f'{cor_dir}/corrupted/severity-{severity}/test.npy'
        ensure_dir(save_data_dir)
        np.save(save_data_dir, test_data)

    # for severity-all directory, match the data shape
    save_data_dir = f'{cor_dir}/corrupted/severity-all/test.npy'
    ensure_dir(save_data_dir)
    test_data_ = np.concatenate([test_data for _ in range(len(severities))], axis=0)
    np.save(save_data_dir, test_data_)


if __name__ == "__main__":
    # get directory names to save data
    dataset = sys.argv[1]
    cor_dirname, org_dirname = get_dirnames(dataset)

    # get full directories
    home_dir = os.getcwd()
    cor_dir = f'{home_dir}/dataset/{cor_dirname}'
    org_dir = f'{home_dir}/dataset/{org_dirname}'

    process_corrupted_data(cor_dir)  # process corrupted data and save to the dedicated directory
    process_original_data(org_dir, dataset)  # process original data and save to the dedicated directory

process_cifar.py
taeckyung/SoTTA
import numpy as np
import torch

from PIL import ImageOps, Image
from torchvision import transforms


## https://github.com/google-research/augmix

def _augmix_aug(x_orig):
    x_orig = preaugment(x_orig)
    x_processed = preprocess(x_orig)
    w = np.float32(np.random.dirichlet([1.0, 1.0, 1.0]))
    m = np.float32(np.random.beta(1.0, 1.0))

    mix = torch.zeros_like(x_processed)
    for i in range(3):
        x_aug = x_orig.copy()
        for _ in range(np.random.randint(1, 4)):
            x_aug = np.random.choice(augmentations)(x_aug)
        mix += w[i] * preprocess(x_aug)
    mix = m * x_processed + (1 - m) * mix
    return mix


def autocontrast(pil_img, level=None):
    return ImageOps.autocontrast(pil_img)


def equalize(pil_img, level=None):
    return ImageOps.equalize(pil_img)


def rotate(pil_img, level):
    degrees = int_parameter(rand_lvl(level), 30)
    if np.random.uniform() > 0.5:
        degrees = -degrees
    return pil_img.rotate(degrees, resample=Image.BILINEAR, fillcolor=128)


def solarize(pil_img, level):
    level = int_parameter(rand_lvl(level), 256)
    return ImageOps.solarize(pil_img, 256 - level)


def shear_x(pil_img, level):
    level = float_parameter(rand_lvl(level), 0.3)
    if np.random.uniform() > 0.5:
        level = -level
    return pil_img.transform((32, 32), Image.AFFINE, (1, level, 0, 0, 1, 0), resample=Image.BILINEAR, fillcolor=128)


def shear_y(pil_img, level):
    level = float_parameter(rand_lvl(level), 0.3)
    if np.random.uniform() > 0.5:
        level = -level
    return pil_img.transform((32, 32), Image.AFFINE, (1, 0, 0, level, 1, 0), resample=Image.BILINEAR, fillcolor=128)


def translate_x(pil_img, level):
    level = int_parameter(rand_lvl(level), 32 / 3)
    if np.random.random() > 0.5:
        level = -level
    return pil_img.transform((32, 32), Image.AFFINE, (1, 0, level, 0, 1, 0), resample=Image.BILINEAR, fillcolor=128)


def translate_y(pil_img, level):
    level = int_parameter(rand_lvl(level), 32 / 3)
    if np.random.random() > 0.5:
        level = -level
    return pil_img.transform((32, 32), Image.AFFINE, (1, 0, 0, 0, 1, level), resample=Image.BILINEAR, fillcolor=128)


def posterize(pil_img, level):
    level = int_parameter(rand_lvl(level), 4)
    return ImageOps.posterize(pil_img, 4 - level)


def int_parameter(level, maxval):
    """Helper function to scale `val` between 0 and maxval .
    Args:
    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
    maxval: Maximum value that the operation can have. This will be scaled
      to level/PARAMETER_MAX.
    Returns:
    An int that results from scaling `maxval` according to `level`.
    """
    return int(level * maxval / 10)


def float_parameter(level, maxval):
    """Helper function to scale `val` between 0 and maxval .
    Args:
    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
    maxval: Maximum value that the operation can have. This will be scaled
      to level/PARAMETER_MAX.
    Returns:
    A float that results from scaling `maxval` according to `level`.
    """
    return float(level) * maxval / 10.


def rand_lvl(n):
    return np.random.uniform(low=0.1, high=n)


augmentations = [
    autocontrast,
    equalize,
    lambda x: rotate(x, 1),
    lambda x: solarize(x, 1),
    lambda x: shear_x(x, 1),
    lambda x: shear_y(x, 1),
    lambda x: translate_x(x, 1),
    lambda x: translate_y(x, 1),
    lambda x: posterize(x, 1),
]

mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])
preaugment = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
])

utils/augmentation.py
taeckyung/SoTTA
import torch
import torch.nn as nn
from copy import deepcopy


class MomentumBN(nn.Module):
    def __init__(self, bn_layer: nn.BatchNorm2d, momentum):
        super().__init__()
        self.num_features = bn_layer.num_features
        self.momentum = momentum
        if bn_layer.track_running_stats and bn_layer.running_var is not None and bn_layer.running_mean is not None:
            self.register_buffer("source_mean", deepcopy(bn_layer.running_mean))
            self.register_buffer("source_var", deepcopy(bn_layer.running_var))
            self.source_num = bn_layer.num_batches_tracked
        self.weight = deepcopy(bn_layer.weight)
        self.bias = deepcopy(bn_layer.bias)

        self.register_buffer("target_mean", torch.zeros_like(self.source_mean))
        self.register_buffer("target_var", torch.ones_like(self.source_var))
        self.eps = bn_layer.eps

        self.current_mu = None
        self.current_sigma = None

    def forward(self, x):
        raise NotImplementedError


class RobustBN1d(MomentumBN):
    def forward(self, x):
        if self.training:
            b_var, b_mean = torch.var_mean(x, dim=0, unbiased=False, keepdim=False)  # (C,)
            mean = (1 - self.momentum) * self.source_mean + self.momentum * b_mean
            var = (1 - self.momentum) * self.source_var + self.momentum * b_var
            self.source_mean, self.source_var = deepcopy(mean.detach()), deepcopy(var.detach())
            mean, var = mean.view(1, -1), var.view(1, -1)
        else:
            mean, var = self.source_mean.view(1, -1), self.source_var.view(1, -1)

        x = (x - mean) / torch.sqrt(var + self.eps)
        weight = self.weight.view(1, -1)
        bias = self.bias.view(1, -1)

        return x * weight + bias


class RobustBN2d(MomentumBN):
    def forward(self, x):
        if self.training:
            b_var, b_mean = torch.var_mean(x, dim=[0, 2, 3], unbiased=False, keepdim=False)  # (C,)
            mean = (1 - self.momentum) * self.source_mean + self.momentum * b_mean
            var = (1 - self.momentum) * self.source_var + self.momentum * b_var
            self.source_mean, self.source_var = deepcopy(mean.detach()), deepcopy(var.detach())
            mean, var = mean.view(1, -1, 1, 1), var.view(1, -1, 1, 1)
        else:
            mean, var = self.source_mean.view(1, -1, 1, 1), self.source_var.view(1, -1, 1, 1)

        x = (x - mean) / torch.sqrt(var + self.eps)
        weight = self.weight.view(1, -1, 1, 1)
        bias = self.bias.view(1, -1, 1, 1)

        return x * weight + bias

utils/bn_layers_rotta.py
taeckyung/SoTTA
from robustbench.model_zoo.enums import ThreatModel
from robustbench.utils import load_model


def build_model(cfg):
    if cfg.CORRUPTION.DATASET in ["cifar10", "cifar100"]:
        base_model = load_model(cfg.MODEL.ARCH, cfg.CKPT_DIR,
                                cfg.CORRUPTION.DATASET, ThreatModel.corruptions).cuda()
    else:
        raise NotImplementedError()

    return base_model

utils/buildmodel_rotta.py
taeckyung/SoTTA
# KATANA: Simple Post-Training Robustness Using Test Time Augmentations
# https://arxiv.org/pdf/2109.08191v1.pdf
import torch
import torchvision.transforms.functional as F
from torchvision.transforms import ColorJitter, Compose, Lambda
from numpy import random


class GaussianNoise(torch.nn.Module):
    def __init__(self, mean=0., std=1.):
        super().__init__()
        self.std = std
        self.mean = mean

    def forward(self, img):
        noise = torch.randn(img.size()) * self.std + self.mean
        noise = noise.to(img.device)
        return img + noise

    def __repr__(self):
        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)


class Clip(torch.nn.Module):
    def __init__(self, min_val=0., max_val=1.):
        super().__init__()
        self.min_val = min_val
        self.max_val = max_val

    def forward(self, img):
        return torch.clip(img, self.min_val, self.max_val)

    def __repr__(self):
        return self.__class__.__name__ + '(min_val={0}, max_val={1})'.format(self.min_val, self.max_val)


class ColorJitterPro(ColorJitter):
    """Randomly change the brightness, contrast, saturation, and gamma correction of an image."""

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, gamma=0):
        super().__init__(brightness, contrast, saturation, hue)
        self.gamma = self._check_input(gamma, 'gamma')

    @staticmethod
    @torch.jit.unused
    def get_params(brightness, contrast, saturation, hue, gamma):
        """Get a randomized transform to be applied on image.

        Arguments are same as that of __init__.

        Returns:
            Transform which randomly adjusts brightness, contrast and
            saturation in a random order.
        """
        transforms = []

        if brightness is not None:
            brightness_factor = random.uniform(brightness[0], brightness[1])
            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))

        if contrast is not None:
            contrast_factor = random.uniform(contrast[0], contrast[1])
            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))

        if saturation is not None:
            saturation_factor = random.uniform(saturation[0], saturation[1])
            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))

        if hue is not None:
            hue_factor = random.uniform(hue[0], hue[1])
            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))

        if gamma is not None:
            gamma_factor = random.uniform(gamma[0], gamma[1])
            transforms.append(Lambda(lambda img: F.adjust_gamma(img, gamma_factor)))

        random.shuffle(transforms)
        transform = Compose(transforms)

        return transform

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Input image.

        Returns:
            PIL Image or Tensor: Color jittered image.
        """
        fn_idx = torch.randperm(5)
        for fn_id in fn_idx:
            if fn_id == 0 and self.brightness is not None:
                brightness = self.brightness
                brightness_factor = torch.tensor(1.0).uniform_(brightness[0], brightness[1]).item()
                img = F.adjust_brightness(img, brightness_factor)

            if fn_id == 1 and self.contrast is not None:
                contrast = self.contrast
                contrast_factor = torch.tensor(1.0).uniform_(contrast[0], contrast[1]).item()
                img = F.adjust_contrast(img, contrast_factor)

            if fn_id == 2 and self.saturation is not None:
                saturation = self.saturation
                saturation_factor = torch.tensor(1.0).uniform_(saturation[0], saturation[1]).item()
                img = F.adjust_saturation(img, saturation_factor)

            if fn_id == 3 and self.hue is not None:
                hue = self.hue
                hue_factor = torch.tensor(1.0).uniform_(hue[0], hue[1]).item()
                img = F.adjust_hue(img, hue_factor)

            if fn_id == 4 and self.gamma is not None:
                gamma = self.gamma
                gamma_factor = torch.tensor(1.0).uniform_(gamma[0], gamma[1]).item()
                img = img.clamp(1e-8, 1.0)  # to fix Nan values in gradients, which happens when applying gamma
                # after contrast
                img = F.adjust_gamma(img, gamma_factor)

        return img

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        format_string += 'brightness={0}'.format(self.brightness)
        format_string += ', contrast={0}'.format(self.contrast)
        format_string += ', saturation={0}'.format(self.saturation)
        format_string += ', hue={0})'.format(self.hue)
        format_string += ', gamma={0})'.format(self.gamma)
        return format_string

utils/cotta_utils.py
taeckyung/SoTTA
import torch
import torchvision.transforms.functional as F
from torchvision.transforms import ColorJitter, Compose, Lambda
from numpy import random
import PIL
import torchvision.transforms as transforms


def get_tta_transforms(input_shape, gaussian_std: float = 0.005, soft=False):
    img_shape = (input_shape[0], input_shape[1], 3)
    n_pixels = img_shape[0]

    clip_min, clip_max = 0.0, 1.0

    p_hflip = 0.5

    tta_transforms = transforms.Compose([
        Clip(0.0, 1.0),
        ColorJitterPro(
            brightness=[0.8, 1.2] if soft else [0.6, 1.4],
            contrast=[0.85, 1.15] if soft else [0.7, 1.3],
            saturation=[0.75, 1.25] if soft else [0.5, 1.5],
            hue=[-0.03, 0.03] if soft else [-0.06, 0.06],
            gamma=[0.85, 1.15] if soft else [0.7, 1.3]
        ),
        transforms.Pad(padding=int(n_pixels / 2), padding_mode='edge'),
        transforms.RandomAffine(
            degrees=[-8, 8] if soft else [-15, 15],
            translate=(1 / 16, 1 / 16),
            scale=(0.95, 1.05) if soft else (0.9, 1.1),
            shear=None,
            resample=PIL.Image.BILINEAR,
            fillcolor=None
        ),
        transforms.GaussianBlur(kernel_size=5, sigma=[0.001, 0.25] if soft else [0.001, 0.5]),
        transforms.CenterCrop(size=n_pixels),
        transforms.RandomHorizontalFlip(p=p_hflip),
        GaussianNoise(0, gaussian_std),
        Clip(clip_min, clip_max)
    ])
    return tta_transforms


class GaussianNoise(torch.nn.Module):
    def __init__(self, mean=0., std=1.):
        super().__init__()
        self.std = std
        self.mean = mean

    def forward(self, img):
        noise = torch.randn(img.size()) * self.std + self.mean
        noise = noise.to(img.device)
        return img + noise

    def __repr__(self):
        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)


class Clip(torch.nn.Module):
    def __init__(self, min_val=0., max_val=1.):
        super().__init__()
        self.min_val = min_val
        self.max_val = max_val

    def forward(self, img):
        return torch.clip(img, self.min_val, self.max_val)

    def __repr__(self):
        return self.__class__.__name__ + '(min_val={0}, max_val={1})'.format(self.min_val, self.max_val)


class ColorJitterPro(ColorJitter):
    """Randomly change the brightness, contrast, saturation, and gamma correction of an image."""

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, gamma=0):
        super().__init__(brightness, contrast, saturation, hue)
        self.gamma = self._check_input(gamma, 'gamma')

    @staticmethod
    @torch.jit.unused
    def get_params(brightness, contrast, saturation, hue, gamma):
        """Get a randomized transform to be applied on image.

        Arguments are same as that of __init__.

        Returns:
            Transform which randomly adjusts brightness, contrast and
            saturation in a random order.
        """
        transforms = []

        if brightness is not None:
            brightness_factor = random.uniform(brightness[0], brightness[1])
            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))

        if contrast is not None:
            contrast_factor = random.uniform(contrast[0], contrast[1])
            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))

        if saturation is not None:
            saturation_factor = random.uniform(saturation[0], saturation[1])
            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))

        if hue is not None:
            hue_factor = random.uniform(hue[0], hue[1])
            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))

        if gamma is not None:
            gamma_factor = random.uniform(gamma[0], gamma[1])
            transforms.append(Lambda(lambda img: F.adjust_gamma(img, gamma_factor)))

        random.shuffle(transforms)
        transform = Compose(transforms)

        return transform

    def forward(self, img):
        """
        Args:
            img (PIL Image or Tensor): Input image.

        Returns:
            PIL Image or Tensor: Color jittered image.
        """
        fn_idx = torch.randperm(5)
        for fn_id in fn_idx:
            if fn_id == 0 and self.brightness is not None:
                brightness = self.brightness
                brightness_factor = torch.tensor(1.0).uniform_(brightness[0], brightness[1]).item()
                img = F.adjust_brightness(img, brightness_factor)

            if fn_id == 1 and self.contrast is not None:
                contrast = self.contrast
                contrast_factor = torch.tensor(1.0).uniform_(contrast[0], contrast[1]).item()
                img = F.adjust_contrast(img, contrast_factor)

            if fn_id == 2 and self.saturation is not None:
                saturation = self.saturation
                saturation_factor = torch.tensor(1.0).uniform_(saturation[0], saturation[1]).item()
                img = F.adjust_saturation(img, saturation_factor)

            if fn_id == 3 and self.hue is not None:
                hue = self.hue
                hue_factor = torch.tensor(1.0).uniform_(hue[0], hue[1]).item()
                img = F.adjust_hue(img, hue_factor)

            if fn_id == 4 and self.gamma is not None:
                gamma = self.gamma
                gamma_factor = torch.tensor(1.0).uniform_(gamma[0], gamma[1]).item()
                img = img.clamp(1e-8, 1.0)  # to fix Nan values in gradients, which happens when applying gamma
                # after contrast
                img = F.adjust_gamma(img, gamma_factor)

        return img

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        format_string += 'brightness={0}'.format(self.brightness)
        format_string += ', contrast={0}'.format(self.contrast)
        format_string += ', saturation={0}'.format(self.saturation)
        format_string += ', hue={0})'.format(self.hue)
        format_string += ', gamma={0})'.format(self.gamma)
        return format_string

utils/custom_transforms.py
taeckyung/SoTTA
import os
from typing import Tuple
# from PIL import Image
import numpy as np

import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets

# from data.hoi_dataset import BongardDataset
try:
    from torchvision.transforms import InterpolationMode

    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from PIL import Image, ImageOps, ImageEnhance

# ImageNet code should change this value
IMAGE_SIZE = 224


def int_parameter(level, maxval):
    """Helper function to scale `val` between 0 and maxval .

  Args:
    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
    maxval: Maximum value that the operation can have. This will be scaled to
      level/PARAMETER_MAX.

  Returns:
    An int that results from scaling `maxval` according to `level`.
  """
    return int(level * maxval / 10)


def float_parameter(level, maxval):
    """Helper function to scale `val` between 0 and maxval.

  Args:
    level: Level of the operation that will be between [0, `PARAMETER_MAX`].
    maxval: Maximum value that the operation can have. This will be scaled to
      level/PARAMETER_MAX.

  Returns:
    A float that results from scaling `maxval` according to `level`.
  """
    return float(level) * maxval / 10.


def sample_level(n):
    return np.random.uniform(low=0.1, high=n)


def autocontrast(pil_img, _):
    return ImageOps.autocontrast(pil_img)


def equalize(pil_img, _):
    return ImageOps.equalize(pil_img)


def posterize(pil_img, level):
    level = int_parameter(sample_level(level), 4)
    return ImageOps.posterize(pil_img, 4 - level)


def rotate(pil_img, level):
    degrees = int_parameter(sample_level(level), 30)
    if np.random.uniform() > 0.5:
        degrees = -degrees
    return pil_img.rotate(degrees, resample=Image.BILINEAR)


def solarize(pil_img, level):
    level = int_parameter(sample_level(level), 256)
    return ImageOps.solarize(pil_img, 256 - level)


def shear_x(pil_img, level):
    level = float_parameter(sample_level(level), 0.3)
    if np.random.uniform() > 0.5:
        level = -level
    return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
                             Image.AFFINE, (1, level, 0, 0, 1, 0),
                             resample=Image.BILINEAR)


def shear_y(pil_img, level):
    level = float_parameter(sample_level(level), 0.3)
    if np.random.uniform() > 0.5:
        level = -level
    return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
                             Image.AFFINE, (1, 0, 0, level, 1, 0),
                             resample=Image.BILINEAR)


def translate_x(pil_img, level):
    level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
    if np.random.random() > 0.5:
        level = -level
    return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
                             Image.AFFINE, (1, 0, level, 0, 1, 0),
                             resample=Image.BILINEAR)


def translate_y(pil_img, level):
    level = int_parameter(sample_level(level), IMAGE_SIZE / 3)
    if np.random.random() > 0.5:
        level = -level
    return pil_img.transform((IMAGE_SIZE, IMAGE_SIZE),
                             Image.AFFINE, (1, 0, 0, 0, 1, level),
                             resample=Image.BILINEAR)


# operation that overlaps with ImageNet-C's test set
def color(pil_img, level):
    level = float_parameter(sample_level(level), 1.8) + 0.1
    return ImageEnhance.Color(pil_img).enhance(level)


# operation that overlaps with ImageNet-C's test set
def contrast(pil_img, level):
    level = float_parameter(sample_level(level), 1.8) + 0.1
    return ImageEnhance.Contrast(pil_img).enhance(level)


# operation that overlaps with ImageNet-C's test set
def brightness(pil_img, level):
    level = float_parameter(sample_level(level), 1.8) + 0.1
    return ImageEnhance.Brightness(pil_img).enhance(level)


# operation that overlaps with ImageNet-C's test set
def sharpness(pil_img, level):
    level = float_parameter(sample_level(level), 1.8) + 0.1
    return ImageEnhance.Sharpness(pil_img).enhance(level)


augmentations = [
    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
    translate_x, translate_y
]

augmentations_all = [
    autocontrast, equalize, posterize, rotate, solarize, shear_x, shear_y,
    translate_x, translate_y, color, contrast, brightness, sharpness
]


# AugMix Transforms
def get_preaugment():
    return transforms.Compose([
        transforms.RandomResizedCrop(IMAGE_SIZE),
        transforms.RandomHorizontalFlip(),
    ])


def augmix(image, preprocess, aug_list, severity=1):
    preaugment = get_preaugment()
    x_orig = preaugment(image)
    x_processed = preprocess(x_orig)
    if len(aug_list) == 0:
        return x_processed
    w = np.float32(np.random.dirichlet([1.0, 1.0, 1.0]))
    m = np.float32(np.random.beta(1.0, 1.0))

    mix = torch.zeros_like(x_processed)
    for i in range(3):
        x_aug = x_orig.copy()
        for _ in range(np.random.randint(1, 4)):
            x_aug = np.random.choice(aug_list)(x_aug, severity)
        mix += w[i] * preprocess(x_aug)
    mix = m * x_processed + (1 - m) * mix
    return mix


class AugMixAugmenter(object):
    def __init__(self, base_transform, preprocess, n_views=2, augmix=False,
                 severity=1):
        self.base_transform = base_transform
        self.preprocess = preprocess
        self.n_views = n_views
        if augmix:
            self.aug_list = augmentations
        else:
            self.aug_list = []
        self.severity = severity

    def __call__(self, x):
        image = self.preprocess(self.base_transform(x))
        views = [augmix(x, self.preprocess, self.aug_list, self.severity) for _ in range(self.n_views)]
        return [image] + views

utils/dataaux.py
taeckyung/SoTTA
import numpy as np


# https://stackoverflow.com/questions/10097477/python-json-array-newlines
def to_json(o, level=0, indent=3, space=" ", newline="\n"):
    INDENT = indent
    SPACE = space
    NEWLINE = newline

    ret = ""
    if isinstance(o, dict):
        ret += "{" + NEWLINE
        comma = ""
        for k, v in o.items():
            ret += comma
            comma = ",\n"
            ret += SPACE * INDENT * (level + 1)
            ret += '"' + str(k) + '":' + SPACE
            ret += to_json(v, level + 1)

        ret += NEWLINE + SPACE * INDENT * level + "}"
    elif isinstance(o, str):
        ret += '"' + o + '"'
    elif isinstance(o, list):
        ret += "[" + ",".join([to_json(e, level + 1) for e in o]) + "]"
    # Tuples are interpreted as lists
    elif isinstance(o, tuple):
        ret += "[" + ",".join(to_json(e, level + 1) for e in o) + "]"
    elif isinstance(o, bool):
        ret += "true" if o else "false"
    elif isinstance(o, int):
        ret += str(o)
    elif isinstance(o, float):
        ret += '%.7g' % o
    elif isinstance(o, np.float32):
        ret += '%.7g' % o
    elif isinstance(o, np.ndarray) and np.issubdtype(o.dtype, np.integer):
        ret += "[" + ','.join(map(str, o.flatten().tolist())) + "]"
    elif isinstance(o, np.ndarray) and np.issubdtype(o.dtype, np.inexact):
        ret += "[" + ','.join(map(lambda x: '%.7g' % x, o.flatten().tolist())) + "]"
    elif o is None:
        ret += 'null'
    else:
        raise TypeError("Unknown type '%s' for json serialization" % str(type(o)))
    return ret

utils/logging.py
taeckyung/SoTTA
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np


# https://discuss.pytorch.org/t/calculating-the-entropy-loss/14510
# but there is a bug in the original code: it sums up the entropy over a batch. so I take mean instead of sum
class HLoss(nn.Module):
    def __init__(self, temp_factor=1.0):
        super(HLoss, self).__init__()
        self.temp_factor = temp_factor

    def forward(self, x):
        softmax = F.softmax(x / self.temp_factor, dim=1)
        entropy = -softmax * torch.log(softmax + 1e-6)
        b = entropy.mean()

        return b


@torch.jit.script
def calc_energy(x: torch.Tensor, temp_factor: float = 1.0) -> torch.Tensor:
    return temp_factor * torch.logsumexp(x / temp_factor, dim=1)


class EnergyLoss(nn.Module):
    def __init__(self, temp_factor=1.0):
        super(EnergyLoss, self).__init__()
        self.temp_factor = temp_factor

    def forward(self, x):
        e = calc_energy(x, self.temp_factor)
        # energy = 1.0 / torch.linalg.vector_norm(6.0 - energy, 2)
        e = 1.0 / e.mean()
        return e


class JSDLoss(nn.Module):
    def __init__(self):
        super(JSDLoss, self).__init__()
        self.kl = nn.KLDivLoss(reduction='batchmean', log_target=True)

    def forward(self, p: torch.tensor, q: torch.tensor):
        p = F.softmax(p, dim=1)
        q = F.softmax(q, dim=1)
        p, q = p.view(-1, p.size(-1)), q.view(-1, q.size(-1))
        m = (0.5 * (p + q)).log()
        return 0.5 * (self.kl(m, p.log()) + self.kl(m, q.log()))


def marginal_entropy(outputs):
    logits = outputs - outputs.logsumexp(dim=-1, keepdim=True)
    avg_logits = logits.logsumexp(dim=0) - np.log(logits.shape[0])
    min_real = torch.finfo(avg_logits.dtype).min
    avg_logits = torch.clamp(avg_logits, min=min_real)
    return -(avg_logits * torch.exp(avg_logits)).sum(dim=-1), avg_logits


@torch.jit.script
def softmax_entropy(x: torch.Tensor) -> torch.Tensor:
    """Entropy of softmax distribution from logits."""
    softmax = x.softmax(1)
    return -(softmax * torch.log(softmax + 1e-6)).sum(1)


@torch.jit.script
def softmax_entropy_rotta(x, x_ema):
    return -(x_ema.softmax(1) * x.log_softmax(1)).sum(1)

utils/loss_functions.py
taeckyung/SoTTA
import random

import numpy as np
import torch

import conf

device = torch.device("cuda:{:d}".format(conf.args.gpu_idx) if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(
    conf.args.gpu_idx)  # this prevents unnecessary gpu memory allocation to cuda:0 when using estimator


class FIFO:
    def __init__(self, capacity):
        self.data = [[], [], []]
        self.capacity = capacity
        pass

    def set_memory(self, state_dict):  # for tta_attack
        self.data = [ls[:] for ls in state_dict['data']]
        if 'capacity' in state_dict.keys():
            self.capacity = state_dict['capacity']

    def save_state_dict(self):
        dic = {}
        dic['data'] = [ls[:] for ls in self.data]
        dic['capacity'] = self.capacity
        return dic

    def get_memory(self):
        return self.data

    def get_occupancy(self):
        return len(self.data[0])

    def add_instance(self, instance):
        assert (len(instance) == 3)

        if self.get_occupancy() >= self.capacity:
            self.remove_instance()

        for i, dim in enumerate(self.data):
            dim.append(instance[i])

    def remove_instance(self):
        for dim in self.data:
            dim.pop(0)
        pass


class HUS:
    def __init__(self, capacity, threshold=None):
        self.data = [[[], [], [], []] for _ in
                     range(conf.args.opt['num_class'])]  # feat, pseudo_cls, domain, conf
        self.counter = [0] * conf.args.opt['num_class']
        self.marker = [''] * conf.args.opt['num_class']
        self.capacity = capacity
        self.threshold = threshold

    def set_memory(self, state_dict):  # for tta_attack
        self.data = [[l[:] for l in ls] for ls in state_dict['data']]
        self.counter = state_dict['counter'][:]
        self.marker = state_dict['marker'][:]
        self.capacity = state_dict['capacity']
        self.threshold = state_dict['threshold']

    def save_state_dict(self):
        dic = {}
        dic['data'] = [[l[:] for l in ls] for ls in self.data]
        dic['counter'] = self.counter[:]
        dic['marker'] = self.marker[:]
        dic['capacity'] = self.capacity
        dic['threshold'] = self.threshold

        return dic

    def print_class_dist(self):
        print(self.get_occupancy_per_class())

    def print_real_class_dist(self):
        occupancy_per_class = [0] * conf.args.opt['num_class']
        for i, data_per_cls in enumerate(self.data):
            for cls in data_per_cls[3]:
                occupancy_per_class[cls] += 1
        print(occupancy_per_class)

    def get_memory(self):
        data = self.data

        tmp_data = [[], [], []]
        for data_per_cls in data:
            feats, cls, dls, _ = data_per_cls
            tmp_data[0].extend(feats)
            tmp_data[1].extend(cls)
            tmp_data[2].extend(dls)

        return tmp_data

    def get_occupancy(self):
        occupancy = 0
        for data_per_cls in self.data:
            occupancy += len(data_per_cls[0])
        return occupancy

    def get_occupancy_per_class(self):
        occupancy_per_class = [0] * conf.args.opt['num_class']
        for i, data_per_cls in enumerate(self.data):
            occupancy_per_class[i] = len(data_per_cls[0])
        return occupancy_per_class

    def add_instance(self, instance):
        assert (len(instance) == 4)
        cls = instance[1]
        self.counter[cls] += 1
        is_add = True

        if self.threshold is not None and instance[3] < self.threshold:
            is_add = False
        elif self.get_occupancy() >= self.capacity:
            is_add = self.remove_instance(cls)

        if is_add:
            for i, dim in enumerate(self.data[cls]):
                dim.append(instance[i])

    def get_largest_indices(self):
        occupancy_per_class = self.get_occupancy_per_class()
        max_value = max(occupancy_per_class)
        largest_indices = []
        for i, oc in enumerate(occupancy_per_class):
            if oc == max_value:
                largest_indices.append(i)
        return largest_indices

    def get_average_confidence(self):
        conf_list = []
        for i, data_per_cls in enumerate(self.data):
            for confidence in data_per_cls[3]:
                conf_list.append(confidence)
        if len(conf_list) > 0:
            return np.average(conf_list)
        else:
            return 0

    def get_target_index(self, data):
        return random.randrange(0, len(data))

    def remove_instance(self, cls):
        largest_indices = self.get_largest_indices()
        if cls not in largest_indices:  # instance is stored in the place of another instance that belongs to the largest class
            largest = random.choice(largest_indices)  # select only one largest class
            tgt_idx = self.get_target_index(self.data[largest][3])
            for dim in self.data[largest]:
                dim.pop(tgt_idx)
        else:  # replaces a randomly selected stored instance of the same class
            tgt_idx = self.get_target_index(self.data[cls][3])
            for dim in self.data[cls]:
                dim.pop(tgt_idx)
        return True

    def reset_value(self, feats, cls, aux):
        self.data = [[[], [], [], []] for _ in range(conf.args.opt['num_class'])]  # feat, pseudo_cls, domain, conf

        for i in range(len(feats)):
            tgt_idx = cls[i]
            self.data[tgt_idx][0].append(feats[i])
            self.data[tgt_idx][1].append(cls[i])
            self.data[tgt_idx][2].append(0)
            self.data[tgt_idx][3].append(aux[i])


class ConfFIFO:
    def __init__(self, capacity, threshold):
        self.data = [[], [], [], []]
        self.capacity = capacity
        self.threshold = threshold
        pass

    def set_memory(self, state_dict):  # for tta_attack
        self.data = [ls[:] for ls in state_dict['data']]
        self.threshold = state_dict['threshold']
        if 'capacity' in state_dict.keys():
            self.capacity = state_dict['capacity']

    def save_state_dict(self):
        dic = {}
        dic['data'] = [ls[:] for ls in self.data]
        dic['capacity'] = self.capacity
        dic['threshold'] = self.threshold
        return dic

    def get_memory(self):
        return self.data[:3]

    def get_occupancy(self):
        return len(self.data[0])

    def add_instance(self, instance):
        assert (len(instance) == 4)

        if instance[3] < self.threshold:
            return

        if self.get_occupancy() >= self.capacity:
            self.remove_instance()

        for i, dim in enumerate(self.data):
            dim.append(instance[i])

    def remove_instance(self):
        for dim in self.data:
            dim.pop(0)
        pass

    def reset_value(self, feats, cls, aux):
        self.data = [[], [], [], []]

utils/memory.py
taeckyung/SoTTA
import math


class MemoryItem:
    def __init__(self, data=None, uncertainty=0, age=0):
        self.data = data
        self.uncertainty = uncertainty
        self.age = age

    def increase_age(self):
        if not self.empty():
            self.age += 1

    def get_data(self):
        return self.data, self.uncertainty, self.age

    def empty(self):
        return self.data == "empty"


class CSTU:
    def __init__(self, capacity, num_class, lambda_t=1.0, lambda_u=1.0):
        self.capacity = capacity
        self.num_class = num_class
        self.per_class = self.capacity / self.num_class
        self.lambda_t = lambda_t
        self.lambda_u = lambda_u

        self.data: list[list[MemoryItem]] = [[] for _ in range(self.num_class)]

    def set_memory(self, state_dict):  # for tta_attack
        self.capacity = state_dict['capacity']
        self.num_class = state_dict['num_class']
        self.per_class = state_dict['per_class']
        self.lambda_t = state_dict['lambda_t']
        self.lambda_u = state_dict['lambda_u']
        self.data = [ls[:] for ls in state_dict['data']]

    def save_state_dict(self):
        dic = {}
        dic['capacity'] = self.capacity
        dic['num_class'] = self.num_class
        dic['per_class'] = self.per_class
        dic['lambda_t'] = self.lambda_t
        dic['lambda_u'] = self.lambda_u
        dic['data'] = [ls[:] for ls in self.data]

        return dic

    def get_occupancy(self):
        occupancy = 0
        for data_per_cls in self.data:
            occupancy += len(data_per_cls)
        return occupancy

    def per_class_dist(self):
        per_class_occupied = [0] * self.num_class
        for cls, class_list in enumerate(self.data):
            per_class_occupied[cls] = len(class_list)

        return per_class_occupied

    def add_instance(self, instance):
        assert (len(instance) == 3)
        x, prediction, uncertainty = instance
        new_item = MemoryItem(data=x, uncertainty=uncertainty, age=0)
        new_score = self.heuristic_score(0, uncertainty)
        if self.remove_instance(prediction, new_score):
            self.data[prediction].append(new_item)
        self.add_age()

    def remove_instance(self, cls, score):
        class_list = self.data[cls]
        class_occupied = len(class_list)
        all_occupancy = self.get_occupancy()
        if all_occupancy < self.capacity:
            return True
        if class_occupied < self.per_class:
            majority_classes = self.get_majority_classes()
            return self.remove_from_classes(majority_classes, score)
        else:
            return self.remove_from_classes([cls], score)

    def remove_from_classes(self, classes: list[int], score_base):
        max_class = None
        max_index = None
        max_score = None
        for cls in classes:
            for idx, item in enumerate(self.data[cls]):
                uncertainty = item.uncertainty
                age = item.age
                score = self.heuristic_score(age=age, uncertainty=uncertainty)
                if max_score is None or score >= max_score:
                    max_score = score
                    max_index = idx
                    max_class = cls

        if max_class is not None:
            if max_score > score_base:
                self.data[max_class].pop(max_index)
                return True
            else:
                return False
        else:
            return True

    def get_majority_classes(self):
        per_class_dist = self.per_class_dist()
        max_occupied = max(per_class_dist)
        classes = []
        for i, occupied in enumerate(per_class_dist):
            if occupied == max_occupied:
                classes.append(i)

        return classes

    def heuristic_score(self, age, uncertainty):
        return self.lambda_t * 1 / (1 + math.exp(-age / self.capacity)) + self.lambda_u * uncertainty / math.log(
            self.num_class)

    def add_age(self):
        for class_list in self.data:
            for item in class_list:
                item.increase_age()
        return

    def get_memory(self):
        tmp_data = []
        tmp_age = []

        for class_list in self.data:
            for item in class_list:
                tmp_data.append(item.data)
                tmp_age.append(item.age)

        tmp_age = [x / self.capacity for x in tmp_age]

        return tmp_data, tmp_age

utils/memory_rotta.py
taeckyung/SoTTA
import torch

_CIFAR10_MEAN = [0.4914, 0.4822, 0.4465]
_CIFAR10_STDDEV = [0.2471, 0.2435, 0.2616]

_CIFAR100_MEAN = [0.5071, 0.4865, 0.4409]
_CIFAR100_STDDEV = [0.2673, 0.2564, 0.2762]

_IMAGENET_MEAN = [0.485, 0.456, 0.406]
_IMAGENET_STDDEV = [0.229, 0.224, 0.225]


class NormalizeLayer(torch.nn.Module):
    """Standardize the channels of a batch of images by subtracting the dataset mean
      and dividing by the dataset standard deviation.
      """

    def __init__(self, means, sds):
        """
        :param means: the channel means
        :param sds: the channel standard deviations
        """
        super(NormalizeLayer, self).__init__()
        self.register_buffer(
            'mu', torch.tensor(means).view(-1, 1, 1))
        self.register_buffer(
            'sigma', torch.tensor(sds).view(-1, 1, 1))

    def forward(self, input: torch.tensor):
        return (input - self.mu) / self.sigma


def get_normalize_layer(dataset):
    """Return the dataset's normalization layer"""
    if dataset in ["cifar10", "cifar10noisy"]:
        return NormalizeLayer(_CIFAR10_MEAN, _CIFAR10_STDDEV)
    elif dataset in ["cifar100", "cifar100noisy"]:
        return NormalizeLayer(_CIFAR100_MEAN, _CIFAR100_STDDEV)
    elif dataset in ['imagenet', 'imagenetnoisy']:
        return NormalizeLayer(_IMAGENET_MEAN, _IMAGENET_STDDEV)
    else:
        return None

utils/normalize_layer.py
taeckyung/SoTTA
"""
from https://github.com/davda54/sam
"""

import torch
from torch import nn


def sam_collect_params(model, freeze_top=False):
    """Collect the affine scale + shift parameters from norm layers.
    Walk the model's modules and collect all normalization parameters.
    Return the parameters and their names.
    Note: other choices of parameterization are possible!
    """
    params = []
    names = []
    for nm, m in model.named_modules():
        # skip top layers for adaptation: layer4 for ResNets and blocks9-11 for Vit-Base
        if freeze_top:
            if 'layer4' in nm:
                continue
            if 'blocks.9' in nm:
                continue
            if 'blocks.10' in nm:
                continue
            if 'blocks.11' in nm:
                continue
            if 'norm.' in nm:
                continue
            if nm in ['norm']:
                continue

        if isinstance(m, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
            for np, p in m.named_parameters():
                if np in ['weight', 'bias']:  # weight is scale, bias is shift
                    params.append(p)
                    names.append(f"{nm}.{np}")

    return params, names


class SAM(torch.optim.Optimizer):
    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):
        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"

        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
        super(SAM, self).__init__(params, defaults)

        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
        # print(self.base_optimizer, self.param_groups, **kwargs)
        self.param_groups = self.base_optimizer.param_groups
        self.defaults.update(self.base_optimizer.defaults)

    @torch.no_grad()
    def first_step(self, zero_grad=False):
        grad_norm = self._grad_norm()
        for group in self.param_groups:
            scale = group["rho"] / (grad_norm + 1e-12)

            for p in group["params"]:
                if p.grad is None: continue
                self.state[p]["old_p"] = p.data.clone()
                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
                p.add_(e_w)  # climb to the local maximum "w + e(w)"

        if zero_grad: self.zero_grad()

    @torch.no_grad()
    def second_step(self, zero_grad=False):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None: continue
                p.data = self.state[p]["old_p"]  # get back to "w" from "w + e(w)"

        self.base_optimizer.step()  # do the actual "sharpness-aware" update

        if zero_grad: self.zero_grad()

    @torch.no_grad()
    def step(self, closure=None):
        assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass

        self.first_step(zero_grad=True)
        closure()
        self.second_step()

    def _grad_norm(self):
        shared_device = self.param_groups[0]["params"][
            0].device  # put everything on the same device, in case of model parallelism
        norm = torch.norm(
            torch.stack([
                ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for group in self.param_groups for p in group["params"]
                if p.grad is not None
            ]),
            p=2
        )
        return norm

    def load_state_dict(self, state_dict):
        super().load_state_dict(state_dict)
        self.base_optimizer.param_groups = self.param_groups

utils/sam_optimizer.py
taeckyung/SoTTA