Scaling for Training Time and Post-hoc Out-of-distribution Detection Enhancement.
import os
'''
path="./data/images_classic/cinic/valid"
save_path="./data/benchmark_imglist/cifar10/val_cinic10.txt"
prefix="cinic/valid/"
category=["airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"]
with open(save_path,'a') as f:
    for name in category:
        label=category.index(name)
        sub_path=path+'/'+name
        files=os.listdir(sub_path)
        for file in files:
            line=prefix+name+'/'+file+' '+str(label)+'\n'
            f.write(line)
    f.close()       
'''

path="./data/images_classic/cifar100c"
save_path="./data/benchmark_imglist/cifar100/test_cifar100c.txt"
prefix="cifar100c/"
files=os.listdir(path)
with open(save_path,'a') as f:
    for file in files:
        splits=file.split("_")
        label=(splits[1].split("."))[0]
        line=prefix+file+" "+label+'\n'
        f.write(line)
    f.close()  

'''
path="./data/images_largescale/imagenet_v2"
save_path="./data/benchmark_imglist/imagenet/test_imagenetv2.txt"
prefix="imagenet_v2/"
with open(save_path,'a') as f:
    for i in range(0,1000):
        label=str(i)
        sub_path=path+'/'+label
        files=os.listdir(sub_path)
        for file in files:
            line=prefix+label+'/'+file+' '+label+'\n'
            f.write(line)
    f.close() 
'''
imglist_generator.py
kai422/SCALE
from openood.pipelines import get_pipeline
from openood.utils import launch, setup_config


def main(config):
    """Main entrance. Config is all you need to provide to run the code. Config
    should be provided in the format of YAML and can be modified with command
    line.

    Example:
        python main.py \
            --config configs/datasets/mnist_datasets.yml \
            configs/train/mnist_baseline.yml \
            --dataset.image_size 32 \
            --network res18

    Note:
        A config file is the minimum requirement.
        You don't need to add "--config_key new_value"
        if you don't have anything to modify.
    """

    pipeline = get_pipeline(config)
    pipeline.run()


if __name__ == '__main__':

    config = setup_config()
    # generate output directory and save the full config file
    # setup_logger(config)

    # pipeline = get_pipeline(config)
    # pipeline.run()

    launch(
        main,
        config.num_gpus,
        num_machines=config.num_machines,
        machine_rank=config.machine_rank,
        dist_url='auto',
        args=(config, ),
    )

main.py
kai422/SCALE
from .utils import get_dataloader, get_feature_dataloader, get_feature_opengan_dataloader, get_ood_dataloader

openood/datasets/__init__.py
kai422/SCALE
import logging
import random
import traceback

from torch.utils.data import Dataset


class BaseDataset(Dataset):
    def __init__(self, pseudo_index=-1, skip_broken=False, new_index='next'):
        super(BaseDataset, self).__init__()
        self.pseudo_index = pseudo_index
        self.skip_broken = skip_broken
        self.new_index = new_index
        if new_index not in ('next', 'rand'):
            raise ValueError('new_index not one of ("next", "rand")')

    def __getitem__(self, index):
        # in some pytorch versions, input index will be torch.Tensor
        index = int(index)

        # if sampler produce pseudo_index,
        # randomly sample an index, and mark it as pseudo
        if index == self.pseudo_index:
            index = random.randrange(len(self))
            pseudo = 1
        else:
            pseudo = 0

        while True:
            try:
                sample = self.getitem(index)
                break
            except Exception as e:
                if self.skip_broken and not isinstance(e, NotImplementedError):
                    if self.new_index == 'next':
                        new_index = (index + 1) % len(self)
                    else:
                        new_index = random.randrange(len(self))
                    logging.warn(
                        'skip broken index [{}], use next index [{}]'.format(
                            index, new_index))
                    index = new_index
                else:
                    logging.error('index [{}] broken'.format(index))
                    traceback.print_exc()
                    logging.error(e)
                    raise e

        sample['index'] = index
        sample['pseudo'] = pseudo
        return sample

    def getitem(self, index):
        raise NotImplementedError

openood/datasets/base_dataset.py
kai422/SCALE
from torch.utils.data import Dataset


class FeatDataset(Dataset):
    def __init__(self, feat, labels):
        self.data = feat
        self.labels = labels
        self.len = feat.shape[0]
        assert self.len == len(labels)

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        data = self.data[idx]
        label = self.labels[idx]
        return {'data': data, 'label': label}

openood/datasets/feature_dataset.py
kai422/SCALE
import ast
import io
import logging
import os

import torch
from PIL import Image, ImageFile

from .base_dataset import BaseDataset

# to fix "OSError: image file is truncated"
ImageFile.LOAD_TRUNCATED_IMAGES = True


class Convert:
    def __init__(self, mode='RGB'):
        self.mode = mode

    def __call__(self, image):
        return image.convert(self.mode)


class ImglistAugMixDataset(BaseDataset):
    def __init__(self,
                 name,
                 imglist_pth,
                 data_dir,
                 num_classes,
                 preprocessor,
                 data_aux_preprocessor,
                 maxlen=None,
                 dummy_read=False,
                 dummy_size=None,
                 **kwargs):
        super(ImglistAugMixDataset, self).__init__(**kwargs)

        self.name = name
        with open(imglist_pth) as imgfile:
            self.imglist = imgfile.readlines()
        self.data_dir = data_dir
        self.num_classes = num_classes
        self.preprocessor = preprocessor
        self.transform_image = preprocessor
        self.transform_aux_image = data_aux_preprocessor
        self.maxlen = maxlen
        self.dummy_read = dummy_read
        self.dummy_size = dummy_size
        if dummy_read and dummy_size is None:
            raise ValueError(
                'if dummy_read is True, should provide dummy_size')

    def __len__(self):
        if self.maxlen is None:
            return len(self.imglist)
        else:
            return min(len(self.imglist), self.maxlen)

    def getitem(self, index):
        line = self.imglist[index].strip('\n')
        tokens = line.split(' ', 1)
        image_name, extra_str = tokens[0], tokens[1]
        if self.data_dir != '' and image_name.startswith('/'):
            raise RuntimeError('image_name starts with "/"')
        path = os.path.join(self.data_dir, image_name)
        sample = dict()
        sample['image_name'] = image_name
        kwargs = {'name': self.name, 'path': path, 'tokens': tokens}
        try:
            # some preprocessor methods require setup
            self.preprocessor.setup(**kwargs)
        except:
            pass

        try:
            if not self.dummy_read:
                with open(path, 'rb') as f:
                    content = f.read()
                filebytes = content
                buff = io.BytesIO(filebytes)
            if self.dummy_size is not None:
                sample['data'] = torch.rand(self.dummy_size)
            else:
                image = Image.open(buff).convert('RGB')
                orig, aug1, aug2 = self.transform_image(image)
                sample['data'] = orig
                sample['data_aug1'] = aug1
                sample['data_aug2'] = aug2
                sample['data_aux'] = self.transform_aux_image(image)
            extras = ast.literal_eval(extra_str)
            try:
                for key, value in extras.items():
                    sample[key] = value
                # if you use dic the code below will need ['label']
                sample['label'] = 0
            except AttributeError:
                sample['label'] = int(extra_str)
            # Generate Soft Label
            soft_label = torch.Tensor(self.num_classes)
            if sample['label'] < 0:
                soft_label.fill_(1.0 / self.num_classes)
            else:
                soft_label.fill_(0)
                soft_label[sample['label']] = 1
            sample['soft_label'] = soft_label

        except Exception as e:
            logging.error('[{}] broken'.format(path))
            raise e
        return sample

openood/datasets/imglist_augmix_dataset.py
kai422/SCALE
import ast
import io
import logging
import os

import torch
from PIL import Image, ImageFile

from .base_dataset import BaseDataset

# to fix "OSError: image file is truncated"
ImageFile.LOAD_TRUNCATED_IMAGES = True


class Convert:
    def __init__(self, mode='RGB'):
        self.mode = mode

    def __call__(self, image):
        return image.convert(self.mode)


class ImglistDataset(BaseDataset):
    def __init__(self,
                 name,
                 imglist_pth,
                 data_dir,
                 num_classes,
                 preprocessor,
                 data_aux_preprocessor,
                 maxlen=None,
                 dummy_read=False,
                 dummy_size=None,
                 **kwargs):
        super(ImglistDataset, self).__init__(**kwargs)

        self.name = name
        with open(imglist_pth) as imgfile:
            self.imglist = imgfile.readlines()
        self.data_dir = data_dir
        self.num_classes = num_classes
        self.preprocessor = preprocessor
        self.transform_image = preprocessor
        self.transform_aux_image = data_aux_preprocessor
        self.maxlen = maxlen
        self.dummy_read = dummy_read
        self.dummy_size = dummy_size
        if dummy_read and dummy_size is None:
            raise ValueError(
                'if dummy_read is True, should provide dummy_size')

    def __len__(self):
        if self.maxlen is None:
            return len(self.imglist)
        else:
            return min(len(self.imglist), self.maxlen)

    def getitem(self, index):
        line = self.imglist[index].strip('\n')
        tokens = line.split(' ', 1)
        image_name, extra_str = tokens[0], tokens[1]
        if self.data_dir != '' and image_name.startswith('/'):
            raise RuntimeError('image_name starts with "/"')
        path = os.path.join(self.data_dir, image_name)
        sample = dict()
        sample['image_name'] = image_name
        kwargs = {'name': self.name, 'path': path, 'tokens': tokens}
        try:
            # some preprocessor methods require setup
            self.preprocessor.setup(**kwargs)
        except:
            pass

        try:
            if not self.dummy_read:
                with open(path, 'rb') as f:
                    content = f.read()
                filebytes = content
                buff = io.BytesIO(filebytes)
            if self.dummy_size is not None:
                sample['data'] = torch.rand(self.dummy_size)
            else:
                image = Image.open(buff).convert('RGB')
                sample['data'] = self.transform_image(image)
                sample['data_aux'] = self.transform_aux_image(image)
            extras = ast.literal_eval(extra_str)
            try:
                for key, value in extras.items():
                    sample[key] = value
                # if you use dic the code below will need ['label']
                sample['label'] = 0
            except AttributeError:
                sample['label'] = int(extra_str)
            # Generate Soft Label
            soft_label = torch.Tensor(self.num_classes)
            if sample['label'] < 0:
                soft_label.fill_(1.0 / self.num_classes)
            else:
                soft_label.fill_(0)
                soft_label[sample['label']] = 1
            sample['soft_label'] = soft_label

        except Exception as e:
            logging.error('[{}] broken'.format(path))
            raise e
        return sample

openood/datasets/imglist_dataset.py
kai422/SCALE
import ast
import io
import logging
import os

import numpy as np
import torch
from PIL import Image, ImageFile
from torch.utils.data import Sampler

from .base_dataset import BaseDataset

# to fix "OSError: image file is truncated"
ImageFile.LOAD_TRUNCATED_IMAGES = True


class Convert:
    def __init__(self, mode='RGB'):
        self.mode = mode

    def __call__(self, image):
        return image.convert(self.mode)


class ImglistExtraDataDataset(BaseDataset):
    def __init__(self,
                 name,
                 imglist_pth,
                 data_dir,
                 num_classes,
                 preprocessor,
                 data_aux_preprocessor,
                 maxlen=None,
                 dummy_read=False,
                 dummy_size=None,
                 extra_data_pth=None,
                 extra_label_pth=None,
                 extra_percent=100,
                 **kwargs):
        super(ImglistExtraDataDataset, self).__init__(**kwargs)

        self.name = name
        with open(imglist_pth) as imgfile:
            self.imglist = imgfile.readlines()
        self.data_dir = data_dir
        self.num_classes = num_classes
        self.preprocessor = preprocessor
        self.transform_image = preprocessor
        self.transform_aux_image = data_aux_preprocessor
        self.maxlen = maxlen
        self.dummy_read = dummy_read
        self.dummy_size = dummy_size
        if dummy_read and dummy_size is None:
            raise ValueError(
                'if dummy_read is True, should provide dummy_size')

        self.orig_ids = list(range(len(self.imglist)))

        assert extra_data_pth is not None
        assert extra_label_pth is not None
        extra_data = np.load(extra_data_pth)
        extra_labels = np.load(extra_label_pth)
        assert len(extra_data) == len(extra_labels)

        self.extra_num = int(len(extra_labels) * extra_percent / 100.)
        self.total_num = len(self.imglist) + self.extra_num

        rng = np.random.RandomState(0)
        indices = rng.permutation(len(extra_labels))
        self.extra_data = extra_data[indices[:self.extra_num]]
        self.extra_labels = extra_labels[indices[:self.extra_num]]
        self.extra_ids = list(
            set(range(self.total_num)) - set(range(len(self.imglist))))

    def __len__(self):
        return self.total_num

    def getitem(self, index):
        if index in self.orig_ids:
            line = self.imglist[index].strip('\n')
            tokens = line.split(' ', 1)
            image_name, extra_str = tokens[0], tokens[1]
            if self.data_dir != '' and image_name.startswith('/'):
                raise RuntimeError('image_name starts with "/"')
            path = os.path.join(self.data_dir, image_name)
            sample = dict()
            sample['image_name'] = image_name
            kwargs = {'name': self.name, 'path': path, 'tokens': tokens}
            # some preprocessor methods require setup
            self.preprocessor.setup(**kwargs)
            try:
                if not self.dummy_read:
                    with open(path, 'rb') as f:
                        content = f.read()
                    filebytes = content
                    buff = io.BytesIO(filebytes)
                if self.dummy_size is not None:
                    sample['data'] = torch.rand(self.dummy_size)
                else:
                    image = Image.open(buff).convert('RGB')
                    sample['data'] = self.transform_image(image)
                    sample['data_aux'] = self.transform_aux_image(image)
                extras = ast.literal_eval(extra_str)
                try:
                    for key, value in extras.items():
                        sample[key] = value
                    # if you use dic the code below will need ['label']
                    sample['label'] = 0
                except AttributeError:
                    sample['label'] = int(extra_str)
                # Generate Soft Label
                soft_label = torch.Tensor(self.num_classes)
                if sample['label'] < 0:
                    soft_label.fill_(1.0 / self.num_classes)
                else:
                    soft_label.fill_(0)
                    soft_label[sample['label']] = 1
                sample['soft_label'] = soft_label

            except Exception as e:
                logging.error('[{}] broken'.format(path))
                raise e
            return sample
        else:
            ind = index - len(self.imglist)
            image = Image.fromarray(self.extra_data[ind])

            sample = dict()
            sample['image_name'] = str(ind)  # dummy name
            sample['data'] = self.transform_image(image)
            sample['data_aux'] = self.transform_aux_image(image)
            sample['label'] = self.extra_labels[ind]

            # Generate Soft Label
            soft_label = torch.Tensor(self.num_classes)
            if sample['label'] < 0:
                soft_label.fill_(1.0 / self.num_classes)
            else:
                soft_label.fill_(0)
                soft_label[sample['label']] = 1
            sample['soft_label'] = soft_label

            return sample


class TwoSourceSampler(Sampler):
    def __init__(self, real_inds, syn_inds, batch_size, real_ratio=0.5):
        assert len(real_inds) == 50000
        self.real_inds = real_inds
        self.syn_inds = syn_inds
        self.batch_size = batch_size
        self.real_batch_size = int(self.batch_size * real_ratio)
        self.syn_batch_size = self.batch_size - self.real_batch_size
        if real_ratio == 0:
            assert self.real_batch_size == 0
        elif real_ratio == 1:
            assert self.syn_batch_size == 0

        self.num_batches = int(np.ceil(len(self.real_inds) / self.batch_size))
        super().__init__(None)

    def __iter__(self):
        batch_counter = 0
        real_inds_shuffled = [
            self.real_inds[i] for i in torch.randperm(len(self.real_inds))
        ]
        syn_inds_shuffled = [
            self.syn_inds[i] for i in torch.randperm(len(self.syn_inds))
        ]

        real_offset = 0
        syn_offset = 0
        while batch_counter < self.num_batches:
            real_batch = real_inds_shuffled[
                real_offset:min(real_offset +
                                self.real_batch_size, len(real_inds_shuffled))]
            real_offset += self.real_batch_size

            syn_batch = syn_inds_shuffled[
                syn_offset:min(syn_offset +
                               self.syn_batch_size, len(syn_inds_shuffled))]
            syn_offset += self.syn_batch_size

            batch = real_batch + syn_batch
            np.random.shuffle(batch)
            yield batch
            batch_counter += 1

    def __len__(self):
        return self.num_batches

openood/datasets/imglist_extradata_dataset.py
kai422/SCALE
import ast
import io
import logging
import os

import numpy as np
import torch
from PIL import Image, ImageFile

from .imglist_dataset import ImglistDataset

# to fix "OSError: image file is truncated"
ImageFile.LOAD_TRUNCATED_IMAGES = True


class UDGDataset(ImglistDataset):
    def __init__(self,
                 name,
                 imglist_pth,
                 data_dir,
                 num_classes,
                 preprocessor,
                 data_aux_preprocessor,
                 maxlen=None,
                 dummy_read=False,
                 dummy_size=None,
                 **kwargs):
        super(UDGDataset,
              self).__init__(name, imglist_pth, data_dir, num_classes,
                             preprocessor, data_aux_preprocessor, maxlen,
                             dummy_read, dummy_size, **kwargs)

        self.cluster_id = np.zeros(len(self.imglist), dtype=int)
        self.cluster_reweight = np.ones(len(self.imglist), dtype=float)

        # use pseudo labels for unlabeled dataset during training
        self.pseudo_label = np.array(-1 * np.ones(len(self.imglist)),
                                     dtype=int)
        self.ood_conf = np.ones(len(self.imglist), dtype=float)

    def getitem(self, index):
        line = self.imglist[index].strip('\n')
        tokens = line.split(' ', 1)
        image_name, extra_str = tokens[0], tokens[1]
        if self.data_dir != '' and image_name.startswith('/'):
            raise RuntimeError('root not empty but image_name starts with "/"')
        path = os.path.join(self.data_dir, image_name)
        sample = dict()
        sample['image_name'] = image_name
        try:
            if not self.dummy_read:
                with open(path, 'rb') as f:
                    content = f.read()
                filebytes = content
                buff = io.BytesIO(filebytes)
            if self.dummy_size is not None:
                sample['data'] = torch.rand(self.dummy_size)
            else:
                image = Image.open(buff).convert('RGB')
                sample['data'] = self.transform_image(image)
                sample['data_aux'] = self.transform_aux_image(image)
            extras = ast.literal_eval(extra_str)
            try:
                for key, value in extras.items():
                    sample[key] = value
            except AttributeError:
                sample['label'] = int(extra_str)
            # Generate Soft Label
            soft_label = torch.Tensor(self.num_classes)
            if sample['label'] < 0:
                soft_label.fill_(1.0 / self.num_classes)
            else:
                soft_label.fill_(0)
                soft_label[sample['label']] = 1
            sample['soft_label'] = soft_label
            # Deep Clustering Aux Label Assignment for
            # both labeled/unlabeled data
            sample['cluster_id'] = self.cluster_id[index]
            sample['cluster_reweight'] = self.cluster_reweight[index]

            # Deep Clustering Pseudo Label Assignment for unlabeled data
            sample['pseudo_label'] = self.pseudo_label[index]
            soft_pseudo_label = torch.Tensor(len(sample['soft_label']))
            if sample['pseudo_label'] == -1:
                soft_pseudo_label.fill_(1.0 / len(sample['soft_label']))
            else:
                soft_pseudo_label.fill_(0.0)
                soft_pseudo_label[sample['pseudo_label']] = 1.0
            sample['pseudo_softlabel'] = soft_pseudo_label
            sample['ood_conf'] = self.ood_conf[index]

        except Exception as e:
            logging.error('[{}] broken'.format(path))
            raise e
        return sample

openood/datasets/udg_dataset.py
kai422/SCALE
import os
import torch
from numpy import load
from torch.utils.data import DataLoader

from openood.preprocessors.test_preprocessor import TestStandardPreProcessor
from openood.preprocessors.utils import get_preprocessor
from openood.utils.config import Config

from .feature_dataset import FeatDataset
from .imglist_dataset import ImglistDataset
from .imglist_augmix_dataset import ImglistAugMixDataset
from .imglist_extradata_dataset import ImglistExtraDataDataset, TwoSourceSampler
from .udg_dataset import UDGDataset


def get_dataloader(config: Config):
    # prepare a dataloader dictionary
    dataset_config = config.dataset
    dataloader_dict = {}
    for split in dataset_config.split_names:
        split_config = dataset_config[split]
        preprocessor = get_preprocessor(config, split)
        # weak augmentation for data_aux
        data_aux_preprocessor = TestStandardPreProcessor(config)

        if split_config.dataset_class == 'ImglistExtraDataDataset':
            dataset = ImglistExtraDataDataset(
                name=dataset_config.name + '_' + split,
                imglist_pth=split_config.imglist_pth,
                data_dir=split_config.data_dir,
                num_classes=dataset_config.num_classes,
                preprocessor=preprocessor,
                data_aux_preprocessor=data_aux_preprocessor,
                extra_data_pth=split_config.extra_data_pth,
                extra_label_pth=split_config.extra_label_pth,
                extra_percent=split_config.extra_percent)

            batch_sampler = TwoSourceSampler(dataset.orig_ids,
                                             dataset.extra_ids,
                                             split_config.batch_size,
                                             split_config.orig_ratio)

            dataloader = DataLoader(
                dataset,
                batch_sampler=batch_sampler,
                num_workers=dataset_config.num_workers,
            )
        elif split_config.dataset_class == 'ImglistAugMixDataset':
            dataset = ImglistAugMixDataset(
                name=dataset_config.name + '_' + split,
                imglist_pth=split_config.imglist_pth,
                data_dir=split_config.data_dir,
                num_classes=dataset_config.num_classes,
                preprocessor=preprocessor,
                data_aux_preprocessor=data_aux_preprocessor)
            sampler = None
            if dataset_config.num_gpus * dataset_config.num_machines > 1:
                sampler = torch.utils.data.distributed.DistributedSampler(
                    dataset)
                split_config.shuffle = False

            dataloader = DataLoader(dataset,
                                    batch_size=split_config.batch_size,
                                    shuffle=split_config.shuffle,
                                    num_workers=dataset_config.num_workers,
                                    sampler=sampler)
        else:
            CustomDataset = eval(split_config.dataset_class)
            dataset = CustomDataset(
                name=dataset_config.name + '_' + split,
                imglist_pth=split_config.imglist_pth,
                data_dir=split_config.data_dir,
                num_classes=dataset_config.num_classes,
                preprocessor=preprocessor,
                data_aux_preprocessor=data_aux_preprocessor)
            sampler = None
            if dataset_config.num_gpus * dataset_config.num_machines > 1:
                sampler = torch.utils.data.distributed.DistributedSampler(
                    dataset)
                split_config.shuffle = False

            dataloader = DataLoader(dataset,
                                    batch_size=split_config.batch_size,
                                    shuffle=split_config.shuffle,
                                    num_workers=dataset_config.num_workers,
                                    sampler=sampler)

        dataloader_dict[split] = dataloader
    return dataloader_dict


def get_ood_dataloader(config: Config):
    # specify custom dataset class
    ood_config = config.ood_dataset
    CustomDataset = eval(ood_config.dataset_class)
    dataloader_dict = {}
    for split in ood_config.split_names:
        split_config = ood_config[split]
        preprocessor = get_preprocessor(config, split)
        data_aux_preprocessor = TestStandardPreProcessor(config)
        if split == 'val':
            # validation set
            dataset = CustomDataset(
                name=ood_config.name + '_' + split,
                imglist_pth=split_config.imglist_pth,
                data_dir=split_config.data_dir,
                num_classes=ood_config.num_classes,
                preprocessor=preprocessor,
                data_aux_preprocessor=data_aux_preprocessor)
            dataloader = DataLoader(dataset,
                                    batch_size=ood_config.batch_size,
                                    shuffle=ood_config.shuffle,
                                    num_workers=ood_config.num_workers)
            dataloader_dict[split] = dataloader
        else:
            # dataloaders for csid, nearood, farood
            sub_dataloader_dict = {}
            for dataset_name in split_config.datasets:
                dataset_config = split_config[dataset_name]
                dataset = CustomDataset(
                    name=ood_config.name + '_' + split,
                    imglist_pth=dataset_config.imglist_pth,
                    data_dir=dataset_config.data_dir,
                    num_classes=ood_config.num_classes,
                    preprocessor=preprocessor,
                    data_aux_preprocessor=data_aux_preprocessor)
                dataloader = DataLoader(dataset,
                                        batch_size=ood_config.batch_size,
                                        shuffle=ood_config.shuffle,
                                        num_workers=ood_config.num_workers)
                sub_dataloader_dict[dataset_name] = dataloader
            dataloader_dict[split] = sub_dataloader_dict

    return dataloader_dict


def get_feature_dataloader(dataset_config: Config):
    # load in the cached feature
    loaded_data = load(dataset_config.feat_path, allow_pickle=True)
    total_feat = torch.from_numpy(loaded_data['feat_list'])
    del loaded_data
    # reshape the vector to fit in to the network
    total_feat.unsqueeze_(-1).unsqueeze_(-1)
    # let's see what we got here should be something like:
    # torch.Size([total_num, channel_size, 1, 1])
    print('Loaded feature size: {}'.format(total_feat.shape))

    split_config = dataset_config['train']

    dataset = FeatDataset(feat=total_feat)
    dataloader = DataLoader(dataset,
                            batch_size=split_config.batch_size,
                            shuffle=split_config.shuffle,
                            num_workers=dataset_config.num_workers)

    return dataloader


def get_feature_opengan_dataloader(dataset_config: Config):
    feat_root = dataset_config.feat_root

    dataloader_dict = {}
    for d in ['id_train', 'id_val', 'ood_val']:
        # load in the cached feature
        loaded_data = load(os.path.join(feat_root, f'{d}.npz'),
                           allow_pickle=True)
        total_feat = torch.from_numpy(loaded_data['feat_list'])
        total_labels = loaded_data['label_list']
        del loaded_data
        # reshape the vector to fit in to the network
        total_feat.unsqueeze_(-1).unsqueeze_(-1)
        # let's see what we got here should be something like:
        # torch.Size([total_num, channel_size, 1, 1])
        print('Loaded feature size: {}'.format(total_feat.shape))

        if d == 'id_train':
            split_config = dataset_config['train']
        else:
            split_config = dataset_config['val']

        dataset = FeatDataset(feat=total_feat, labels=total_labels)
        dataloader = DataLoader(dataset,
                                batch_size=split_config.batch_size,
                                shuffle=split_config.shuffle,
                                num_workers=dataset_config.num_workers)
        dataloader_dict[d] = dataloader

    return dataloader_dict

openood/datasets/utils.py
kai422/SCALE
import os
import gdown
import zipfile

from torch.utils.data import DataLoader
import torchvision as tvs
if tvs.__version__ >= '0.13':
    tvs_new = True
else:
    tvs_new = False

from openood.datasets.imglist_dataset import ImglistDataset
from openood.preprocessors import BasePreprocessor

from .preprocessor import get_default_preprocessor, ImageNetCPreProcessor

DATA_INFO = {
    'imagenet': {
        'num_classes': 1000,
        'id': {
            'train': {
                'data_dir': 'images_largescale/',
                'imglist_path': 'benchmark_imglist/imagenet/train_imagenet.txt'
            },
            'val': {
                'data_dir': 'images_largescale/',
                'imglist_path': 'benchmark_imglist/imagenet/val_imagenet.txt'
            },
            'test': {
                'data_dir': 'images_largescale/',
                'imglist_path': 'benchmark_imglist/imagenet/test_imagenet.txt'
            }
        },
        'ood': {
            'val': {
                'data_dir': 'images_largescale/',
                'imglist_path':
                'benchmark_imglist/imagenet/val_openimage_o.txt'
            },
            'near': {
                'datasets': ['ssb_hard', 'ninco'],
                'ssb_hard': {
                    'data_dir': 'images_largescale/',
                    'imglist_path':
                    'benchmark_imglist/imagenet/test_ssb_hard.txt'
                },
                'ninco': {
                    'data_dir': 'images_largescale/',
                    'imglist_path': 'benchmark_imglist/imagenet/test_ninco.txt'
                }
            },
            'far': {
                'datasets': ['inaturalist', 'textures', 'openimage_o'],
                'inaturalist': {
                    'data_dir':
                    'images_largescale/',
                    'imglist_path':
                    'benchmark_imglist/imagenet/test_inaturalist.txt'
                },
                'textures': {
                    'data_dir': 'images_classic/',
                    'imglist_path':
                    'benchmark_imglist/imagenet/test_textures.txt'
                },
                'openimage_o': {
                    'data_dir':
                    'images_largescale/',
                    'imglist_path':
                    'benchmark_imglist/imagenet/test_openimage_o.txt'
                },
            },
        }
    },
}

download_id_dict = {
    'textures': '1OSz1m3hHfVWbRdmMwKbUzoU8Hg9UKcam',
    'imagenet_1k': '1i1ipLDFARR-JZ9argXd2-0a6DXwVhXEj',
    'ssb_hard': '1PzkA-WGG8Z18h0ooL_pDdz9cO-DCIouE',
    'ninco': '1Z82cmvIB0eghTehxOGP5VTdLt7OD3nk6',
    'inaturalist': '1zfLfMvoUD0CUlKNnkk7LgxZZBnTBipdj',
    'openimage_o': '1VUFXnB_z70uHfdgJG2E_pjYOcEgqM7tE',
    'benchmark_imglist': '1XKzBdWCqg3vPoj-D32YixJyJJ0hL63gP'
}

dir_dict = {
    'images_classic/': [ 'texture'],
    'images_largescale/': [
        'imagenet_1k',
        'species_sub',
        'ssb_hard',
        'ninco',
        'inaturalist',
        'openimage_o',
    ],
}

benchmarks_dict = {
    'imagenet': [
        'imagenet_1k', 'ssb_hard', 'ninco', 'inaturalist', 'texture',
        'openimage_o'
    ],
}


def require_download(filename, path):
    for item in os.listdir(path):
        if item.startswith(filename) or filename.startswith(
                item) or path.endswith(filename):
            return False

    else:
        print(filename + ' needs download:')
        return True


def download_dataset(dataset, data_root):
    for key in dir_dict.keys():
        if dataset in dir_dict[key]:
            store_path = os.path.join(data_root, key, dataset)
            if not os.path.exists(store_path):
                os.makedirs(store_path)
            break
    else:
        print('Invalid dataset detected {}'.format(dataset))
        return

    if require_download(dataset, store_path):
        print(store_path)
        if not store_path.endswith('/'):
            store_path = store_path + '/'
        gdown.download(id=download_id_dict[dataset], output=store_path)

        file_path = os.path.join(store_path, dataset + '.zip')
        with zipfile.ZipFile(file_path, 'r') as zip_file:
            zip_file.extractall(store_path)
        os.remove(file_path)


def data_setup(data_root, id_data_name):
    if not data_root.endswith('/'):
        data_root = data_root + '/'

    if not os.path.exists(os.path.join(data_root, 'benchmark_imglist')):
        gdown.download(id=download_id_dict['benchmark_imglist'],
                       output=data_root)
        file_path = os.path.join(data_root, 'benchmark_imglist.zip')
        with zipfile.ZipFile(file_path, 'r') as zip_file:
            zip_file.extractall(data_root)
        os.remove(file_path)

    for dataset in benchmarks_dict[id_data_name]:
        download_dataset(dataset, data_root)


def get_id_ood_dataloader(id_name, data_root, preprocessor, **loader_kwargs):
    if 'imagenet' in id_name:
        if tvs_new:
            if isinstance(preprocessor,
                          tvs.transforms._presets.ImageClassification):
                mean, std = preprocessor.mean, preprocessor.std
            elif isinstance(preprocessor, tvs.transforms.Compose):
                temp = preprocessor.transforms[-1]
                mean, std = temp.mean, temp.std
            elif isinstance(preprocessor, BasePreprocessor):
                temp = preprocessor.transform.transforms[-1]
                mean, std = temp.mean, temp.std
            else:
                raise TypeError
        else:
            if isinstance(preprocessor, tvs.transforms.Compose):
                temp = preprocessor.transforms[-1]
                mean, std = temp.mean, temp.std
            elif isinstance(preprocessor, BasePreprocessor):
                temp = preprocessor.transform.transforms[-1]
                mean, std = temp.mean, temp.std
            else:
                raise TypeError
        imagenet_c_preprocessor = ImageNetCPreProcessor(mean, std)

    # weak augmentation for data_aux
    test_standard_preprocessor = get_default_preprocessor(id_name)

    dataloader_dict = {}
    data_info = DATA_INFO[id_name]

    # id
    sub_dataloader_dict = {}
    for split in data_info['id'].keys():
        dataset = ImglistDataset(
            name='_'.join((id_name, split)),
            imglist_pth=os.path.join(data_root,
                                     data_info['id'][split]['imglist_path']),
            data_dir=os.path.join(data_root,
                                  data_info['id'][split]['data_dir']),
            num_classes=data_info['num_classes'],
            preprocessor=preprocessor,
            data_aux_preprocessor=test_standard_preprocessor)
        dataloader = DataLoader(dataset, **loader_kwargs)
        sub_dataloader_dict[split] = dataloader
    dataloader_dict['id'] = sub_dataloader_dict

    # ood
    dataloader_dict['ood'] = {}
    for split in data_info['ood'].keys():
        split_config = data_info['ood'][split]

        if split == 'val':
            # validation set
            dataset = ImglistDataset(
                name='_'.join((id_name, 'ood', split)),
                imglist_pth=os.path.join(data_root,
                                         split_config['imglist_path']),
                data_dir=os.path.join(data_root, split_config['data_dir']),
                num_classes=data_info['num_classes'],
                preprocessor=preprocessor,
                data_aux_preprocessor=test_standard_preprocessor)
            dataloader = DataLoader(dataset, **loader_kwargs)
            dataloader_dict['ood'][split] = dataloader
        else:
            # dataloaders for nearood, farood
            sub_dataloader_dict = {}
            for dataset_name in split_config['datasets']:
                dataset_config = split_config[dataset_name]
                dataset = ImglistDataset(
                    name='_'.join((id_name, 'ood', dataset_name)),
                    imglist_pth=os.path.join(data_root,
                                             dataset_config['imglist_path']),
                    data_dir=os.path.join(data_root,
                                          dataset_config['data_dir']),
                    num_classes=data_info['num_classes'],
                    preprocessor=preprocessor,
                    data_aux_preprocessor=test_standard_preprocessor)
                dataloader = DataLoader(dataset, **loader_kwargs)
                sub_dataloader_dict[dataset_name] = dataloader
            dataloader_dict['ood'][split] = sub_dataloader_dict

    return dataloader_dict

openood/evaluation_api/datasets.py
kai422/SCALE
from typing import Callable, List, Type

import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.evaluators.metrics import compute_all_metrics
from openood.postprocessors import BasePostprocessor
from openood.networks.scale_net import ScaleNet

from .datasets import DATA_INFO, data_setup, get_id_ood_dataloader
from .postprocessor import get_postprocessor
from .preprocessor import get_default_preprocessor


class Evaluator:
    def __init__(
        self,
        net: nn.Module,
        id_name: str,
        data_root: str = './data',
        config_root: str = './configs',
        preprocessor: Callable = None,
        postprocessor_name: str = None,
        postprocessor: Type[BasePostprocessor] = None,
        batch_size: int = 200,
        shuffle: bool = False,
        num_workers: int = 4,
    ) -> None:
        """A unified, easy-to-use API for evaluating (most) discriminative OOD
        detection methods.

        Args:
            net (nn.Module):
                The base classifier.
            id_name (str):
                The name of the in-distribution dataset.
            data_root (str, optional):
                The path of the data folder. Defaults to './data'.
            config_root (str, optional):
                The path of the config folder. Defaults to './configs'.
            preprocessor (Callable, optional):
                The preprocessor of input images.
                Passing None will use the default preprocessor
                following convention. Defaults to None.
            postprocessor_name (str, optional):
                The name of the postprocessor that obtains OOD score.
                Ignored if an actual postprocessor is passed.
                Defaults to None.
            postprocessor (Type[BasePostprocessor], optional):
                An actual postprocessor instance which inherits
                OpenOOD's BasePostprocessor. Defaults to None.
            batch_size (int, optional):
                The batch size of samples. Defaults to 200.
            shuffle (bool, optional):
                Whether shuffling samples. Defaults to False.
            num_workers (int, optional):
                The num_workers argument that will be passed to
                data loaders. Defaults to 4.

        Raises:
            ValueError:
                If both postprocessor_name and postprocessor are None.
            ValueError:
                If the specified ID dataset {id_name} is not supported.
            TypeError:
                If the passed postprocessor does not inherit BasePostprocessor.
        """
        # check the arguments
        if postprocessor_name is None and postprocessor is None:
            raise ValueError('Please pass postprocessor_name or postprocessor')
        if postprocessor_name is not None and postprocessor is not None:
            print(
                'Postprocessor_name is ignored because postprocessor is passed'
            )
        if id_name not in DATA_INFO:
            raise ValueError(f'Dataset [{id_name}] is not supported')

        # get data preprocessor
        if preprocessor is None:
            preprocessor = get_default_preprocessor(id_name)

        # set up config root
        if config_root is None:
            filepath = os.path.dirname(os.path.abspath(__file__))
            config_root = os.path.join(*filepath.split('/')[:-2], 'configs')

        # get postprocessor
        if postprocessor is None:
            postprocessor = get_postprocessor(config_root, postprocessor_name,
                                              id_name)
        if not isinstance(postprocessor, BasePostprocessor):
            raise TypeError(
                'postprocessor should inherit BasePostprocessor in OpenOOD')

        # load data
        data_setup(data_root, id_name)
        loader_kwargs = {
            'batch_size': batch_size,
            'shuffle': shuffle,
            'num_workers': num_workers
        }
        dataloader_dict = get_id_ood_dataloader(id_name, data_root,
                                                preprocessor, **loader_kwargs)

        # wrap base model to work with certain postprocessors
        if postprocessor_name == 'scale':
            net = ScaleNet(net)

        # postprocessor setup
        postprocessor.setup(net, dataloader_dict['id'], dataloader_dict['ood'])

        self.id_name = id_name
        self.net = net
        self.preprocessor = preprocessor
        self.postprocessor = postprocessor
        self.dataloader_dict = dataloader_dict
        self.metrics = {
            'id_acc': None,
            'csid_acc': None,
            'ood': None,
            'fsood': None
        }
        self.scores = {
            'id': {
                'train': None,
                'val': None,
                'test': None
            },
            'ood': {
                'val': None,
                'near':
                {k: None
                 for k in dataloader_dict['ood']['near'].keys()},
                'far': {k: None
                        for k in dataloader_dict['ood']['far'].keys()},
            },
            'id_preds': None,
            'id_labels': None,
        }
        # perform hyperparameter search if have not done so
        if (self.postprocessor.APS_mode
                and not self.postprocessor.hyperparam_search_done):
            self.hyperparam_search()

        self.net.eval()

        # how to ensure the postprocessors can work with
        # models whose definition doesn't align with OpenOOD

    def _classifier_inference(self,
                              data_loader: DataLoader,
                              msg: str = 'Acc Eval',
                              progress: bool = True):
        self.net.eval()

        all_preds = []
        all_labels = []
        with torch.no_grad():
            for batch in tqdm(data_loader, desc=msg, disable=not progress):
                data = batch['data'].cuda()
                logits = self.net(data)
                preds = logits.argmax(1)
                all_preds.append(preds.cpu())
                all_labels.append(batch['label'])

        all_preds = torch.cat(all_preds)
        all_labels = torch.cat(all_labels)
        return all_preds, all_labels

    def eval_acc(self, data_name: str = 'id') -> float:
        if data_name == 'id':
            if self.metrics['id_acc'] is not None:
                return self.metrics['id_acc']
            else:
                if self.scores['id_preds'] is None:
                    all_preds, all_labels = self._classifier_inference(
                        self.dataloader_dict['id']['test'], 'ID Acc Eval')
                    self.scores['id_preds'] = all_preds
                    self.scores['id_labels'] = all_labels
                else:
                    all_preds = self.scores['id_preds']
                    all_labels = self.scores['id_labels']
                assert len(all_preds) == len(all_labels)
                correct = (all_preds == all_labels).sum().item()
                acc = correct / len(all_labels) * 100
                self.metrics['id_acc'] = acc
                return acc
        elif data_name == 'csid':
            if self.metrics['csid_acc'] is not None:
                return self.metrics['csid_acc']
            else:
                correct, total = 0, 0
                for _, (dataname, dataloader) in enumerate(
                        self.dataloader_dict['csid'].items()):
                    if self.scores['csid_preds'][dataname] is None:
                        all_preds, all_labels = self._classifier_inference(
                            dataloader, f'CSID {dataname} Acc Eval')
                        self.scores['csid_preds'][dataname] = all_preds
                        self.scores['csid_labels'][dataname] = all_labels
                    else:
                        all_preds = self.scores['csid_preds'][dataname]
                        all_labels = self.scores['csid_labels'][dataname]

                    assert len(all_preds) == len(all_labels)
                    c = (all_preds == all_labels).sum().item()
                    t = len(all_labels)
                    correct += c
                    total += t

                if self.scores['id_preds'] is None:
                    all_preds, all_labels = self._classifier_inference(
                        self.dataloader_dict['id']['test'], 'ID Acc Eval')
                    self.scores['id_preds'] = all_preds
                    self.scores['id_labels'] = all_labels
                else:
                    all_preds = self.scores['id_preds']
                    all_labels = self.scores['id_labels']

                correct += (all_preds == all_labels).sum().item()
                total += len(all_labels)

                acc = correct / total * 100
                self.metrics['csid_acc'] = acc
                return acc
        else:
            raise ValueError(f'Unknown data name {data_name}')

    def eval_ood(self, fsood: bool = False, progress: bool = True):
        id_name = 'id' if not fsood else 'csid'
        task = 'ood' if not fsood else 'fsood'
        if self.metrics[task] is None:
            self.net.eval()

            # id score
            if self.scores['id']['test'] is None:
                print(f'Performing inference on {self.id_name} test set...',
                      flush=True)
                id_pred, id_conf, id_gt = self.postprocessor.inference(
                    self.net, self.dataloader_dict['id']['test'], progress)
                self.scores['id']['test'] = [id_pred, id_conf, id_gt]
            else:
                id_pred, id_conf, id_gt = self.scores['id']['test']

            if fsood:
                csid_pred, csid_conf, csid_gt = [], [], []
                for i, dataset_name in enumerate(self.scores['csid'].keys()):
                    if self.scores['csid'][dataset_name] is None:
                        print(
                            f'Performing inference on {self.id_name} '
                            f'(cs) test set [{i+1}]: {dataset_name}...',
                            flush=True)
                        temp_pred, temp_conf, temp_gt = \
                            self.postprocessor.inference(
                                self.net,
                                self.dataloader_dict['csid'][dataset_name],
                                progress)
                        self.scores['csid'][dataset_name] = [
                            temp_pred, temp_conf, temp_gt
                        ]

                    csid_pred.append(self.scores['csid'][dataset_name][0])
                    csid_conf.append(self.scores['csid'][dataset_name][1])
                    csid_gt.append(self.scores['csid'][dataset_name][2])

                csid_pred = np.concatenate(csid_pred)
                csid_conf = np.concatenate(csid_conf)
                csid_gt = np.concatenate(csid_gt)

                id_pred = np.concatenate((id_pred, csid_pred))
                id_conf = np.concatenate((id_conf, csid_conf))
                id_gt = np.concatenate((id_gt, csid_gt))

            self.scores['id_labels'] = torch.LongTensor(id_gt)
            self.scores['id_preds'] = torch.LongTensor(id_pred)
            # load nearood data and compute ood metrics
            near_metrics = self._eval_ood([id_pred, id_conf, id_gt],
                                          ood_split='near',
                                          progress=progress)
            # load farood data and compute ood metrics
            far_metrics = self._eval_ood([id_pred, id_conf, id_gt],
                                         ood_split='far',
                                         progress=progress)

            if self.metrics[f'{id_name}_acc'] is None:
                self.eval_acc(id_name)
            near_metrics[:, -1] = np.array([self.metrics[f'{id_name}_acc']] *
                                           len(near_metrics))
            far_metrics[:, -1] = np.array([self.metrics[f'{id_name}_acc']] *
                                          len(far_metrics))

            self.metrics[task] = pd.DataFrame(
                np.concatenate([near_metrics, far_metrics], axis=0),
                index=list(self.dataloader_dict['ood']['near'].keys()) +
                ['nearood'] + list(self.dataloader_dict['ood']['far'].keys()) +
                ['farood'],
                columns=['FPR@95', 'AUROC', 'AUPR_IN', 'AUPR_OUT', 'ACC'],
            )
        else:
            print('Evaluation has already been done!')

        with pd.option_context(
                'display.max_rows', None, 'display.max_columns', None,
                'display.float_format',
                '{:,.2f}'.format):  # more options can be specified also
            print(self.metrics[task])

        return self.metrics[task]

    def _eval_ood(self,
                  id_list: List[np.ndarray],
                  ood_split: str = 'near',
                  progress: bool = True):
        print(f'Processing {ood_split} ood...', flush=True)
        [id_pred, id_conf, id_gt] = id_list
        metrics_list = []
        for dataset_name, ood_dl in self.dataloader_dict['ood'][
                ood_split].items():
            if self.scores['ood'][ood_split][dataset_name] is None:
                print(f'Performing inference on {dataset_name} dataset...',
                      flush=True)
                ood_pred, ood_conf, ood_gt = self.postprocessor.inference(
                    self.net, ood_dl, progress)
                self.scores['ood'][ood_split][dataset_name] = [
                    ood_pred, ood_conf, ood_gt
                ]
            else:
                print(
                    'Inference has been performed on '
                    f'{dataset_name} dataset...',
                    flush=True)
                [ood_pred, ood_conf,
                 ood_gt] = self.scores['ood'][ood_split][dataset_name]

            ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood
            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])
        

            print(f'Computing metrics on {dataset_name} dataset...')
            ood_metrics = compute_all_metrics(conf, label, pred)
            metrics_list.append(ood_metrics)
            self._print_metrics(ood_metrics)

        print('Computing mean metrics...', flush=True)
        metrics_list = np.array(metrics_list)
        metrics_mean = np.mean(metrics_list, axis=0, keepdims=True)
        self._print_metrics(list(metrics_mean[0]))
        return np.concatenate([metrics_list, metrics_mean], axis=0) * 100

    def _print_metrics(self, metrics):
        [fpr, auroc, aupr_in, aupr_out, _] = metrics

        # print ood metric results
        print('FPR@95: {:.2f}, AUROC: {:.2f}'.format(100 * fpr, 100 * auroc),
              end=' ',
              flush=True)
        print('AUPR_IN: {:.2f}, AUPR_OUT: {:.2f}'.format(
            100 * aupr_in, 100 * aupr_out),
              flush=True)
        print(u'\u2500' * 70, flush=True)
        print('', flush=True)

    def hyperparam_search(self):
        print('Starting automatic parameter search...')
        max_auroc = 0
        hyperparam_names = []
        hyperparam_list = []
        count = 0

        for name in self.postprocessor.args_dict.keys():
            hyperparam_names.append(name)
            count += 1

        for name in hyperparam_names:
            hyperparam_list.append(self.postprocessor.args_dict[name])

        hyperparam_combination = self.recursive_generator(
            hyperparam_list, count)

        final_index = None
        for i, hyperparam in enumerate(hyperparam_combination):
            self.postprocessor.set_hyperparam(hyperparam)

            id_pred, id_conf, id_gt = self.postprocessor.inference(
                self.net, self.dataloader_dict['id']['val'])
            ood_pred, ood_conf, ood_gt = self.postprocessor.inference(
                self.net, self.dataloader_dict['ood']['val'])

            ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood
            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])
            ood_metrics = compute_all_metrics(conf, label, pred)
            auroc = ood_metrics[1]

            print('Hyperparam: {}, auroc: {}'.format(hyperparam, auroc))
            if auroc > max_auroc:
                final_index = i
                max_auroc = auroc

        self.postprocessor.set_hyperparam(hyperparam_combination[final_index])
        print('Final hyperparam: {}'.format(
            self.postprocessor.get_hyperparam()))
        self.postprocessor.hyperparam_search_done = True

    def recursive_generator(self, list, n):
        if n == 1:
            results = []
            for x in list[0]:
                k = []
                k.append(x)
                results.append(k)
            return results
        else:
            results = []
            temp = self.recursive_generator(list, n - 1)
            for x in list[n - 1]:
                for y in temp:
                    k = y.copy()
                    k.append(x)
                    results.append(k)
            return results

openood/evaluation_api/evaluator.py
kai422/SCALE
import os
import urllib.request

from openood.postprocessors import (
    ScalePostprocessor, BasePostprocessor)
from openood.utils.config import Config, merge_configs

postprocessors = {
    'scale': ScalePostprocessor,
}

link_prefix = 'https://raw.githubusercontent.com/Jingkang50/OpenOOD/main/configs/postprocessors/'


def get_postprocessor(config_root: str, postprocessor_name: str,
                      id_data_name: str):
    postprocessor_config_path = os.path.join(config_root, 'postprocessors',
                                             f'{postprocessor_name}.yml')
    if not os.path.exists(postprocessor_config_path):
        os.makedirs(os.path.dirname(postprocessor_config_path), exist_ok=True)
        urllib.request.urlretrieve(link_prefix + f'{postprocessor_name}.yml',
                                   postprocessor_config_path)

    config = Config(postprocessor_config_path)
    config = merge_configs(config,
                           Config(**{'dataset': {
                               'name': id_data_name
                           }}))
    postprocessor = postprocessors[postprocessor_name](config)
    postprocessor.APS_mode = config.postprocessor.APS_mode
    postprocessor.hyperparam_search_done = False
    return postprocessor

openood/evaluation_api/postprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.preprocessors import BasePreprocessor
from openood.utils import Config

INTERPOLATION = tvs_trans.InterpolationMode.BILINEAR

default_preprocessing_dict = {
    'cifar10': {
        'pre_size': 32,
        'img_size': 32,
        'normalization': [[0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]],
    },
    'cifar100': {
        'pre_size': 32,
        'img_size': 32,
        'normalization': [[0.5071, 0.4867, 0.4408], [0.2675, 0.2565, 0.2761]],
    },
    'imagenet': {
        'pre_size': 256,
        'img_size': 224,
        'normalization': [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
    },
    'imagenet200': {
        'pre_size': 256,
        'img_size': 224,
        'normalization': [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
    },
    'aircraft': {
        'pre_size': 512,
        'img_size': 448,
        'normalization': [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
    },
    'cub': {
        'pre_size': 512,
        'img_size': 448,
        'normalization': [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
    }
}


class Convert:
    def __init__(self, mode='RGB'):
        self.mode = mode

    def __call__(self, image):
        return image.convert(self.mode)


class TestStandardPreProcessor(BasePreprocessor):
    """For test and validation dataset standard image transformation."""
    def __init__(self, config: Config):
        self.transform = tvs_trans.Compose([
            Convert('RGB'),
            tvs_trans.Resize(config.pre_size, interpolation=INTERPOLATION),
            tvs_trans.CenterCrop(config.img_size),
            tvs_trans.ToTensor(),
            tvs_trans.Normalize(*config.normalization),
        ])


class ImageNetCPreProcessor(BasePreprocessor):
    def __init__(self, mean, std):
        self.transform = tvs_trans.Compose([
            tvs_trans.ToTensor(),
            tvs_trans.Normalize(mean, std),
        ])


def get_default_preprocessor(data_name: str):
    # TODO: include fine-grained datasets proposed in Vaze et al.?

    if data_name not in default_preprocessing_dict:
        raise NotImplementedError(f'The dataset {data_name} is not supported')

    config = Config(**default_preprocessing_dict[data_name])
    preprocessor = TestStandardPreProcessor(config)

    return preprocessor

openood/evaluation_api/preprocessor.py
kai422/SCALE
from .utils import get_evaluator

openood/evaluators/__init__.py
kai422/SCALE
import numpy as np
import torch
from sklearn.metrics import auc, roc_curve

from openood.utils import Config


class ADEvaluator():
    def __init__(self, config: Config):
        self.config = config

    def eval_ood(self,
                 net,
                 id_data_loader,
                 ood_data_loaders,
                 postprocessor,
                 epoch_idx: int = -1):
        with torch.no_grad():
            if type(net) is dict:
                for subnet in net.values():
                    subnet.eval()
            else:
                net.eval()
            auroc = self.get_auroc(net, id_data_loader['test'],
                                   ood_data_loaders['val'], postprocessor)
            metrics = {
                'epoch_idx': epoch_idx,
                'image_auroc': auroc,
            }
            return metrics

    def report(self, test_metrics):

        print('Complete Evaluation:\n'
              '{}\n'
              '==============================\n'
              'AUC Image: {:.2f} \n'
              '=============================='.format(
                  self.config.dataset.name,
                  100.0 * test_metrics['image_auroc']),
              flush=True)
        print('Completed!', flush=True)

    def get_auroc(self, net, id_data_loader, ood_data_loader, postprocessor):
        _, id_conf, id_gt = postprocessor.inference(net, id_data_loader)
        _, ood_conf, ood_gt = postprocessor.inference(net, ood_data_loader)
        ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood

        conf = np.concatenate([id_conf, ood_conf])
        label = np.concatenate([id_gt, ood_gt])

        ind_indicator = np.zeros_like(label)
        ind_indicator[label != -1] = 1

        fpr, tpr, _ = roc_curve(ind_indicator, conf)

        auroc = auc(fpr, tpr)

        return auroc

openood/evaluators/ad_evaluator.py
kai422/SCALE
from typing import Dict

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.postprocessors import BasePostprocessor
from openood.utils import Config

from .ood_evaluator import OODEvaluator


class ARPLEvaluator(OODEvaluator):
    def __init__(self, config: Config):
        self.config = config

    def eval_acc(self,
                 net: dict,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1,
                 fsood: bool = False,
                 csid_data_loaders: Dict[str, DataLoader] = None):
        criterion = net['criterion']
        net = net['netF']
        net.eval()

        loss_avg = 0.0
        correct = 0
        with torch.no_grad():
            for batch in tqdm(data_loader,
                              desc='Eval: ',
                              position=0,
                              leave=True):
                # prepare data
                data = batch['data'].cuda()
                target = batch['label'].cuda()

                # forward
                _, feat = net(data, return_feature=True)
                output, loss = criterion(feat, target)

                # accuracy
                pred = output.data.max(1)[1]
                correct += pred.eq(target.data).sum().item()

                # test loss average
                loss_avg += float(loss.data)

        if not fsood:
            metrics = {}
            metrics['epoch_idx'] = epoch_idx
            metrics['loss'] = loss_avg / len(data_loader)
            metrics['acc'] = correct / len(data_loader.dataset)
            return metrics
        else:
            all_correct = 0
            all_total = 0
            all_correct += correct
            all_total += len(data_loader.dataset)

            assert csid_data_loaders is not None
            for dataset_name, csid_dl in csid_data_loaders.items():
                correct = 0
                with torch.no_grad():
                    for batch in tqdm(csid_dl,
                                      desc='Eval: ',
                                      position=0,
                                      leave=True):
                        # prepare data
                        data = batch['data'].cuda()
                        target = batch['label'].cuda()

                        # forward
                        _, feat = net(data, return_feature=True)
                        output, loss = criterion(feat, target)

                        # accuracy
                        pred = output.data.max(1)[1]
                        correct += pred.eq(target.data).sum().item()

                all_correct += correct
                all_total += len(csid_dl.dataset)

            metrics = {}
            metrics['epoch_idx'] = epoch_idx
            metrics['acc'] = all_correct / all_total
            return metrics

    def eval_ood(self,
                 net: dict,
                 id_data_loader: DataLoader,
                 ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                 postprocessor: BasePostprocessor,
                 fsood: bool = False):
        criterion = net['criterion']
        net = net['netF']
        net = nn.Sequential(
            net,
            criterion,
        )
        net.eval()
        # load training in-distribution data
        assert 'test' in id_data_loader, \
            'id_data_loaders should have the key: test!'
        dataset_name = self.config.dataset.name
        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_pred, id_conf, id_gt = postprocessor.inference(
            net, id_data_loader['test'])
        if self.config.recorder.save_scores:
            self._save_scores(id_pred, id_conf, id_gt, dataset_name)

        if fsood:
            # load csid data and compute confidence
            for dataset_name, csid_dl in ood_data_loaders['csid'].items():
                print(f'Performing inference on {dataset_name} dataset...',
                      flush=True)
                csid_pred, csid_conf, csid_gt = postprocessor.inference(
                    net, csid_dl)
                if self.config.recorder.save_scores:
                    self._save_scores(csid_pred, csid_conf, csid_gt,
                                      dataset_name)
                id_pred = np.concatenate([id_pred, csid_pred])
                id_conf = np.concatenate([id_conf, csid_conf])
                id_gt = np.concatenate([id_gt, csid_gt])

        # load nearood data and compute ood metrics
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='nearood')

        # load farood data and compute ood metrics
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='farood')

openood/evaluators/arpl_evaluator.py
kai422/SCALE
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.postprocessors import BasePostprocessor
from openood.utils import Config


def to_np(x):
    return x.data.cpu().numpy()


class BaseEvaluator:
    def __init__(self, config: Config):
        self.config = config

    def eval_acc(self,
                 net: nn.Module,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1):
        net.eval()

        loss_avg = 0.0
        correct = 0
        with torch.no_grad():
            for batch in tqdm(data_loader,
                              desc='Eval: ',
                              position=0,
                              leave=True,
                              disable=not comm.is_main_process()):
                # prepare data
                data = batch['data'].cuda()
                target = batch['label'].cuda()

                # forward
                output = net(data)

                loss = F.cross_entropy(output, target)

                # accuracy
                pred = output.data.max(1)[1]
                correct += pred.eq(target.data).sum().item()

                # test loss average
                loss_avg += float(loss.data)

        loss = loss_avg / len(data_loader)
        acc = correct / len(data_loader.dataset)

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss)
        metrics['acc'] = self.save_metrics(acc)
        return metrics

    def extract(self,
                net: nn.Module,
                data_loader: DataLoader,
                filename: str = 'feature'):
        net.eval()
        feat_list, label_list = [], []

        with torch.no_grad():
            for batch in tqdm(data_loader,
                              desc='Feature Extracting: ',
                              position=0,
                              leave=True,
                              disable=not comm.is_main_process()):
                data = batch['data'].cuda()
                label = batch['label']

                _, feat = net(data, return_feature=True)
                feat_list.extend(to_np(feat))
                label_list.extend(to_np(label))

        feat_list = np.array(feat_list)
        label_list = np.array(label_list)

        save_dir = self.config.output_dir
        os.makedirs(save_dir, exist_ok=True)
        np.savez(os.path.join(save_dir, filename),
                 feat_list=feat_list,
                 label_list=label_list)

    def save_metrics(self, value):
        all_values = comm.gather(value)
        temp = 0
        for i in all_values:
            temp = temp + i
        # total_value = np.add([x for x in all_values])s

        return temp

openood/evaluators/base_evaluator.py
kai422/SCALE
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.postprocessors import BasePostprocessor
from openood.utils import Config
from .base_evaluator import BaseEvaluator


class ECEEvaluator(BaseEvaluator):
    def __init__(self, config: Config):
        """OOD Evaluator.

        Args:
            config (Config): Config file from
        """
        super(ECEEvaluator, self).__init__(config)

    def eval_acc(self,
                 net: nn.Module,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1,
                 num_bins: int = 15):
        net.eval()
        """Calculates ECE.
        Args:
          num_bins: the number of bins to partition all samples. we set it as 15.
        Returns:
          ece: the calculated ECE value.
        """

        loss_avg = 0.0
        correct = 0
        total_scores = []
        total_preds = []
        total_labels = []
        with torch.no_grad():
            for batch in tqdm(data_loader,
                              desc='Eval: ',
                              position=0,
                              leave=True):
                # prepare data
                data = batch['data'].cuda()
                target = batch['label'].cuda()

                # forward
                output = net(data)
                loss = F.cross_entropy(output, target)

                # accuracy
                pred = output.data.max(1)[1]
                score = output.data.max(1)[0]
                correct += pred.eq(target.data).sum().item()

                # test loss average
                loss_avg += float(loss.data)

                total_preds.append(pred.cpu().numpy().reshape(-1))
                total_scores.append(score.cpu().numpy().reshape(-1))
                total_labels.append(target.data.cpu().numpy().reshape(-1))

        scores_np = np.reshape(total_scores, -1)
        preds_np = np.reshape(total_preds, -1)
        labels_np = np.reshape(total_labels, -1)
        acc_tab = np.zeros(num_bins)  # Empirical (true) confidence
        mean_conf = np.zeros(num_bins)  # Predicted confidence
        nb_items_bin = np.zeros(num_bins)  # Number of items in the bins
        tau_tab = np.linspace(0, 1, num_bins + 1)  # Confidence bins
        for i in np.arange(num_bins):  # Iterates over the bins
            # Selects the items where the predicted max probability falls in the bin
            # [tau_tab[i], tau_tab[i + 1)]
            sec = (tau_tab[i + 1] > scores_np) & (scores_np >= tau_tab[i])
            nb_items_bin[i] = np.sum(sec)  # Number of items in the bin
            # Selects the predicted classes, and the true classes
            class_pred_sec, y_sec = preds_np[sec], labels_np[sec]
            # Averages of the predicted max probabilities
            mean_conf[i] = np.mean(
                scores_np[sec]) if nb_items_bin[i] > 0 else np.nan
            # Computes the empirical confidence
            acc_tab[i] = np.mean(
                class_pred_sec == y_sec) if nb_items_bin[i] > 0 else np.nan
        # Cleaning
        mean_conf = mean_conf[nb_items_bin > 0]
        acc_tab = acc_tab[nb_items_bin > 0]
        nb_items_bin = nb_items_bin[nb_items_bin > 0]
        if sum(nb_items_bin) != 0:
            ece = np.average(
                np.absolute(mean_conf - acc_tab),
                weights=nb_items_bin.astype(np.float) / np.sum(nb_items_bin))
        else:
            ece = 0.0

        loss = loss_avg / len(data_loader)
        acc = correct / len(data_loader.dataset)

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss)
        metrics['acc'] = self.save_metrics(acc)
        metrics['ece'] = self.save_metrics(ece)
        return metrics

openood/evaluators/ece_evaluator.py
kai422/SCALE
import csv
import os
from typing import Dict, List

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from openood.postprocessors import BasePostprocessor

from .ood_evaluator import OODEvaluator


class FSOODEvaluator(OODEvaluator):
    def eval_csid_acc(self, net: nn.Module,
                      csid_loaders: Dict[str, Dict[str, DataLoader]]):
        # ensure the networks in eval mode
        net.eval()
        for dataset_name, csid_dl in csid_loaders.items():
            print(f'Computing accuracy on {dataset_name} dataset...')
            correct = 0
            with torch.no_grad():
                for batch in csid_dl:
                    data = batch['data'].cuda()
                    target = batch['label'].cuda()
                    # forward
                    output = net(data)
                    # accuracy
                    pred = output.data.max(1)[1]
                    correct += pred.eq(target.data).sum().item()
            acc = correct / len(csid_dl.dataset)
            if self.config.recorder.save_csv:
                self._save_acc_results(acc, dataset_name)
        print(u'\u2500' * 70, flush=True)

    def _save_acc_results(self, acc, dataset_name):
        write_content = {
            'dataset': dataset_name,
            'FPR@95': '-',
            'AUROC': '-',
            'AUPR_IN': '-',
            'AUPR_OUT': '-',
            'ACC': '{:.2f}'.format(100 * acc),
        }
        fieldnames = list(write_content.keys())
        # print csid metric results
        print('CSID[{}] accuracy: {:.2f}%'.format(dataset_name, 100 * acc),
              flush=True)
        csv_path = os.path.join(self.config.output_dir, 'csid.csv')
        if not os.path.exists(csv_path):
            with open(csv_path, 'w', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerow(write_content)
        else:
            with open(csv_path, 'a', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writerow(write_content)

    def eval_ood(self, net: nn.Module, id_data_loader: List[DataLoader],
                 ood_data_loaders: List[DataLoader],
                 postprocessor: BasePostprocessor):
        # ensure the networks in eval mode
        net.eval()
        # load training in-distribution data
        assert 'test' in id_data_loader, \
            'id_data_loaders should have the key: test!'
        dataset_name = self.config.dataset.name
        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_pred, id_conf, id_gt = postprocessor.inference(
            net, id_data_loader['test'])
        if self.config.recorder.save_scores:
            self._save_scores(id_pred, id_conf, id_gt, dataset_name)

        # load csid data and compute confidence
        for dataset_name, csid_dl in ood_data_loaders['csid'].items():
            print(f'Performing inference on {dataset_name} dataset...',
                  flush=True)
            csid_pred, csid_conf, csid_gt = postprocessor.inference(
                net, csid_dl)
            if self.config.recorder.save_scores:
                self._save_scores(csid_pred, csid_conf, csid_gt, dataset_name)
            id_pred = np.concatenate([id_pred, csid_pred])
            id_conf = np.concatenate([id_conf, csid_conf])
            id_gt = np.concatenate([id_gt, csid_gt])

        # compute accuracy on csid
        print(u'\u2500' * 70, flush=True)
        self.eval_csid_acc(net, ood_data_loaders['csid'])

        # load nearood data and compute ood metrics
        print(u'\u2500' * 70, flush=True)
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='nearood')

        # load farood data and compute ood metrics
        print(u'\u2500' * 70, flush=True)
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='farood')

openood/evaluators/fsood_evaluator.py
kai422/SCALE
import numpy as np
from sklearn import metrics


def compute_all_metrics(conf, label, pred):
    np.set_printoptions(precision=3)
    recall = 0.95

    print(conf, label, pred)
    auroc, aupr_in, aupr_out, fpr = auc_and_fpr_recall(conf, label, recall)

    accuracy = acc(pred, label)

    results = [fpr, auroc, aupr_in, aupr_out, accuracy]

    return results


# accuracy
def acc(pred, label):
    ind_pred = pred[label != -1]
    ind_label = label[label != -1]

    num_tp = np.sum(ind_pred == ind_label)
    acc = num_tp / len(ind_label)

    return acc


# fpr_recall
def fpr_recall(conf, label, tpr):
    gt = np.ones_like(label)
    gt[label == -1] = 0

    fpr_list, tpr_list, threshold_list = metrics.roc_curve(gt, conf)
    fpr = fpr_list[np.argmax(tpr_list >= tpr)]
    thresh = threshold_list[np.argmax(tpr_list >= tpr)]
    return fpr, thresh


# auc
def auc_and_fpr_recall(conf, label, tpr_th):
    # following convention in ML we treat OOD as positive
    ood_indicator = np.zeros_like(label)
    ood_indicator[label == -1] = 1

    # in the postprocessor we assume ID samples will have larger
    # "conf" values than OOD samples
    # therefore here we need to negate the "conf" values
    fpr_list, tpr_list, thresholds = metrics.roc_curve(ood_indicator, -conf)
    fpr = fpr_list[np.argmax(tpr_list >= tpr_th)]

    precision_in, recall_in, thresholds_in \
        = metrics.precision_recall_curve(ood_indicator, -conf)

    precision_out, recall_out, thresholds_out \
        = metrics.precision_recall_curve(1 - ood_indicator, conf)

    auroc = metrics.auc(fpr_list, tpr_list)
    aupr_in = metrics.auc(recall_in, precision_in)
    aupr_out = metrics.auc(recall_out, precision_out)

    return auroc, aupr_in, aupr_out, fpr


# ccr_fpr
def ccr_fpr(conf, fpr, pred, label):
    ind_conf = conf[label != -1]
    ind_pred = pred[label != -1]
    ind_label = label[label != -1]

    ood_conf = conf[label == -1]

    num_ind = len(ind_conf)
    num_ood = len(ood_conf)

    fp_num = int(np.ceil(fpr * num_ood))
    thresh = np.sort(ood_conf)[-fp_num]
    num_tp = np.sum((ind_conf > thresh) * (ind_pred == ind_label))
    ccr = num_tp / num_ind

    return ccr


def detection(ind_confidences,
              ood_confidences,
              n_iter=100000,
              return_data=False):
    # calculate the minimum detection error
    Y1 = ood_confidences
    X1 = ind_confidences

    start = np.min([np.min(X1), np.min(Y1)])
    end = np.max([np.max(X1), np.max(Y1)])
    gap = (end - start) / n_iter

    best_error = 1.0
    best_delta = None
    all_thresholds = []
    all_errors = []
    for delta in np.arange(start, end, gap):
        tpr = np.sum(np.sum(X1 < delta)) / np.float(len(X1))
        error2 = np.sum(np.sum(Y1 > delta)) / np.float(len(Y1))
        detection_error = (tpr + error2) / 2.0

        if return_data:
            all_thresholds.append(delta)
            all_errors.append(detection_error)

        if detection_error < best_error:
            best_error = np.minimum(best_error, detection_error)
            best_delta = delta

    if return_data:
        return best_error, best_delta, all_errors, all_thresholds
    else:
        return best_error, best_delta

openood/evaluators/metrics.py
kai422/SCALE
import csv
import os
from typing import Dict, List

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.postprocessors import BasePostprocessor
from openood.utils import Config

from .base_evaluator import BaseEvaluator
from .metrics import compute_all_metrics


def topk(output, target, ks=(1, )):
    """Returns one boolean vector for each k, whether the target is within the
    output's top-k."""
    _, pred = output.topk(max(ks), 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))
    return [correct[:k].max(0)[0] for k in ks]


def get_group_slices(classes_per_group):
    group_slices = []
    start = 0
    for num_cls in classes_per_group:
        end = start + num_cls + 1
        group_slices.append([start, end])
        start = end
    return torch.LongTensor(group_slices)


def cal_ood_score(logits, group_slices):
    num_groups = group_slices.shape[0]

    all_group_ood_score_MOS = []
    for i in range(num_groups):
        group_logit = logits[:, group_slices[i][0]:group_slices[i][1]]

        group_softmax = F.softmax(group_logit, dim=-1)
        group_others_score = group_softmax[:, 0]

        all_group_ood_score_MOS.append(-group_others_score)

    all_group_ood_score_MOS = torch.stack(all_group_ood_score_MOS, dim=1)
    final_max_score_MOS, _ = torch.max(all_group_ood_score_MOS, dim=1)
    return final_max_score_MOS.data.cpu().numpy()


def iterate_data(data_loader, model, group_slices):
    confs_mos = []
    dataiter = iter(data_loader)

    with torch.no_grad():
        for _ in tqdm(range(1,
                            len(dataiter) + 1),
                      desc='Batches',
                      position=0,
                      leave=True,
                      disable=not comm.is_main_process()):
            batch = next(dataiter)
            data = batch['data'].cuda()

            logits = model(data)
            conf_mos = cal_ood_score(logits, group_slices)
            confs_mos.extend(conf_mos)

    return np.array(confs_mos)


def calc_group_softmax_acc(logits, labels, group_slices):
    num_groups = group_slices.shape[0]
    loss = 0
    num_samples = logits.shape[0]

    all_group_max_score, all_group_max_class = [], []

    smax = torch.nn.Softmax(dim=-1).cuda()
    cri = torch.nn.CrossEntropyLoss(reduction='none').cuda()

    for i in range(num_groups):
        group_logit = logits[:, group_slices[i][0]:group_slices[i][1]]
        group_label = labels[:, i]
        loss += cri(group_logit, group_label)

        group_softmax = smax(group_logit)
        group_softmax = group_softmax[:, 1:]  # disregard others category
        group_max_score, group_max_class = torch.max(group_softmax, dim=1)
        group_max_class += 1  # shift the class index by 1

        all_group_max_score.append(group_max_score)
        all_group_max_class.append(group_max_class)

    all_group_max_score = torch.stack(all_group_max_score, dim=1)
    all_group_max_class = torch.stack(all_group_max_class, dim=1)

    final_max_score, max_group = torch.max(all_group_max_score, dim=1)

    pred_cls_within_group = all_group_max_class[torch.arange(num_samples),
                                                max_group]

    gt_class, gt_group = torch.max(labels, dim=1)

    selected_groups = (max_group == gt_group)

    pred_acc = torch.zeros(logits.shape[0]).bool().cuda()

    pred_acc[selected_groups] = (
        pred_cls_within_group[selected_groups] == gt_class[selected_groups])

    return loss, pred_acc


def run_eval_acc(model, data_loader, group_slices, num_group):
    # switch to evaluate mode
    model.eval()

    print('Running validation...')

    all_c, all_top1 = [], []

    train_dataiter = iter(data_loader)
    for train_step in tqdm(range(1,
                                 len(train_dataiter) + 1),
                           desc='Test: ',
                           position=0,
                           leave=True,
                           disable=not comm.is_main_process()):
        batch = next(train_dataiter)
        data = batch['data'].cuda()
        group_label = batch['group_label'].cuda()
        class_label = batch['class_label'].cuda()
        labels = []
        for i in range(len(group_label)):
            label = torch.zeros(num_group, dtype=torch.int64)
            label[group_label[i]] = class_label[i] + 1
            labels.append(label.unsqueeze(0))
        labels = torch.cat(labels, dim=0).cuda()

        with torch.no_grad():
            # compute output, measure accuracy and record loss.
            logits = model(data)
            if group_slices is not None:
                c, top1 = calc_group_softmax_acc(logits, labels, group_slices)
            else:
                c = torch.nn.CrossEntropyLoss(reduction='none')(logits, labels)
                top1 = topk(logits, labels, ks=(1, ))[0]

            all_c.extend(c.cpu())  # Also ensures a sync point.
            all_top1.extend(top1.cpu())

    model.train()
    # all_c is val loss
    # all_top1 is val top1 acc
    return all_c, all_top1


class MOSEvaluator(BaseEvaluator):
    def __init__(self, config: Config):
        super(MOSEvaluator, self).__init__(config)
        self.config = config
        self.num_groups = None
        self.group_slices = None
        self.acc = None

    def cal_group_slices(self, train_loader):
        config = self.config
        # if specified group_config
        if (config.trainer.group_config.endswith('npy')):
            classes_per_group = np.load(config.trainer.group_config)
        elif (config.trainer.group_config.endswith('txt')):
            classes_per_group = np.loadtxt(config.trainer.group_config,
                                           dtype=int)
        else:
            # cal group config
            config = self.config
            group = {}
            train_dataiter = iter(train_loader)
            for train_step in tqdm(range(1,
                                         len(train_dataiter) + 1),
                                   desc='cal group_config',
                                   position=0,
                                   leave=True,
                                   disable=not comm.is_main_process()):
                batch = next(train_dataiter)
                group_label = batch['group_label']
                class_label = batch['class_label']

                for i in range(len(class_label)):
                    gl = group_label[i].item()
                    cl = class_label[i].item()

                    try:
                        group[str(gl)]
                    except:
                        group[str(gl)] = []

                    if cl not in group[str(gl)]:
                        group[str(gl)].append(cl)

            classes_per_group = []
            for i in range(len(group)):
                classes_per_group.append(max(group[str(i)]) + 1)

        self.num_groups = len(classes_per_group)
        self.group_slices = get_group_slices(classes_per_group)
        self.group_slices = self.group_slices.cuda()

    def eval_ood(self,
                 net: nn.Module,
                 id_data_loader: DataLoader,
                 ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                 postprocessor=None,
                 fsood=False):
        net.eval()
        if self.group_slices is None or self.num_groups is None:
            self.cal_group_slices(id_data_loader['train'])
        dataset_name = self.config.dataset.name

        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_conf = iterate_data(id_data_loader['test'], net, self.group_slices)
        # dummy pred and gt
        # the accuracy will be handled by self.eval_acc
        id_pred = np.zeros_like(id_conf)
        id_gt = np.zeros_like(id_conf)

        if fsood:
            # load csid data and compute confidence
            for dataset_name, csid_dl in ood_data_loaders['csid'].items():
                print(f'Performing inference on {dataset_name} dataset...',
                      flush=True)
                csid_conf = iterate_data(csid_dl, net, self.group_slices)
                # dummy pred and gt
                # the accuracy will be handled by self.eval_acc
                csid_pred = np.zeros_like(csid_conf)
                csid_gt = np.zeros_like(csid_conf)
                if self.config.recorder.save_scores:
                    self._save_scores(csid_pred, csid_conf, csid_gt,
                                      dataset_name)
                id_pred = np.concatenate([id_pred, csid_pred])
                id_conf = np.concatenate([id_conf, csid_conf])
                id_gt = np.concatenate([id_gt, csid_gt])

        # load nearood data and compute ood metrics
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       ood_split='nearood')
        # load farood data and compute ood metrics
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       ood_split='farood')

    def _eval_ood(self,
                  net: nn.Module,
                  id_list: List[np.ndarray],
                  ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                  ood_split: str = 'nearood'):
        print(f'Processing {ood_split}...', flush=True)
        [id_pred, id_conf, id_gt] = id_list
        metrics_list = []
        for dataset_name, ood_dl in ood_data_loaders[ood_split].items():
            print(f'Performing inference on {dataset_name} dataset...',
                  flush=True)
            ood_conf = iterate_data(ood_dl, net, self.group_slices)
            ood_gt = -1 * np.ones_like(ood_conf)  # hard set to -1 as ood
            # dummy pred
            ood_pred = np.zeros_like(ood_conf)
            if self.config.recorder.save_scores:
                self._save_scores(ood_pred, ood_conf, ood_gt, dataset_name)

            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])

            print(f'Computing metrics on {dataset_name} dataset...')

            ood_metrics = compute_all_metrics(conf, label, pred)
            # the acc here is not reliable
            # since we use dummy pred and gt for id samples
            # so we use the acc computed by self.eval_acc
            ood_metrics[-1] = self.acc

            if self.config.recorder.save_csv:
                self._save_csv(ood_metrics, dataset_name=dataset_name)
            metrics_list.append(ood_metrics)

        print('Computing mean metrics...', flush=True)
        metrics_list = np.array(metrics_list)
        metrics_mean = np.mean(metrics_list, axis=0)
        if self.config.recorder.save_csv:
            self._save_csv(metrics_mean, dataset_name=ood_split)

    def _save_csv(self, metrics, dataset_name):
        [fpr, auroc, aupr_in, aupr_out, accuracy] = metrics

        write_content = {
            'dataset': dataset_name,
            'FPR@95': '{:.2f}'.format(100 * fpr),
            'AUROC': '{:.2f}'.format(100 * auroc),
            'AUPR_IN': '{:.2f}'.format(100 * aupr_in),
            'AUPR_OUT': '{:.2f}'.format(100 * aupr_out),
            'ACC': '{:.2f}'.format(100 * accuracy)
        }

        fieldnames = list(write_content.keys())

        # print ood metric results
        print('FPR@95: {:.2f}, AUROC: {:.2f}'.format(100 * fpr, 100 * auroc),
              end=' ',
              flush=True)
        print('AUPR_IN: {:.2f}, AUPR_OUT: {:.2f}'.format(
            100 * aupr_in, 100 * aupr_out),
              flush=True)
        print('ACC: {:.2f}'.format(accuracy * 100), flush=True)
        print(u'\u2500' * 70, flush=True)

        csv_path = os.path.join(self.config.output_dir, 'ood.csv')
        if not os.path.exists(csv_path):
            with open(csv_path, 'w', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerow(write_content)
        else:
            with open(csv_path, 'a', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writerow(write_content)

    def _save_scores(self, pred, conf, gt, save_name):
        save_dir = os.path.join(self.config.output_dir, 'scores')
        os.makedirs(save_dir, exist_ok=True)
        np.savez(os.path.join(save_dir, save_name),
                 pred=pred,
                 conf=conf,
                 label=gt)

    def eval_acc(self,
                 net: nn.Module,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1,
                 num_groups: int = None,
                 group_slices: torch.Tensor = None,
                 fsood: bool = False,
                 csid_data_loaders: DataLoader = None):
        net.eval()
        if num_groups is None or group_slices is None:
            self.cal_group_slices(data_loader)
        else:
            self.num_groups = num_groups
            self.group_slices = group_slices.cuda()

        loss, top1 = run_eval_acc(net, data_loader, self.group_slices,
                                  self.num_groups)

        if fsood:
            assert csid_data_loaders is not None
            for dataset_name, csid_dl in csid_data_loaders.items():
                _, temp = run_eval_acc(net, csid_dl, self.group_slices,
                                       self.num_groups)
                top1.extend(temp)

        metrics = {}
        metrics['acc'] = np.mean(top1)
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = np.mean(loss)
        self.acc = metrics['acc']

        return metrics

    def report(self, test_metrics):
        print('Completed!', flush=True)

openood/evaluators/mos_evaluator.py
kai422/SCALE
import csv
import os
from typing import Dict, List

import numpy as np
import torch.nn as nn
from torch.utils.data import DataLoader

from openood.postprocessors import BasePostprocessor
from openood.utils import Config

from .base_evaluator import BaseEvaluator
from .metrics import compute_all_metrics


class OODEvaluator(BaseEvaluator):
    def __init__(self, config: Config):
        """OOD Evaluator.

        Args:
            config (Config): Config file from
        """
        super(OODEvaluator, self).__init__(config)
        self.id_pred = None
        self.id_conf = None
        self.id_gt = None

    def eval_ood(self,
                 net: nn.Module,
                 id_data_loaders: Dict[str, DataLoader],
                 ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                 postprocessor: BasePostprocessor,
                 fsood: bool = False):
        if type(net) is dict:
            for subnet in net.values():
                subnet.eval()
        else:
            net.eval()
        assert 'test' in id_data_loaders, \
            'id_data_loaders should have the key: test!'
        dataset_name = self.config.dataset.name

        if self.config.postprocessor.APS_mode:
            assert 'val' in id_data_loaders
            assert 'val' in ood_data_loaders
            self.hyperparam_search(net, id_data_loaders['val'],
                                   ood_data_loaders['val'], postprocessor)

        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_pred, id_conf, id_gt = postprocessor.inference(
            net, id_data_loaders['test'])
        if self.config.recorder.save_scores:
            self._save_scores(id_pred, id_conf, id_gt, dataset_name)

        if fsood:
            # load csid data and compute confidence
            for dataset_name, csid_dl in ood_data_loaders['csid'].items():
                print(f'Performing inference on {dataset_name} dataset...',
                      flush=True)
                csid_pred, csid_conf, csid_gt = postprocessor.inference(
                    net, csid_dl)
                if self.config.recorder.save_scores:
                    self._save_scores(csid_pred, csid_conf, csid_gt,
                                      dataset_name)
                id_pred = np.concatenate([id_pred, csid_pred])
                id_conf = np.concatenate([id_conf, csid_conf])
                id_gt = np.concatenate([id_gt, csid_gt])

        # load nearood data and compute ood metrics
        print(u'\u2500' * 70, flush=True)
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='nearood')

        # load farood data and compute ood metrics
        print(u'\u2500' * 70, flush=True)
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='farood')

    def _eval_ood(self,
                  net: nn.Module,
                  id_list: List[np.ndarray],
                  ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                  postprocessor: BasePostprocessor,
                  ood_split: str = 'nearood'):
        print(f'Processing {ood_split}...', flush=True)
        [id_pred, id_conf, id_gt] = id_list
        metrics_list = []
        for dataset_name, ood_dl in ood_data_loaders[ood_split].items():
            print(f'Performing inference on {dataset_name} dataset...',
                  flush=True)
            ood_pred, ood_conf, ood_gt = postprocessor.inference(net, ood_dl)
            ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood
            if self.config.recorder.save_scores:
                self._save_scores(ood_pred, ood_conf, ood_gt, dataset_name)

            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])

            print(f'Computing metrics on {dataset_name} dataset...')

            ood_metrics = compute_all_metrics(conf, label, pred)
            if self.config.recorder.save_csv:
                self._save_csv(ood_metrics, dataset_name=dataset_name)
            metrics_list.append(ood_metrics)

        print('Computing mean metrics...', flush=True)
        metrics_list = np.array(metrics_list)
        metrics_mean = np.mean(metrics_list, axis=0)
        if self.config.recorder.save_csv:
            self._save_csv(metrics_mean, dataset_name=ood_split)

    def eval_ood_val(self, net: nn.Module, id_data_loaders: Dict[str,
                                                                 DataLoader],
                     ood_data_loaders: Dict[str, DataLoader],
                     postprocessor: BasePostprocessor):
        if type(net) is dict:
            for subnet in net.values():
                subnet.eval()
        else:
            net.eval()
        assert 'val' in id_data_loaders
        assert 'val' in ood_data_loaders
        if self.config.postprocessor.APS_mode:
            val_auroc = self.hyperparam_search(net, id_data_loaders['val'],
                                               ood_data_loaders['val'],
                                               postprocessor)
        else:
            id_pred, id_conf, id_gt = postprocessor.inference(
                net, id_data_loaders['val'])
            ood_pred, ood_conf, ood_gt = postprocessor.inference(
                net, ood_data_loaders['val'])
            ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood
            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])
            ood_metrics = compute_all_metrics(conf, label, pred)
            val_auroc = ood_metrics[1]
        return {'auroc': 100 * val_auroc}

    def _save_csv(self, metrics, dataset_name):
        [fpr, auroc, aupr_in, aupr_out, accuracy] = metrics

        write_content = {
            'dataset': dataset_name,
            'FPR@95': '{:.2f}'.format(100 * fpr),
            'AUROC': '{:.2f}'.format(100 * auroc),
            'AUPR_IN': '{:.2f}'.format(100 * aupr_in),
            'AUPR_OUT': '{:.2f}'.format(100 * aupr_out),
            'ACC': '{:.2f}'.format(100 * accuracy)
        }

        fieldnames = list(write_content.keys())

        # print ood metric results
        print('FPR@95: {:.2f}, AUROC: {:.2f}'.format(100 * fpr, 100 * auroc),
              end=' ',
              flush=True)
        print('AUPR_IN: {:.2f}, AUPR_OUT: {:.2f}'.format(
            100 * aupr_in, 100 * aupr_out),
              flush=True)
        print('ACC: {:.2f}'.format(accuracy * 100), flush=True)
        print(u'\u2500' * 70, flush=True)

        csv_path = os.path.join(self.config.output_dir, 'ood.csv')
        if not os.path.exists(csv_path):
            with open(csv_path, 'w', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerow(write_content)
        else:
            with open(csv_path, 'a', newline='') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writerow(write_content)

    def _save_scores(self, pred, conf, gt, save_name):
        save_dir = os.path.join(self.config.output_dir, 'scores')
        os.makedirs(save_dir, exist_ok=True)
        np.savez(os.path.join(save_dir, save_name),
                 pred=pred,
                 conf=conf,
                 label=gt)

    def eval_acc(self,
                 net: nn.Module,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1,
                 fsood: bool = False,
                 csid_data_loaders: DataLoader = None):
        """Returns the accuracy score of the labels and predictions.

        :return: float
        """
        if type(net) is dict:
            net['backbone'].eval()
        else:
            net.eval()
        self.id_pred, self.id_conf, self.id_gt = postprocessor.inference(
            net, data_loader)

        if fsood:
            assert csid_data_loaders is not None
            for dataset_name, csid_dl in csid_data_loaders.items():
                csid_pred, csid_conf, csid_gt = postprocessor.inference(
                    net, csid_dl)
                self.id_pred = np.concatenate([self.id_pred, csid_pred])
                self.id_conf = np.concatenate([self.id_conf, csid_conf])
                self.id_gt = np.concatenate([self.id_gt, csid_gt])

        metrics = {}
        metrics['acc'] = sum(self.id_pred == self.id_gt) / len(self.id_pred)
        metrics['epoch_idx'] = epoch_idx
        return metrics

    def report(self, test_metrics):
        print('Completed!', flush=True)

    def hyperparam_search(
        self,
        net: nn.Module,
        id_data_loader,
        ood_data_loader,
        postprocessor: BasePostprocessor,
    ):
        print('Starting automatic parameter search...')
        aps_dict = {}
        max_auroc = 0
        hyperparam_names = []
        hyperparam_list = []
        count = 0
        for name in postprocessor.args_dict.keys():
            hyperparam_names.append(name)
            count += 1
        for name in hyperparam_names:
            hyperparam_list.append(postprocessor.args_dict[name])
        hyperparam_combination = self.recursive_generator(
            hyperparam_list, count)
        for hyperparam in hyperparam_combination:
            postprocessor.set_hyperparam(hyperparam)
            id_pred, id_conf, id_gt = postprocessor.inference(
                net, id_data_loader)
            ood_pred, ood_conf, ood_gt = postprocessor.inference(
                net, ood_data_loader)
            ood_gt = -1 * np.ones_like(ood_gt)  # hard set to -1 as ood
            pred = np.concatenate([id_pred, ood_pred])
            conf = np.concatenate([id_conf, ood_conf])
            label = np.concatenate([id_gt, ood_gt])
            ood_metrics = compute_all_metrics(conf, label, pred)
            index = hyperparam_combination.index(hyperparam)
            aps_dict[index] = ood_metrics[1]
            print('Hyperparam:{}, auroc:{}'.format(hyperparam,
                                                   aps_dict[index]))
            if ood_metrics[1] > max_auroc:
                max_auroc = ood_metrics[1]
        for key in aps_dict.keys():
            if aps_dict[key] == max_auroc:
                postprocessor.set_hyperparam(hyperparam_combination[key])
        print('Final hyperparam: {}'.format(postprocessor.get_hyperparam()))
        return max_auroc

    def recursive_generator(self, list, n):
        if n == 1:
            results = []
            for x in list[0]:
                k = []
                k.append(x)
                results.append(k)
            return results
        else:
            results = []
            temp = self.recursive_generator(list, n - 1)
            for x in list[n - 1]:
                for y in temp:
                    k = y.copy()
                    k.append(x)
                    results.append(k)
            return results

openood/evaluators/ood_evaluator.py
kai422/SCALE
from typing import Dict

import torch.nn as nn
from torch.utils.data import DataLoader

from openood.postprocessors import BasePostprocessor
from openood.utils import Config

from .ood_evaluator import OODEvaluator


class OSREvaluator(OODEvaluator):
    def __init__(self, config: Config):
        super(OSREvaluator, self).__init__(config)

    def eval_ood(self, net: nn.Module, id_data_loader: DataLoader,
                 ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                 postprocessor: BasePostprocessor):
        if type(net) is dict:
            for subnet in net.values():
                subnet.eval()
        else:
            net.eval()

        # load training in-distribution data
        assert 'test' in id_data_loader, \
            'id_data_loaders should have the key: test!'
        dataset_name = self.config.dataset.name
        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_pred, id_conf, id_gt = postprocessor.inference(
            net, id_data_loader['test'])
        if self.config.recorder.save_scores:
            self._save_scores(id_pred, id_conf, id_gt, dataset_name)

        # load nearood data and compute ood metrics
        self._eval_ood(net, [id_pred, id_conf, id_gt],
                       ood_data_loaders,
                       postprocessor,
                       ood_split='osr')

openood/evaluators/osr_evaluator.py
kai422/SCALE
import os
from typing import Dict

import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from scipy.ndimage import gaussian_filter
from sklearn.metrics import roc_auc_score
from torch.utils.data import DataLoader
from torchvision import transforms

from openood.postprocessors import BasePostprocessor
from openood.utils import Config

from .base_evaluator import BaseEvaluator


class PatchCoreEvaluator(BaseEvaluator):
    def __init__(self, config: Config):
        super(PatchCoreEvaluator, self).__init__(config)
        self.config = config

    def eval_ood(self, net: nn.Module, id_data_loader: DataLoader,
                 ood_data_loaders: Dict[str, Dict[str, DataLoader]],
                 postprocessor: BasePostprocessor):
        net.eval()

        dataset_name = self.config.dataset.name
        print(f'Performing inference on {dataset_name} dataset...', flush=True)
        id_pred, id_conf, id_gt = postprocessor.inference(
            net, ood_data_loaders['val'])  # not good
        good_pred, good_conf, good_gt = postprocessor.inference(
            net, id_data_loader['test'])  # good

        # pred = np.concatenate([id_pred, good_pred])
        conf = np.concatenate([id_conf, good_conf])
        gt = np.concatenate([id_gt, good_gt])

        self.gt_transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.CenterCrop(224)
        ])
        mean_train = [0.485, 0.456, 0.406]
        std_train = [0.229, 0.224, 0.225]
        self.transform = transforms.Compose([
            transforms.Resize((256, 256), Image.ANTIALIAS),
            transforms.ToTensor(),
            transforms.CenterCrop(224),
            transforms.Normalize(mean=mean_train, std=std_train)
        ])
        count = 0
        self.gt_list_px_lvl = []

        for batch in id_data_loader['trainGT']:
            #data = batch['data'].cuda()
            data = []
            label = batch['label'].cuda()
            name = batch['image_name']
            for i in name:
                path = os.path.join('./data/images/', i)
                gt_img = Image.open(path)
                gt_img = self.gt_transform(gt_img)
                gt_img = torch.unsqueeze(gt_img, 0)

                # gt_img = self.gt_transform(gt_img)
                gt_np = gt_img.cpu().numpy()[0, 0].astype(int)
                self.gt_list_px_lvl.extend(gt_np.ravel())

        self.pred_list_px_lvl = []
        self.pred_list_img_lvl = []

        for patchscore in conf:

            anomaly_map = patchscore[:, 0].reshape((28, 28))
            N_b = patchscore[np.argmax(patchscore[:, 0])]
            w = (1 - (np.max(np.exp(N_b)) / np.sum(np.exp(N_b))))
            score = w * max(patchscore[:, 0])  # Image-level score

            anomaly_map_resized = cv2.resize(anomaly_map, (224, 224))
            anomaly_map_resized_blur = gaussian_filter(anomaly_map_resized,
                                                       sigma=4)
            self.pred_list_px_lvl.extend(anomaly_map_resized_blur.ravel())
            self.pred_list_img_lvl.append(score)

        print('Total image-level auc-roc score :')
        img_auc = roc_auc_score(gt, self.pred_list_img_lvl)
        print(img_auc)

        if (test_pix):
            print('Total pixel-level auc-roc score :')
            pixel_auc = roc_auc_score(self.gt_list_px_lvl,
                                      self.pred_list_px_lvl)
            print(pixel_auc)

    def eval_acc(self,
                 net: nn.Module,
                 data_loader: DataLoader,
                 postprocessor: BasePostprocessor = None,
                 epoch_idx: int = -1):
        net.eval()
        id_pred, _, id_gt = postprocessor.inference(net, data_loader)
        metrics = {}
        metrics['acc'] = sum(id_pred == id_gt) / len(id_pred)
        metrics['epoch_idx'] = epoch_idx
        return metrics

    def report(self, test_metrics):
        print('Completed!', flush=True)

openood/evaluators/patchcore_evaluator.py
kai422/SCALE
from openood.evaluators.mos_evaluator import MOSEvaluator
from openood.utils import Config

from .ad_evaluator import ADEvaluator
from .arpl_evaluator import ARPLEvaluator
from .base_evaluator import BaseEvaluator
from .ece_evaluator import ECEEvaluator
from .fsood_evaluator import FSOODEvaluator
from .ood_evaluator import OODEvaluator
from .osr_evaluator import OSREvaluator
from .patchcore_evaluator import PatchCoreEvaluator


def get_evaluator(config: Config):
    evaluators = {
        'base': BaseEvaluator,
        'ood': OODEvaluator,
        'fsood': FSOODEvaluator,
        'patch': PatchCoreEvaluator,
        'arpl': ARPLEvaluator,
        'ad': ADEvaluator,
        'mos': MOSEvaluator,
        'ece': ECEEvaluator,
        'osr': OSREvaluator
    }
    return evaluators[config.evaluator.name](config)

openood/evaluators/utils.py
kai422/SCALE
from .scale_net import ScaleNet
from .resnet50 import ResNet50
from .utils import get_network

openood/networks/__init__.py
kai422/SCALE
from torchvision.models.resnet import Bottleneck, ResNet


class SaveFeatures():
    def __init__(self, module):
        self.hook = module.register_forward_hook(self.hook_fn)

    def hook_fn(self, module, input, output):
        # self.input = input.clone()
        self.output = output.clone()

    def close(self):
        self.hook.remove()


class ResNet50(ResNet):
    def __init__(self,
                 block=Bottleneck,
                 layers=[3, 4, 6, 3],
                 num_classes=1000):
        super(ResNet50, self).__init__(block=block,
                                       layers=layers,
                                       num_classes=num_classes)
        self.feature_size = 2048
        self._prerelu = SaveFeatures(self.layer4[2].bn3)

    def forward(self, x, return_feature=False, return_feature_list=False, return_prerelu_feature=False):
        if return_prerelu_feature:
            feature1 = self.relu(self.bn1(self.conv1(x)))
            feature1 = self.maxpool(feature1)
            feature2, feature2_prerelu = self.layer1([feature1])
            feature3, feature3_prerelu = self.layer2([feature2])
            feature4, feature4_prerelu = self.layer3([feature3])
            feature5, feature5_prerelu = self.layer4([feature4])
            feature5 = self.avgpool(feature5)
            feature = feature5.view(feature5.size(0), -1)
            return feature,  self.avgpool(feature5_prerelu).view(feature5.size(0), -1)

        feature1 = self.relu(self.bn1(self.conv1(x)))
        feature1 = self.maxpool(feature1)
        feature2 = self.layer1(feature1)
        feature3 = self.layer2(feature2)
        feature4 = self.layer3(feature3)
        feature5 = self.layer4(feature4)
        feature5 = self.avgpool(feature5)
        feature = feature5.view(feature5.size(0), -1)
        logits_cls = self.fc(feature)

        feature_list = [feature1, feature2, feature3, feature4, feature5]
        if return_feature:
            return logits_cls, feature
        elif return_feature_list:
            return logits_cls, feature_list
        else:
            return logits_cls

    def forward_threshold(self, x, threshold):
        feature1 = self.relu(self.bn1(self.conv1(x)))
        feature1 = self.maxpool(feature1)
        feature2 = self.layer1(feature1)
        feature3 = self.layer2(feature2)
        feature4 = self.layer3(feature3)
        feature5 = self.layer4(feature4)
        feature5 = self.avgpool(feature5)
        feature = feature5.clip(max=threshold)
        feature = feature.view(feature.size(0), -1)
        logits_cls = self.fc(feature)

        return logits_cls

    def intermediate_forward(self, x, layer_index):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.maxpool(out)

        out = self.layer1(out)
        if layer_index == 1:
            return out

        out = self.layer2(out)
        if layer_index == 2:
            return out

        out = self.layer3(out)
        if layer_index == 3:
            return out

        out = self.layer4(out)
        if layer_index == 4:
            return out

        raise ValueError

    def get_fc(self):
        fc = self.fc
        return fc.weight.cpu().detach().numpy(), fc.bias.cpu().detach().numpy()

    def get_fc_layer(self):
        return self.fc

openood/networks/resnet50.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import seaborn as sns


class ScaleNet(nn.Module):
    def __init__(self, backbone):
        super(ScaleNet, self).__init__()
        self.backbone = backbone

    def forward(self, x, return_feature=False, return_feature_list=False):
        try:
            return self.backbone(x, return_feature, return_feature_list)
        except TypeError:
            return self.backbone(x, return_feature)

    def forward_threshold(self, x, percentile):
        _, feature = self.backbone(x, return_feature=True)
        feature = scale(feature.view(feature.size(0), -1, 1, 1), percentile)
        feature = feature.view(feature.size(0), -1)
        logits_cls = self.backbone.get_fc_layer()(feature)
        return logits_cls

    def get_fc(self):
        fc = self.backbone.fc
        return fc.weight.cpu().detach().numpy(), fc.bias.cpu().detach().numpy()

def scale(x, percentile=65):
    input = x.clone()
    assert x.dim() == 4
    assert 0 <= percentile <= 100
    b, c, h, w = x.shape

    # calculate the sum of the input per sample
    s1 = x.sum(dim=[1, 2, 3])
    n = x.shape[1:].numel()
    k = n - int(np.round(n * percentile / 100.0))
    t = x.view((b, c * h * w))
    v, i = torch.topk(t, k, dim=1)
    t.zero_().scatter_(dim=1, index=i, src=v)

    # calculate new sum of the input per sample after pruning
    s2 = x.sum(dim=[1, 2, 3])

    # apply sharpening
    scale = s1 / s2
    
    return input * torch.exp(scale[:, None, None, None])

openood/networks/scale_net.py
kai422/SCALE
# import mmcv
from copy import deepcopy
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
# from mmcls.apis import init_model

import openood.utils.comm as comm

from .resnet50 import ResNet50


def get_network(network_config):

    num_classes = network_config.num_classes

    if network_config.name == 'resnet50':
        net = ResNet50(num_classes=num_classes)

    if network_config.pretrained:
        if type(net) is dict:
            if isinstance(network_config.checkpoint, list):
                for subnet, checkpoint in zip(net.values(),
                                              network_config.checkpoint):
                    if checkpoint is not None:
                        if checkpoint != 'none':
                            subnet.load_state_dict(torch.load(checkpoint),
                                                   strict=False)
            elif isinstance(network_config.checkpoint, str):
                ckpt = torch.load(network_config.checkpoint)
                subnet_ckpts = {k: {} for k in net.keys()}
                for k, v in ckpt.items():
                    for subnet_name in net.keys():
                        if k.startwith(subnet_name):
                            subnet_ckpts[subnet_name][k.replace(
                                subnet_name + '.', '')] = v
                            break

                for subnet_name, subnet in net.items():
                    subnet.load_state_dict(subnet_ckpts[subnet_name])

        elif network_config.name == 'bit' and not network_config.normal_load:
            net.load_from(np.load(network_config.checkpoint))
        elif network_config.name == 'vit':
            pass
        else:
            try:
                net.load_state_dict(torch.load(network_config.checkpoint),
                                    strict=False)
            except RuntimeError:
                # sometimes fc should not be loaded
                loaded_pth = torch.load(network_config.checkpoint)
                loaded_pth.pop('fc.weight')
                loaded_pth.pop('fc.bias')
                net.load_state_dict(loaded_pth, strict=False)
        print('Model Loading {} Completed!'.format(network_config.name))

    if network_config.num_gpus > 1:
        if type(net) is dict:
            for key, subnet in zip(net.keys(), net.values()):
                net[key] = torch.nn.parallel.DistributedDataParallel(
                    subnet.cuda(),
                    device_ids=[comm.get_local_rank()],
                    broadcast_buffers=True)
        else:
            net = torch.nn.parallel.DistributedDataParallel(
                net.cuda(),
                device_ids=[comm.get_local_rank()],
                broadcast_buffers=True)

    if network_config.num_gpus > 0:
        if type(net) is dict:
            for subnet in net.values():
                subnet.cuda()
        else:
            net.cuda()

    cudnn.benchmark = True
    return net

openood/networks/utils.py
kai422/SCALE
from .utils import get_pipeline

openood/pipelines/__init__.py
kai422/SCALE
from openood.datasets import get_dataloader, get_ood_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.utils import setup_logger


class FeatExtractOpenGANPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        id_loader_dict = get_dataloader(self.config)
        ood_loader_dict = get_ood_dataloader(self.config)
        assert 'train' in id_loader_dict
        assert 'val' in id_loader_dict
        assert 'val' in ood_loader_dict

        # init network
        net = get_network(self.config.network)

        # init evaluator
        evaluator = get_evaluator(self.config)

        # sanity check on id val accuracy
        print('\nStart evaluation on ID val data...', flush=True)
        test_metrics = evaluator.eval_acc(net, id_loader_dict['val'])
        print('\nComplete Evaluation, accuracy {:.2f}%'.format(
            100 * test_metrics['acc']),
              flush=True)

        # start extracting features
        print('\nStart Feature Extraction...', flush=True)
        print('\t ID training data...')
        evaluator.extract(net, id_loader_dict['train'], 'id_train')

        print('\t ID val data...')
        evaluator.extract(net, id_loader_dict['val'], 'id_val')

        print('\t OOD val data...')
        evaluator.extract(net, ood_loader_dict['val'], 'ood_val')
        print('\nComplete Feature Extraction!')

openood/pipelines/feat_extract_opengan_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.utils import setup_logger


class FeatExtractPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        test_loader = loader_dict[self.config.pipeline.extract_target]

        # init network
        net = get_network(self.config.network)

        # init evaluator
        evaluator = get_evaluator(self.config)

        # start calculating accuracy
        print('\nStart evaluation...', flush=True)
        test_metrics = evaluator.eval_acc(net, test_loader)
        print('\nComplete Evaluation, accuracy {:.2f}%'.format(
            100 * test_metrics['acc']),
              flush=True)

        # start extracting features
        print('\nStart Feature Extraction...', flush=True)
        evaluator.extract(net, test_loader)
        print('\nComplete Feature Extraction!')

openood/pipelines/feat_extract_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class FinetunePipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        train_loader, val_loader = loader_dict['train'], loader_dict['val']
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, train_loader, self.config)
        evaluator = get_evaluator(self.config)

        # init recorder
        recorder = get_recorder(self.config)

        # trainer setup
        trainer.setup()
        print('\n' + u'\u2500' * 70, flush=True)

        print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train and eval the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            val_metrics = evaluator.eval_acc(net, val_loader, None, epoch_idx)
            # save model and report the result
            recorder.save_model(net, val_metrics)
            recorder.report(train_metrics, val_metrics)
        recorder.summary()
        print(u'\u2500' * 70, flush=True)

        # evaluate on test set
        print('Start testing...', flush=True)
        test_metrics = evaluator.eval_acc(net, test_loader)
        print('\nComplete Evaluation, accuracy {:.2f}'.format(
            100.0 * test_metrics['acc']),
              flush=True)
        print('Completed!', flush=True)

openood/pipelines/finetune_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.utils import setup_logger


class TestAccPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init evaluator
        evaluator = get_evaluator(self.config)

        # start calculating accuracy
        print('\nStart evaluation...', flush=True)
        test_metrics = evaluator.eval_acc(net, test_loader)
        print('\nComplete Evaluation, accuracy {:.2f}%'.format(
            100 * test_metrics['acc']),
              flush=True)

openood/pipelines/test_acc_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader, get_ood_dataloader
from openood.evaluators.utils import get_evaluator
from openood.networks.utils import get_network
from openood.postprocessors import get_postprocessor
from openood.utils import setup_logger


class TestAdPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        id_loader_dict = get_dataloader(self.config)
        ood_loader_dict = get_ood_dataloader(self.config)

        # init network
        net = get_network(self.config.network)

        # init evaluator
        evaluator = get_evaluator(self.config)

        postprocessor = get_postprocessor(self.config)
        # setup for distance-based methods
        postprocessor.setup(net, id_loader_dict, ood_loader_dict)

        print('Start testing...', flush=True)
        test_metrics = evaluator.eval_ood(net, id_loader_dict, ood_loader_dict,
                                          postprocessor)
        evaluator.report(test_metrics)

openood/pipelines/test_ad_pipeline.py
kai422/SCALE
import time

from openood.datasets import get_dataloader, get_ood_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.postprocessors import get_postprocessor
from openood.utils import setup_logger


class TestOODPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        id_loader_dict = get_dataloader(self.config)
        ood_loader_dict = get_ood_dataloader(self.config)

        # init network
        net = get_network(self.config.network)

        # init ood evaluator
        evaluator = get_evaluator(self.config)

        # init ood postprocessor
        postprocessor = get_postprocessor(self.config)
        # setup for distance-based methods
        postprocessor.setup(net, id_loader_dict, ood_loader_dict)
        print('\n', flush=True)
        print(u'\u2500' * 70, flush=True)

        # start calculating accuracy
        print('\nStart evaluation...', flush=True)
        if self.config.evaluator.ood_scheme == 'fsood':
            acc_metrics = evaluator.eval_acc(
                net,
                id_loader_dict['test'],
                postprocessor,
                fsood=True,
                csid_data_loaders=ood_loader_dict['csid'])
        else:
            acc_metrics = evaluator.eval_acc(net, id_loader_dict['test'],
                                             postprocessor)
        print('\nAccuracy {:.2f}%'.format(100 * acc_metrics['acc']),
              flush=True)
        print(u'\u2500' * 70, flush=True)

        # start evaluating ood detection methods
        timer = time.time()
        if self.config.evaluator.ood_scheme == 'fsood':
            evaluator.eval_ood(net,
                               id_loader_dict,
                               ood_loader_dict,
                               postprocessor,
                               fsood=True)
        else:
            evaluator.eval_ood(net, id_loader_dict, ood_loader_dict,
                               postprocessor)
        print('Time used for eval_ood: {:.0f}s'.format(time.time() - timer))
        print('Completed!', flush=True)

openood/pipelines/test_ood_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader, get_ood_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.postprocessors import get_postprocessor
from openood.utils import setup_logger


class TestOODPipelineAPS:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        id_loader_dict = get_dataloader(self.config)
        ood_loader_dict = get_ood_dataloader(self.config)

        # init network
        net = get_network(self.config.network)

        # init ood evaluator
        evaluator = get_evaluator(self.config)

        # init ood postprocessor
        postprocessor = get_postprocessor(self.config)
        # setup for distance-based methods
        postprocessor.setup(net, id_loader_dict, ood_loader_dict)
        print('\n', flush=True)
        print(u'\u2500' * 70, flush=True)

        # start calculating accuracy
        print('\nStart evaluation...', flush=True)
        acc_metrics = evaluator.eval_acc(net, id_loader_dict['test'],
                                         postprocessor)
        print('\nAccuracy {:.2f}%'.format(100 * acc_metrics['acc']),
              flush=True)
        print(u'\u2500' * 70, flush=True)

        # start evaluating ood detection methods
        evaluator.eval_ood(net, id_loader_dict, ood_loader_dict, postprocessor)
        print('Completed!', flush=True)

openood/pipelines/test_ood_pipeline_aps.py
kai422/SCALE
from openood.datasets import get_dataloader, get_ood_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.postprocessors import get_postprocessor
from openood.preprocessors.utils import get_preprocessor
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainAdPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        id_loader_dict = get_dataloader(self.config)
        ood_loader_dict = get_ood_dataloader(self.config)
        train_loader = id_loader_dict['train']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, train_loader, self.config)
        evaluator = get_evaluator(self.config)

        postprocessor = get_postprocessor(self.config)
        # setup for distance-based methods
        postprocessor.setup(net, id_loader_dict, ood_loader_dict)

        # init recorder
        recorder = get_recorder(self.config)

        print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            test_metrics = evaluator.eval_ood(net,
                                              id_loader_dict,
                                              ood_loader_dict,
                                              postprocessor=postprocessor,
                                              epoch_idx=epoch_idx)
            # save model and report the result
            recorder.save_model(net, test_metrics)
            recorder.report(train_metrics, test_metrics)
        recorder.summary()

        # evaluate on test set
        print('Start testing...', flush=True)
        test_metrics = evaluator.eval_ood(net,
                                          id_loader_dict,
                                          ood_loader_dict,
                                          postprocessor=postprocessor)
        evaluator.report(test_metrics)

openood/pipelines/train_ad_pipeline.py
kai422/SCALE
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainARPLGANPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        train_loader, val_loader = loader_dict['train'], loader_dict['val']
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, train_loader, self.config)
        self.config.trainer.name = 'arpl'
        trainer_aux = get_trainer(net, train_loader, self.config)
        evaluator = get_evaluator(self.config)

        # init recorder
        recorder = get_recorder(self.config)

        print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train and eval the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            net, train_aux_metrics = trainer_aux.train_epoch(epoch_idx)
            train_metrics['loss'] = train_aux_metrics['loss']
            val_metrics = evaluator.eval_acc(net, val_loader, None, epoch_idx)
            trainer.scheduler.step()
            # save model and report the result
            recorder.save_model(net, val_metrics)
            recorder.report(train_metrics, val_metrics)
        recorder.summary()
        print(u'\u2500' * 70, flush=True)

        # evaluate on test set
        print('Start testing...', flush=True)
        test_metrics = evaluator.eval_acc(net, trainer.criterion, test_loader)
        print('\nComplete Evaluation, Last accuracy {:.2f}'.format(
            100.0 * test_metrics['acc']),
              flush=True)
        print('Completed!', flush=True)

openood/pipelines/train_aux_pipeline.py
kai422/SCALE
import openood.utils.comm as comm
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        train_loader, val_loader = loader_dict['train'], loader_dict['val']
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, train_loader, self.config)
        evaluator = get_evaluator(self.config)

        if comm.is_main_process():
            # init recorder
            recorder = get_recorder(self.config)

            print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train and eval the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            val_metrics = evaluator.eval_acc(net, val_loader, None, epoch_idx)
            comm.synchronize()
            if comm.is_main_process():
                # save model and report the result
                recorder.save_model(net, val_metrics)
                recorder.report(train_metrics, val_metrics)

        if comm.is_main_process():
            recorder.summary()
            print(u'\u2500' * 70, flush=True)

            # evaluate on test set
            print('Start testing...', flush=True)

        test_metrics = evaluator.eval_acc(net, test_loader)

        if comm.is_main_process():
            print('\nComplete Evaluation, Last accuracy {:.2f}'.format(
                100.0 * test_metrics['acc']),
                  flush=True)
            print('Completed!', flush=True)

openood/pipelines/train_ddt_pipeline.py
kai422/SCALE
import numpy as np
import torch

import openood.utils.comm as comm
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainOEPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # set random seed
        torch.manual_seed(self.config.seed)
        np.random.seed(self.config.seed)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        train_loader, val_loader = loader_dict['train'], loader_dict['val']
        train_oe_loader = loader_dict['oe']
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, [train_loader, train_oe_loader], None,
                              self.config)
        evaluator = get_evaluator(self.config)

        if comm.is_main_process():
            # init recorder
            recorder = get_recorder(self.config)

            print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train and eval the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            val_metrics = evaluator.eval_acc(net, val_loader, None, epoch_idx)
            comm.synchronize()
            if comm.is_main_process():
                # save model and report the result
                recorder.save_model(net, val_metrics)
                recorder.report(train_metrics, val_metrics)

        if comm.is_main_process():
            recorder.summary()
            print(u'\u2500' * 70, flush=True)

            # evaluate on test set
            print('Start testing...', flush=True)

        test_metrics = evaluator.eval_acc(net, test_loader)

        if comm.is_main_process():
            print('\nComplete Evaluation, Last accuracy {:.2f}'.format(
                100.0 * test_metrics['acc']),
                  flush=True)
            print('Completed!', flush=True)

openood/pipelines/train_oe_pipeline.py
kai422/SCALE
from openood.datasets import get_feature_dataloader
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainOpenGanPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # get dataloader
        feat_loader = get_feature_dataloader(self.config.dataset)

        # init network
        net = get_network(self.config.network)

        # init trainer
        trainer = get_trainer(net, feat_loader, self.config)

        # init recorder
        recorder = get_recorder(self.config)

        print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            recorder.save_model(net, train_metrics)
            recorder.report(train_metrics)
        recorder.summary()

        print('Completed!', flush=True)

openood/pipelines/train_only_pipeline.py
kai422/SCALE
import numpy as np
import torch

from openood.datasets import get_feature_opengan_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.postprocessors import get_postprocessor
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainOpenGanPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # set random seed
        torch.manual_seed(self.config.seed)
        np.random.seed(self.config.seed)

        # get dataloader
        dataloaders = get_feature_opengan_dataloader(self.config.dataset)
        id_loaders = {
            'train': dataloaders['id_train'],
            'val': dataloaders['id_val']
        }  # just for consistency with evaluator
        ood_loaders = {'val': dataloaders['ood_val']}

        # init network
        net = get_network(self.config.network)

        # init trainer
        trainer = get_trainer(net, dataloaders['id_train'],
                              dataloaders['id_val'], self.config)
        evaluator = get_evaluator(self.config)

        # init recorder
        recorder = get_recorder(self.config)

        # init ood postprocessor
        postprocessor = get_postprocessor(self.config)

        print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train the model
            net, train_metrics = trainer.train_epoch(epoch_idx)
            val_metrics = evaluator.eval_ood_val(net, id_loaders, ood_loaders,
                                                 postprocessor)
            val_metrics['epoch_idx'] = train_metrics['epoch_idx']
            recorder.save_model(net, val_metrics)
            recorder.report(train_metrics, val_metrics)
        recorder.summary()

        print('Completed!', flush=True)

openood/pipelines/train_opengan_pipeline.py
kai422/SCALE
import numpy as np
import torch

import openood.utils.comm as comm
from openood.datasets import get_dataloader
from openood.evaluators import get_evaluator
from openood.networks import get_network
from openood.recorders import get_recorder
from openood.trainers import get_trainer
from openood.utils import setup_logger


class TrainPipeline:
    def __init__(self, config) -> None:
        self.config = config

    def run(self):
        # generate output directory and save the full config file
        setup_logger(self.config)

        # set random seed
        torch.manual_seed(self.config.seed)
        np.random.seed(self.config.seed)

        # get dataloader
        loader_dict = get_dataloader(self.config)
        train_loader, val_loader = loader_dict['train'], loader_dict['val']
        test_loader = loader_dict['test']

        # init network
        net = get_network(self.config.network)

        # init trainer and evaluator
        trainer = get_trainer(net, train_loader, val_loader, self.config)
        evaluator = get_evaluator(self.config)

        if comm.is_main_process():
            # init recorder
            recorder = get_recorder(self.config)

            print('Start training...', flush=True)
        for epoch_idx in range(1, self.config.optimizer.num_epochs + 1):
            # train and eval the model
            if self.config.trainer.name == 'mos':
                net, train_metrics, num_groups, group_slices = \
                    trainer.train_epoch(epoch_idx)
                val_metrics = evaluator.eval_acc(net,
                                                 val_loader,
                                                 train_loader,
                                                 epoch_idx,
                                                 num_groups=num_groups,
                                                 group_slices=group_slices)
            elif self.config.trainer.name in ['cider', 'npos']:
                net, train_metrics = trainer.train_epoch(epoch_idx)
                # cider and npos only trains the backbone
                # cannot evaluate ID acc without training the fc layer
                val_metrics = train_metrics
            else:
                net, train_metrics = trainer.train_epoch(epoch_idx)
                val_metrics = evaluator.eval_acc(net, val_loader, None,
                                                 epoch_idx)
            comm.synchronize()
            if comm.is_main_process():
                # save model and report the result
                recorder.save_model(net, val_metrics)
                recorder.report(train_metrics, val_metrics)

        if comm.is_main_process():
            recorder.summary()
            print(u'\u2500' * 70, flush=True)

            # evaluate on test set
            print('Start testing...', flush=True)

        test_metrics = evaluator.eval_acc(net, test_loader)

        if comm.is_main_process():
            print('\nComplete Evaluation, Last accuracy {:.2f}'.format(
                100.0 * test_metrics['acc']),
                  flush=True)
            print('Completed!', flush=True)

openood/pipelines/train_pipeline.py
kai422/SCALE
from openood.utils import Config

from .feat_extract_pipeline import FeatExtractPipeline
from .feat_extract_opengan_pipeline import FeatExtractOpenGANPipeline
from .finetune_pipeline import FinetunePipeline
from .test_acc_pipeline import TestAccPipeline
from .test_ad_pipeline import TestAdPipeline
from .test_ood_pipeline import TestOODPipeline
from .train_ad_pipeline import TrainAdPipeline
from .train_aux_pipeline import TrainARPLGANPipeline
from .train_oe_pipeline import TrainOEPipeline
# from .train_only_pipeline import TrainOpenGanPipeline
from .train_opengan_pipeline import TrainOpenGanPipeline
from .train_pipeline import TrainPipeline
from .test_ood_pipeline_aps import TestOODPipelineAPS


def get_pipeline(config: Config):
    pipelines = {
        'train': TrainPipeline,
        'finetune': FinetunePipeline,
        'test_acc': TestAccPipeline,
        'feat_extract': FeatExtractPipeline,
        'feat_extract_opengan': FeatExtractOpenGANPipeline,
        'test_ood': TestOODPipeline,
        'test_ad': TestAdPipeline,
        'train_ad': TrainAdPipeline,
        'train_oe': TrainOEPipeline,
        'train_opengan': TrainOpenGanPipeline,
        'train_arplgan': TrainARPLGANPipeline,
        'test_ood_aps': TestOODPipelineAPS
    }

    return pipelines[config.pipeline.name](config)

openood/pipelines/utils.py
kai422/SCALE
from .scale_postprocessor import ScalePostprocessor
from .base_postprocessor import BasePostprocessor
from .utils import get_postprocessor

openood/postprocessors/__init__.py
kai422/SCALE
from typing import Any
from tqdm import tqdm

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

import openood.utils.comm as comm


class BasePostprocessor:
    def __init__(self, config):
        self.config = config

    def setup(self, net: nn.Module, id_loader_dict, ood_loader_dict):
        pass

    @torch.no_grad()
    def postprocess(self, net: nn.Module, data: Any):
        output = net(data)
        score = torch.softmax(output, dim=1)
        conf, pred = torch.max(score, dim=1)
        return pred, conf

    def inference(self,
                  net: nn.Module,
                  data_loader: DataLoader,
                  progress: bool = True):
        pred_list, conf_list, label_list = [], [], []
        for batch in tqdm(data_loader,
                          disable=not progress or not comm.is_main_process()):
            data = batch['data'].cuda()
            label = batch['label'].cuda()
            pred, conf = self.postprocess(net, data)

            pred_list.append(pred.cpu())
            conf_list.append(conf.cpu())
            label_list.append(label.cpu())

        # convert values into numpy array
        pred_list = torch.cat(pred_list).numpy().astype(int)
        conf_list = torch.cat(conf_list).numpy()
        label_list = torch.cat(label_list).numpy().astype(int)

        return pred_list, conf_list, label_list

openood/postprocessors/base_postprocessor.py
kai422/SCALE
num_classes_dict = {
    'imagenet': 1000
}

openood/postprocessors/info.py
kai422/SCALE
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .base_postprocessor import BasePostprocessor


class ScalePostprocessor(BasePostprocessor):
    def __init__(self, config):
        super(ScalePostprocessor, self).__init__(config)
        self.args = self.config.postprocessor.postprocessor_args
        self.percentile = self.args.percentile
        self.args_dict = self.config.postprocessor.postprocessor_sweep

    @torch.no_grad()
    def postprocess(self, net: nn.Module, data: Any):
        output = net.forward_threshold(data, self.percentile)
        _, pred = torch.max(output, dim=1)
        energyconf = torch.logsumexp(output.data.cpu(), dim=1)
        return pred, energyconf

    def set_hyperparam(self, hyperparam: list):
        self.percentile = hyperparam[0]

    def get_hyperparam(self):
        return self.percentile

openood/postprocessors/scale_postprocessor.py
kai422/SCALE
from openood.utils import Config

from .scale_postprocessor import ScalePostprocessor

def get_postprocessor(config: Config):
    postprocessors = {
        'scale': ScalePostprocessor,
    }

    return postprocessors[config.postprocessor.name](config)

openood/postprocessors/utils.py
kai422/SCALE
from .base_preprocessor import BasePreprocessor
from .cutpaste_preprocessor import CutPastePreprocessor
from .draem_preprocessor import DRAEMPreprocessor
from .pixmix_preprocessor import PixMixPreprocessor
from .test_preprocessor import TestStandardPreProcessor
from .utils import get_preprocessor

openood/preprocessors/__init__.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .transform import Convert, interpolation_modes, normalization_dict


class AugMixPreprocessor():
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        self.severity = config.preprocessor.severity
        self.mixture_width = config.preprocessor.mixture_width
        self.alpha = config.preprocessor.alpha
        self.chain_depth = config.preprocessor.chain_depth
        self.all_ops = config.preprocessor.all_ops
        self.jsd = config.trainer.trainer_args.jsd

        self.augmix = tvs_trans.AugMix(severity=self.severity,
                                       mixture_width=self.mixture_width,
                                       chain_depth=self.chain_depth,
                                       alpha=self.alpha,
                                       all_ops=self.all_ops,
                                       interpolation=self.interpolation)
        self.normalize = tvs_trans.Compose([
            tvs_trans.ToTensor(),
            tvs_trans.Normalize(mean=self.mean, std=self.std),
        ])

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(0.5),
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomCrop(self.image_size, padding=4),
            ])

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        if self.jsd:
            orig = self.transform(image)
            aug1 = self.normalize(self.augmix(orig))
            aug2 = self.normalize(self.augmix(orig))
            return self.normalize(orig), aug1, aug2
        else:
            return self.normalize(self.augmix(self.transform(image)))

openood/preprocessors/augmix_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .transform import Convert, interpolation_modes, normalization_dict


class BasePreprocessor():
    """For train dataset standard transformation."""
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.ColorJitter(brightness=32. / 255., saturation=0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomCrop(self.image_size, padding=4),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        return self.transform(image)

openood/preprocessors/base_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .transform import Convert, interpolation_modes, normalization_dict


class CiderPreprocessor():
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(size=self.image_size,
                                            scale=(0.4, 1.),
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomApply(
                    [tvs_trans.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
                tvs_trans.RandomGrayscale(p=0.2),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.RandomResizedCrop(size=self.image_size,
                                            scale=(0.2, 1.),
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomApply(
                    [tvs_trans.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
                tvs_trans.RandomGrayscale(p=0.2),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])

        self.transform = TwoCropTransform(self.transform)

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        return self.transform(image)


class TwoCropTransform:
    """Create two crops of the same image."""
    def __init__(self, transform):
        self.transform = transform

    def __call__(self, x):
        return [self.transform(x), self.transform(x)]

openood/preprocessors/cider_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .transform import Convert, interpolation_modes, normalization_dict


class CSIPreprocessor():
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                # tvs_trans.RandomHorizontalFlip(0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                # tvs_trans.RandomHorizontalFlip(),
                # tvs_trans.ColorJitter(brightness=32./255., saturation=0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                # tvs_trans.RandomHorizontalFlip(),
                # tvs_trans.RandomCrop(self.image_size, padding=4),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        return self.transform(image)

openood/preprocessors/csi_preprocessor.py
kai422/SCALE
import numpy as np
import torch
import torchvision.transforms as tvs_trans

from openood.utils.config import Config
from .transform import Convert, interpolation_modes, normalization_dict


class CutoutPreprocessor():
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        self.n_holes = config.preprocessor.n_holes
        self.length = config.preprocessor.length

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
                Cutout(n_holes=self.n_holes, length=self.length)
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.ColorJitter(brightness=32. / 255., saturation=0.5),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
                Cutout(n_holes=self.n_holes, length=self.length)
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomCrop(self.image_size, padding=4),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
                Cutout(n_holes=self.n_holes, length=self.length)
            ])

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        return self.transform(image)


class Cutout(object):
    """Randomly mask out one or more patches from an image.

    Args:
        n_holes (int): Number of patches to cut out of each image.
        length (int): The length (in pixels) of each square patch.
    """
    def __init__(self, n_holes, length):
        self.n_holes = n_holes
        self.length = length

    def __call__(self, img):
        """
        Args:
            img (Tensor): Tensor image of size (C, H, W).
        Returns:
            Tensor: Image with n_holes of dimension length x length
            cut out of it.
        """
        h = img.size(1)
        w = img.size(2)

        mask = np.ones((h, w), np.float32)

        for n in range(self.n_holes):
            y = np.random.randint(h)
            x = np.random.randint(w)

            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)

            mask[y1:y2, x1:x2] = 0.

        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img = img * mask

        return img

openood/preprocessors/cutout_preprocessor.py
kai422/SCALE
import math
import random

import torch
import torchvision.transforms as tvs_trans

from .base_preprocessor import BasePreprocessor
from .transform import Convert, normalization_dict


class CutPastePreprocessor(BasePreprocessor):
    def __init__(
            self, config,
            split):  # modify, preprocessors unify to only passing in "config"
        self.args = config.preprocessor.preprocessor_args
        self.area_ratio = self.args.area_ratio
        self.aspect_ratio = self.args.aspect_ratio

        dataset_name = config.dataset.name.split('_')[0]
        image_size = config.dataset.image_size
        pre_size = config.dataset.pre_size
        if dataset_name in normalization_dict.keys():
            mean = normalization_dict[dataset_name][0]
            std = normalization_dict[dataset_name][1]
        else:
            mean = [0.485, 0.456, 0.406]
            std = [0.229, 0.224, 0.225]

        self.before_preprocessor_transform = tvs_trans.Compose([
            Convert('RGB'),
            tvs_trans.Resize(
                pre_size, interpolation=tvs_trans.InterpolationMode.BILINEAR),
            tvs_trans.CenterCrop(image_size),
            tvs_trans.RandomHorizontalFlip(),
            tvs_trans.RandomCrop(image_size, padding=4),
        ])
        self.after_preprocessor_transform = tvs_trans.Compose([
            tvs_trans.ToTensor(),
            tvs_trans.Normalize(mean=mean, std=std),
        ])

    def __call__(self, img):
        img = self.before_preprocessor_transform(img)

        h = img.size[0]
        w = img.size[1]

        # ratio between area_ratio[0] and area_ratio[1]
        ratio_area = random.uniform(0.02, 0.15) * w * h

        # sample in log space
        log_ratio = torch.log(
            torch.tensor((self.aspect_ratio, 1 / self.aspect_ratio)))
        aspect = torch.exp(
            torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()

        cut_w = int(round(math.sqrt(ratio_area * aspect)))
        cut_h = int(round(math.sqrt(ratio_area / aspect)))

        from_location_h = int(random.uniform(0, h - cut_h))
        from_location_w = int(random.uniform(0, w - cut_w))

        box = [
            from_location_w, from_location_h, from_location_w + cut_w,
            from_location_h + cut_h
        ]
        patch = img.crop(box)

        # if self.colorJitter:
        #     patch = self.colorJitter(patch)

        to_location_h = int(random.uniform(0, h - cut_h))
        to_location_w = int(random.uniform(0, w - cut_w))

        insert_box = [
            to_location_w, to_location_h, to_location_w + cut_w,
            to_location_h + cut_h
        ]
        augmented = img.copy()
        augmented.paste(patch, insert_box)

        img = self.after_preprocessor_transform(img)
        augmented = self.after_preprocessor_transform(augmented)

        return img, augmented

openood/preprocessors/cutpaste_preprocessor.py
kai422/SCALE
import glob
import math
import os

import cv2
import imgaug.augmenters as iaa
import numpy as np
import torch

from .base_preprocessor import BasePreprocessor


class DRAEMPreprocessor(BasePreprocessor):
    def __init__(self, config):
        self.config = config
        self.args = self.config.preprocessor.preprocessor_args

        self.resize_shape = [self.args.image_size, self.args.image_size]

        self.anomaly_source_paths = sorted(
            glob.glob(self.args.anomaly_source + '/*/*.jpg'))

        self.augmenters = [
            iaa.GammaContrast((0.5, 2.0), per_channel=True),
            iaa.MultiplyAndAddToBrightness(mul=(0.8, 1.2), add=(-30, 30)),
            iaa.pillike.EnhanceSharpness(),
            iaa.AddToHueAndSaturation((-50, 50), per_channel=True),
            iaa.Solarize(0.5, threshold=(32, 128)),
            iaa.Posterize(),
            iaa.Invert(),
            iaa.pillike.Autocontrast(),
            iaa.pillike.Equalize(),
            iaa.Affine(rotate=(-45, 45))
        ]

        self.rot = iaa.Sequential([iaa.Affine(rotate=(-90, 90))])

        # if config.evaluator.name == 'ood':
        #     assert config.use_gt == False
        # if config.evaluator.name == 'draem':
        #     assert config.use_gt == True

    def transform_test_image(self, image_path, mask_path):
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        if self.resize_shape is not None:
            image = cv2.resize(image,
                               dsize=(self.resize_shape[1],
                                      self.resize_shape[0]))

        image = image / 255.0

        image = np.array(image).reshape(
            (image.shape[0], image.shape[1], 3)).astype(np.float32)

        image = np.transpose(image, (2, 0, 1))
        mask = image
        if self.config.use_gt:
            if mask_path is not None:
                mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            else:
                mask = np.zeros((image.shape[0], image.shape[1]))
            if self.resize_shape is not None:
                mask = cv2.resize(mask,
                                  dsize=(self.resize_shape[1],
                                         self.resize_shape[0]))
            mask = mask / 255.0
            mask = np.array(mask).reshape(
                (mask.shape[0], mask.shape[1], 1)).astype(np.float32)
            mask = np.transpose(mask, (2, 0, 1))

        return image, mask

    def get_test_item(self, path):
        sample = {}
        dir_path, file_name = os.path.split(path)
        base_dir = os.path.basename(dir_path)
        if base_dir == 'good':
            image, mask = self.transform_test_image(path, None)
        else:
            mask_path = os.path.join(dir_path, '../../ground_truth/')
            mask_path = os.path.join(mask_path, base_dir)
            mask_file_name = file_name.split('.')[0] + '_mask.png'
            mask_path = os.path.join(mask_path, mask_file_name)
            image, mask = self.transform_test_image(path, mask_path)

        if self.config.use_gt:
            sample['image'] = image
            sample['mask'] = mask
            return sample
        else:
            return image

    def randAugmenter(self):
        aug_ind = np.random.choice(np.arange(len(self.augmenters)),
                                   3,
                                   replace=False)
        aug = iaa.Sequential([
            self.augmenters[aug_ind[0]], self.augmenters[aug_ind[1]],
            self.augmenters[aug_ind[2]]
        ])
        return aug

    def augment_image(self, image, anomaly_source_path):
        aug = self.randAugmenter()
        perlin_scale = 6
        min_perlin_scale = 0
        anomaly_source_img = cv2.imread(anomaly_source_path)

        anomaly_source_img = cv2.resize(anomaly_source_img,
                                        dsize=(self.resize_shape[1],
                                               self.resize_shape[0]))
        anomaly_img_augmented = aug(image=anomaly_source_img)
        perlin_scalex = 2**(torch.randint(min_perlin_scale, perlin_scale,
                                          (1, )).numpy()[0])
        perlin_scaley = 2**(torch.randint(min_perlin_scale, perlin_scale,
                                          (1, )).numpy()[0])

        perlin_noise = rand_perlin_2d_np(
            (self.resize_shape[0], self.resize_shape[1]),
            (perlin_scalex, perlin_scaley))
        perlin_noise = self.rot(image=perlin_noise)
        threshold = 0.5
        perlin_thr = np.where(perlin_noise > threshold,
                              np.ones_like(perlin_noise),
                              np.zeros_like(perlin_noise))
        perlin_thr = np.expand_dims(perlin_thr, axis=2)

        img_thr = anomaly_img_augmented.astype(np.float32) * perlin_thr / 255.0

        beta = torch.rand(1).numpy()[0] * 0.8

        augmented_image = image * (1 - perlin_thr) + (
            1 - beta) * img_thr + beta * image * (perlin_thr)

        no_anomaly = torch.rand(1).numpy()[0]
        if no_anomaly > 0.5:
            image = image.astype(np.float32)
            return image, np.zeros_like(
                perlin_thr, dtype=np.float32), np.array([0.0],
                                                        dtype=np.float32)
        else:
            augmented_image = augmented_image.astype(np.float32)
            msk = (perlin_thr).astype(np.float32)
            augmented_image = msk * augmented_image + (1 - msk) * image
            has_anomaly = 1.0
            if np.sum(msk) == 0:
                has_anomaly = 0.0
            return augmented_image, msk, np.array([has_anomaly],
                                                  dtype=np.float32)

    def transform_train_image(self, image_path, anomaly_source_path):
        image = cv2.imread(image_path)
        image = cv2.resize(image,
                           dsize=(self.resize_shape[1], self.resize_shape[0]))

        do_aug_orig = torch.rand(1).numpy()[0] > 0.7
        if do_aug_orig:
            image = self.rot(image=image)

        image = np.array(image).reshape(
            (image.shape[0], image.shape[1], image.shape[2])).astype(
                np.float32) / 255.0
        augmented_image, anomaly_mask, has_anomaly = self.augment_image(
            image, anomaly_source_path)
        augmented_image = np.transpose(augmented_image, (2, 0, 1))
        image = np.transpose(image, (2, 0, 1))
        anomaly_mask = np.transpose(anomaly_mask, (2, 0, 1))
        return image, augmented_image, anomaly_mask, has_anomaly

    def get_train_item(self, path):
        sample = {}
        anomaly_source_idx = torch.randint(0, len(self.anomaly_source_paths),
                                           (1, )).item()
        image, augmented_image, anomaly_mask, has_anomaly = \
            self.transform_train_image(
                        path, self.anomaly_source_paths[anomaly_source_idx])
        sample['image'] = image
        sample['anomaly_mask'] = anomaly_mask
        sample['augmented_image'] = augmented_image
        sample['has_anomaly'] = has_anomaly

        return sample

    def __call__(self, img):
        if self.name.endswith('_train'):
            sample = self.get_train_item(self.path)
        else:
            sample = self.get_test_item(self.path)
        return sample

    # some setup so that the preprocessor can get the gt map
    def setup(self, **kwargs):
        self.path = kwargs['path']
        self.name = kwargs['name']

    # append transforms that will apply after the preprocessor
    def concat_transform(self, post_preprocessor_transform=None):
        self.post_preprocessor_transform = post_preprocessor_transform
        return self


def lerp_np(x, y, w):
    fin_out = (y - x) * w + x
    return fin_out


def generate_fractal_noise_2d(shape, res, octaves=1, persistence=0.5):
    noise = np.zeros(shape)
    frequency = 1
    amplitude = 1
    for _ in range(octaves):
        noise += amplitude * generate_perlin_noise_2d(
            shape, (frequency * res[0], frequency * res[1]))
        frequency *= 2
        amplitude *= persistence
    return noise


def generate_perlin_noise_2d(shape, res):
    def f(t):
        return 6 * t**5 - 15 * t**4 + 10 * t**3

    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]].transpose(1, 2,
                                                                    0) % 1
    # Gradients
    angles = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1)
    gradients = np.dstack((np.cos(angles), np.sin(angles)))
    g00 = gradients[0:-1, 0:-1].repeat(d[0], 0).repeat(d[1], 1)
    g10 = gradients[1:, 0:-1].repeat(d[0], 0).repeat(d[1], 1)
    g01 = gradients[0:-1, 1:].repeat(d[0], 0).repeat(d[1], 1)
    g11 = gradients[1:, 1:].repeat(d[0], 0).repeat(d[1], 1)
    # Ramps
    n00 = np.sum(grid * g00, 2)
    n10 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1])) * g10, 2)
    n01 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1] - 1)) * g01, 2)
    n11 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1] - 1)) * g11, 2)
    # Interpolation
    t = f(grid)
    n0 = n00 * (1 - t[:, :, 0]) + t[:, :, 0] * n10
    n1 = n01 * (1 - t[:, :, 0]) + t[:, :, 0] * n11
    return np.sqrt(2) * ((1 - t[:, :, 1]) * n0 + t[:, :, 1] * n1)


def rand_perlin_2d_np(shape,
                      res,
                      fade=lambda t: 6 * t**5 - 15 * t**4 + 10 * t**3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]].transpose(1, 2,
                                                                    0) % 1

    angles = 2 * math.pi * np.random.rand(res[0] + 1, res[1] + 1)
    gradients = np.stack((np.cos(angles), np.sin(angles)), axis=-1)

    tile_grads = lambda slice1, slice2: np.repeat(np.repeat(
        gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]], d[0], axis=0),
                                                  d[1],
                                                  axis=1)
    dot = lambda grad, shift: (np.stack(
        (grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1
                                                        ] + shift[1]),
        axis=-1) * grad[:shape[0], :shape[1]]).sum(axis=-1)

    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])
    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])
    t = fade(grid[:shape[0], :shape[1]])
    return math.sqrt(2) * lerp_np(lerp_np(n00, n10, t[..., 0]),
                                  lerp_np(n01, n11, t[..., 0]), t[..., 1])


def rand_perlin_2d(shape,
                   res,
                   fade=lambda t: 6 * t**5 - 15 * t**4 + 10 * t**3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])

    grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]),
                                      torch.arange(0, res[1], delta[1])),
                       dim=-1) % 1
    angles = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1)
    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1)

    tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[
        0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1)
    dot = lambda grad, shift: (torch.stack(
        (grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1
                                                        ] + shift[1]),
        dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1)

    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])

    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])
    t = fade(grid[:shape[0], :shape[1]])
    return math.sqrt(2) * torch.lerp(torch.lerp(
        n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1])


def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5):
    noise = torch.zeros(shape)
    frequency = 1
    amplitude = 1
    for _ in range(octaves):
        noise += amplitude * rand_perlin_2d(
            shape, (frequency * res[0], frequency * res[1]))
        frequency *= 2
        amplitude *= persistence
    return noise

openood/preprocessors/draem_preprocessor.py
kai422/SCALE
import os

import numpy as np
import torch
import torchvision.transforms as tvs_trans
from PIL import Image as Image
from PIL import ImageEnhance, ImageOps

from .base_preprocessor import BasePreprocessor
from .transform import Convert, interpolation_modes, normalization_dict

resize_list = {
    'osr': 32,
    'mnist': 32,
    'cifar10': 36,
    'cifar100': 36,
    'tin': 72,
    'imagenet': 256,
    'imagenet200': 256,
    'aircraft': 512,
    'cub': 512,
}  # set mnist bymyself, imagenet was set to 224 by author, but 256 here


class PixMixPreprocessor(BasePreprocessor):
    def __init__(self, config):
        self.pre_size = config.dataset.pre_size
        self.dataset_name = config.dataset.name.split('_')[0]
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]
        self.normalize = tvs_trans.Normalize(self.mean, self.std)
        self.tensorize = tvs_trans.ToTensor()

        self.args = config.preprocessor.preprocessor_args

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(0.5),
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomCrop(self.image_size, padding=4),
            ])

        self.mixing_set_transform = tvs_trans.Compose([
            tvs_trans.Resize(resize_list[self.dataset_name]),
            tvs_trans.RandomCrop(self.image_size)
        ])

        with open(self.args.mixing_set_dir, 'r') as f:
            self.mixing_list = f.readlines()

    def __call__(self, image):
        # ? need to add random seed ?
        rnd_idx = np.random.choice(len(self.mixing_list))
        mixing_pic_dir = self.mixing_list[rnd_idx].strip('\n')

        mixing_pic = Image.open(
            os.path.join('./data/images_classic/',
                         mixing_pic_dir)).convert('RGB')
        return self.pixmix(image, mixing_pic)

    def augment_input(self, image):
        aug_list = augmentations_all if self.args.all_ops else augmentations
        op = np.random.choice(aug_list)
        return op(image.copy(), self.args.aug_severity, self.image_size)

    def pixmix(self, orig, mixing_pic):
        mixings = [add, multiply]
        orig = self.transform(orig)

        # do basic augmentation first
        mixing_pic = self.mixing_set_transform(mixing_pic)

        if np.random.random() < 0.5:
            mixed = self.tensorize(self.augment_input(orig))
        else:
            mixed = self.tensorize(orig)

        for _ in range(np.random.randint(self.args.k + 1)):

            if np.random.random() < 0.5:
                aug_image_copy = self.tensorize(self.augment_input(orig))
            else:
                aug_image_copy = self.tensorize(mixing_pic)

            mixed_op = np.random.choice(mixings)

            mixed = mixed_op(mixed, aug_image_copy, self.args.beta)
            mixed = torch.clip(mixed, 0, 1)

        return self.normalize(mixed)


"""Base augmentations operators."""

#########################################################
#################### AUGMENTATIONS ######################
#########################################################


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, _, IMAGE_SIZE):
    return ImageOps.autocontrast(pil_img)


def equalize(pil_img, _, IMAGE_SIZE):
    return ImageOps.equalize(pil_img)


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


def rotate(pil_img, level, IMAGE_SIZE):
    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, IMAGE_SIZE):
    level = int_parameter(sample_level(level), 256)
    return ImageOps.solarize(pil_img, 256 - level)


def shear_x(pil_img, level, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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, IMAGE_SIZE):
    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
]

#########################################################
######################## MIXINGS ########################
#########################################################


def get_ab(beta):
    if np.random.random() < 0.5:
        a = np.float32(np.random.beta(beta, 1))
        b = np.float32(np.random.beta(1, beta))
    else:
        a = 1 + np.float32(np.random.beta(1, beta))
        b = -np.float32(np.random.beta(1, beta))
    return a, b


def add(img1, img2, beta):
    a, b = get_ab(beta)
    img1, img2 = img1 * 2 - 1, img2 * 2 - 1
    out = a * img1 + b * img2
    return (out + 1) / 2


def multiply(img1, img2, beta):
    a, b = get_ab(beta)
    img1, img2 = img1 * 2, img2 * 2
    out = (img1**a) * (img2.clip(1e-37)**b)
    return out / 2


########################################
##### EXTRA MIXIMGS (EXPREIMENTAL) #####
########################################


def invert(img):
    return 1 - img


def screen(img1, img2, beta):
    img1, img2 = invert(img1), invert(img2)
    out = multiply(img1, img2, beta)
    return invert(out)


def overlay(img1, img2, beta):
    case1 = multiply(img1, img2, beta)
    case2 = screen(img1, img2, beta)
    if np.random.random() < 0.5:
        cond = img1 < 0.5
    else:
        cond = img1 > 0.5
    return torch.where(cond, case1, case2)


def darken_or_lighten(img1, img2, beta):
    if np.random.random() < 0.5:
        cond = img1 < img2
    else:
        cond = img1 > img2
    return torch.where(cond, img1, img2)


def swap_channel(img1, img2, beta):
    channel = np.random.randint(3)
    img1[channel] = img2[channel]
    return img1

openood/preprocessors/pixmix_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .transform import Convert, interpolation_modes, normalization_dict


class RandAugmentPreprocessor():
    def __init__(self, config: Config):
        self.pre_size = config.dataset.pre_size
        self.image_size = config.dataset.image_size
        self.interpolation = interpolation_modes[config.dataset.interpolation]
        normalization_type = config.dataset.normalization_type
        if normalization_type in normalization_dict.keys():
            self.mean = normalization_dict[normalization_type][0]
            self.std = normalization_dict[normalization_type][1]
        else:
            self.mean = [0.5, 0.5, 0.5]
            self.std = [0.5, 0.5, 0.5]

        self.n = config.preprocessor.n
        self.m = config.preprocessor.m

        if 'imagenet' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.RandomResizedCrop(self.image_size,
                                            interpolation=self.interpolation),
                tvs_trans.RandomHorizontalFlip(0.5),
                tvs_trans.RandAugment(num_ops=self.n,
                                      magnitude=self.m,
                                      interpolation=self.interpolation),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        elif 'aircraft' in config.dataset.name or 'cub' in config.dataset.name:
            self.transform = tvs_trans.Compose([
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.RandomCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandAugment(num_ops=self.n,
                                      magnitude=self.m,
                                      interpolation=self.interpolation),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])
        else:
            self.transform = tvs_trans.Compose([
                Convert('RGB'),
                tvs_trans.RandAugment(num_ops=self.n,
                                      magnitude=self.m,
                                      interpolation=self.interpolation),
                tvs_trans.Resize(self.pre_size,
                                 interpolation=self.interpolation),
                tvs_trans.CenterCrop(self.image_size),
                tvs_trans.RandomHorizontalFlip(),
                tvs_trans.RandomCrop(self.image_size, padding=4),
                tvs_trans.ToTensor(),
                tvs_trans.Normalize(mean=self.mean, std=self.std),
            ])

    def setup(self, **kwargs):
        pass

    def __call__(self, image):
        return self.transform(image)

openood/preprocessors/randaugment_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

from openood.utils.config import Config

from .base_preprocessor import BasePreprocessor
from .transform import Convert


class TestStandardPreProcessor(BasePreprocessor):
    """For test and validation dataset standard image transformation."""
    def __init__(self, config: Config):
        super(TestStandardPreProcessor, self).__init__(config)
        self.transform = tvs_trans.Compose([
            Convert('RGB'),
            tvs_trans.Resize(self.pre_size, interpolation=self.interpolation),
            tvs_trans.CenterCrop(self.image_size),
            tvs_trans.ToTensor(),
            tvs_trans.Normalize(mean=self.mean, std=self.std),
        ])

openood/preprocessors/test_preprocessor.py
kai422/SCALE
import torchvision.transforms as tvs_trans

normalization_dict = {
    'cifar10': [[0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]],
    'cifar100': [[0.5071, 0.4867, 0.4408], [0.2675, 0.2565, 0.2761]],
    'imagenet': [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
    'imagenet200': [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
    'covid': [[0.4907, 0.4907, 0.4907], [0.2697, 0.2697, 0.2697]],
    'aircraft': [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
    'cub': [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
    'cars': [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
}

interpolation_modes = {
    'nearest': tvs_trans.InterpolationMode.NEAREST,
    'bilinear': tvs_trans.InterpolationMode.BILINEAR,
}


class Convert:
    def __init__(self, mode='RGB'):
        self.mode = mode

    def __call__(self, image):
        return image.convert(self.mode)


# More transform classes shall be written here

openood/preprocessors/transform.py
kai422/SCALE
from openood.utils import Config

from .base_preprocessor import BasePreprocessor
from .cider_preprocessor import CiderPreprocessor
from .csi_preprocessor import CSIPreprocessor
from .cutpaste_preprocessor import CutPastePreprocessor
from .draem_preprocessor import DRAEMPreprocessor
from .augmix_preprocessor import AugMixPreprocessor
from .pixmix_preprocessor import PixMixPreprocessor
from .randaugment_preprocessor import RandAugmentPreprocessor
from .cutout_preprocessor import CutoutPreprocessor
from .test_preprocessor import TestStandardPreProcessor


def get_preprocessor(config: Config, split):
    train_preprocessors = {
        'base': BasePreprocessor,
        'draem': DRAEMPreprocessor,
        'cutpaste': CutPastePreprocessor,
        'augmix': AugMixPreprocessor,
        'pixmix': PixMixPreprocessor,
        'randaugment': RandAugmentPreprocessor,
        'cutout': CutoutPreprocessor,
        'csi': CSIPreprocessor,
        'cider': CiderPreprocessor,
    }
    test_preprocessors = {
        'base': TestStandardPreProcessor,
        'draem': DRAEMPreprocessor,
        'cutpaste': CutPastePreprocessor,
    }

    if split == 'train':
        return train_preprocessors[config.preprocessor.name](config)
    else:
        try:
            return test_preprocessors[config.preprocessor.name](config)
        except KeyError:
            return test_preprocessors['base'](config)

openood/preprocessors/utils.py
kai422/SCALE
from .utils import get_recorder

openood/recorders/__init__.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch

from .base_recorder import BaseRecorder


class ADRecorder(BaseRecorder):
    def __init__(self, config) -> None:
        super(ADRecorder, self).__init__(config)

        self.best_epoch_idx = 0
        self.best_result = 0

        self.begin_time = time.time()

    def report(self, train_metrics, test_metrics):
        print('Epoch {:03d} | Time {:5d}s | Train Loss {:.4f} | '
              'Auroc {:.4f}\n'.format(train_metrics['epoch_idx'],
                                      int(time.time() - self.begin_time),
                                      train_metrics['loss'],
                                      100.0 * test_metrics['image_auroc']),
              flush=True)

    def save_model(self, net, test_metrics):
        if self.config.recorder.save_all_models:
            torch.save(
                net.state_dict(),
                os.path.join(
                    self.output_dir,
                    'model_epoch{}.ckpt'.format(test_metrics['epoch_idx'])))

        # enter only if lower loss occurs
        if test_metrics['image_auroc'] >= self.best_result:

            # delete the depreciated best model
            old_fname = 'best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_result = test_metrics['image_auroc']
            torch.save(net.state_dict(),
                       os.path.join(self.output_dir, 'best.ckpt'))

            save_fname = 'best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net.state_dict(), save_pth)

    def summary(self):
        print('Training Completed!\n '
              'Best Auroc: {:.4f} at epoch {:d}\n'.format(
                  100.0 * self.best_result, self.best_epoch_idx),
              flush=True)

openood/recorders/ad_recorder.py
kai422/SCALE
import copy
import os
import time

import torch

from .base_recorder import BaseRecorder


class ARPLRecorder(BaseRecorder):
    def __init__(self, config) -> None:
        super().__init__(config)

    def report(self, train_metrics, val_metrics):
        if 'lossD' in train_metrics.keys():
            print('\nEpoch {:03d} | Time {:5d}s | D Loss {:.4f} | '
                  'G Loss {:.4f} | Train Loss {:.4f} | '
                  'Val Loss {:.3f} | Val Acc {:.2f}'.format(
                      (train_metrics['epoch_idx']),
                      int(time.time() - self.begin_time),
                      train_metrics['lossD'], train_metrics['lossG'],
                      train_metrics['loss'], val_metrics['loss'],
                      100.0 * val_metrics['acc']),
                  flush=True)
        else:
            print('\nEpoch {:03d} | Time {:5d}s | Train Loss {:.4f} | '
                  'Val Loss {:.3f} | Val Acc {:.2f}'.format(
                      (train_metrics['epoch_idx']),
                      int(time.time() - self.begin_time),
                      train_metrics['loss'], val_metrics['loss'],
                      100.0 * val_metrics['acc']),
                  flush=True)

    def save_model(self, net, val_metrics):

        netF = net['netF']
        criterion = net['criterion']
        epoch_idx = val_metrics['epoch_idx']

        try:
            netF_wts = copy.deepcopy(netF.module.state_dict())
            criterion_wts = copy.deepcopy(criterion.module.state_dict())
        except AttributeError:
            netF_wts = copy.deepcopy(netF.state_dict())
            criterion_wts = copy.deepcopy(criterion.state_dict())

        if self.config.recorder.save_all_models:
            save_pth = os.path.join(self.save_dir,
                                    'epoch-{}_NetF.ckpt'.format(epoch_idx))
            torch.save(netF_wts, save_pth)
            save_pth = os.path.join(
                self.save_dir, 'epoch-{}_criterion.ckpt'.format(epoch_idx))
            torch.save(criterion_wts, save_pth)

        # enter only if better accuracy occurs
        if val_metrics['acc'] >= self.best_acc:

            # update the best model
            self.best_epoch_idx = val_metrics['epoch_idx']
            self.best_acc = val_metrics['acc']

            torch.save(netF_wts, os.path.join(self.output_dir,
                                              'best_NetF.ckpt'))
            torch.save(criterion_wts,
                       os.path.join(self.output_dir, 'best_criterion.ckpt'))

openood/recorders/arpl_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch


class BaseRecorder:
    def __init__(self, config) -> None:
        self.config = config

        self.best_acc = 0.0
        self.best_epoch_idx = 0

        self.begin_time = time.time()
        self.output_dir = config.output_dir

    def report(self, train_metrics, val_metrics):
        print('\nEpoch {:03d} | Time {:5d}s | Train Loss {:.4f} | '
              'Val Loss {:.3f} | Val Acc {:.2f}'.format(
                  (train_metrics['epoch_idx']),
                  int(time.time() - self.begin_time), train_metrics['loss'],
                  val_metrics['loss'], 100.0 * val_metrics['acc']),
              flush=True)

    def save_model(self, net, val_metrics):
        try:
            state_dict = net.module.state_dict()
        except AttributeError:
            state_dict = net.state_dict()

        if self.config.recorder.save_all_models:
            torch.save(
                state_dict,
                os.path.join(
                    self.output_dir,
                    'model_epoch{}.ckpt'.format(val_metrics['epoch_idx'])))

        # enter only if better accuracy occurs
        if val_metrics['acc'] >= self.best_acc:
            # delete the depreciated best model
            old_fname = 'best_epoch{}_acc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_acc)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = val_metrics['epoch_idx']
            self.best_acc = val_metrics['acc']
            torch.save(state_dict, os.path.join(self.output_dir, 'best.ckpt'))

            save_fname = 'best_epoch{}_acc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_acc)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(state_dict, save_pth)

        # save last path
        if val_metrics['epoch_idx'] == self.config.optimizer.num_epochs:
            save_fname = 'last_epoch{}_acc{:.4f}.ckpt'.format(
                val_metrics['epoch_idx'], val_metrics['acc'])
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(state_dict, save_pth)

    def summary(self):
        print('Training Completed! '
              'Best accuracy: {:.2f} '
              'at epoch {:d}'.format(100 * self.best_acc, self.best_epoch_idx),
              flush=True)

openood/recorders/base_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch


class CiderRecorder:
    def __init__(self, config) -> None:
        self.config = config

        self.best_loss = float('inf')
        self.best_epoch_idx = 0

        self.begin_time = time.time()
        self.output_dir = config.output_dir

    def report(self, train_metrics, val_metrics):
        print('\nEpoch {:03d} | Time {:5d}s | Train Loss {:.4f}'.format(
            (train_metrics['epoch_idx']), int(time.time() - self.begin_time),
            train_metrics['loss']),
              flush=True)

    def save_model(self, net, train_metrics):
        try:
            state_dict = net.module.state_dict()
        except AttributeError:
            state_dict = net.state_dict()

        if self.config.recorder.save_all_models:
            torch.save(
                state_dict,
                os.path.join(
                    self.output_dir,
                    'model_epoch{}.ckpt'.format(train_metrics['epoch_idx'])))

        # enter only if better accuracy occurs
        if train_metrics['loss'] <= self.best_loss:
            # delete the depreciated best model
            old_fname = 'best_epoch{}_loss{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_loss)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = train_metrics['epoch_idx']
            self.best_loss = train_metrics['loss']
            torch.save(state_dict, os.path.join(self.output_dir, 'best.ckpt'))

            save_fname = 'best_epoch{}_loss{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_loss)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(state_dict, save_pth)

        # save last path
        if train_metrics['epoch_idx'] == self.config.optimizer.num_epochs:
            save_fname = 'last_epoch{}_loss{:.4f}.ckpt'.format(
                train_metrics['epoch_idx'], train_metrics['loss'])
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(state_dict, save_pth)

    def summary(self):
        print('Training Completed! '
              'Best loss: {:.4f} '
              'at epoch {:d}'.format(self.best_loss, self.best_epoch_idx),
              flush=True)

openood/recorders/cider_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch


class CutpasteRecorder:
    def __init__(self, config) -> None:
        self.config = config

        self.best_auroc = 0.0
        self.best_epoch_idx = 0

        self.begin_time = time.time()
        self.output_dir = config.output_dir

    def report(self, train_metrics, val_metrics):
        print('\nEpoch {:03d} | Time {:5d}s | Train Loss {:.4f} | '
              'AUROC {:.3f}'.format((val_metrics['epoch_idx']),
                                    int(time.time() - self.begin_time),
                                    train_metrics['loss'],
                                    val_metrics['image_auroc']),
              flush=True)

    def save_model(self, net, val_metrics):
        if self.config.recorder.save_all_models:
            torch.save(
                net.state_dict(),
                os.path.join(
                    self.output_dir,
                    'model_epoch{}.ckpt'.format(val_metrics['epoch_idx'])))

        # enter only if best auroc occurs
        if val_metrics['image_auroc'] >= self.best_auroc:

            # delete the depreciated best model
            old_fname = 'best_epoch{}_auroc{}.ckpt'.format(
                self.best_epoch_idx, self.best_auroc)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = val_metrics['epoch_idx']
            self.best_auroc = val_metrics['image_auroc']
            torch.save(net.state_dict(),
                       os.path.join(self.output_dir, 'best.ckpt'))

            save_fname = 'best_epoch{}_auroc{}.ckpt'.format(
                self.best_epoch_idx, self.best_auroc)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net.state_dict(), save_pth)

    def summary(self):
        print('Training Completed! '
              'Best auroc: {:.2f} '
              'at epoch {:d}'.format(self.best_auroc, self.best_epoch_idx),
              flush=True)

openood/recorders/cutpaste_recorder.py
kai422/SCALE
import os
from pathlib import Path

import torch

from .ad_recorder import ADRecorder


class DRAEMRecorder(ADRecorder):
    def __init__(self, config) -> None:
        super(DRAEMRecorder, self).__init__(config)

        self.best_model_basis = self.config.recorder.best_model_basis

        self.run_name = ('draem_test_' + str(self.config.optimizer.lr) + '_' +
                         str(self.config.optimizer.num_epochs) + '_bs' +
                         str(self.config.dataset.train.batch_size) + '_' +
                         self.config.dataset.name)

    def save_model(self, net, test_metrics):
        if self.config.recorder.save_all_models:

            save_fname = self.run_name + '_model_epoch{}'.format(
                test_metrics['epoch_idx'])
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net['generative'].state_dict(), save_pth + '.ckpt')
            torch.save(net['discriminative'].state_dict(),
                       save_pth + '_seg.ckpt')

        # enter only if lower loss occurs
        if test_metrics[self.best_model_basis] >= self.best_result:

            # delete the depreciated best model
            old_fname = self.run_name + '_best_epoch{}_loss{:.4f}'.format(
                self.best_epoch_idx, self.best_result)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth + '.ckpt').unlink(missing_ok=True)
            Path(old_pth + '_seg.ckpt').unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_result = test_metrics[self.best_model_basis]

            save_fname = self.run_name + '_best_epoch{}_loss{:.4f}'.format(
                self.best_epoch_idx, self.best_result)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net['generative'].state_dict(), save_pth + '.ckpt')
            torch.save(net['discriminative'].state_dict(),
                       save_pth + '_seg.ckpt')

        if test_metrics['epoch_idx'] == self.config.optimizer.num_epochs:
            save_fname = self.run_name + '_latest_checkpoint'
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net['generative'].state_dict(), save_pth + '.ckpt')
            torch.save(net['discriminative'].state_dict(),
                       save_pth + '_seg.ckpt')

openood/recorders/draem_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch


class DCAERecorder:
    def __init__(self, config) -> None:
        self.config = config
        self.output_dir = config.output_dir
        self.best_roc_auc = 0.0
        self.best_epoch_idx = 0
        self.begin_time = time.time()

    def report(self, train_metrics, test_metrics):
        print('epoch [{}],time:{:5d}s,loss:{:.4f},roc_auc:{:.2f}'.format(
            train_metrics['epoch_idx'], int(time.time() - self.begin_time),
            train_metrics['epoch_loss'], test_metrics['roc_auc']))

    def save_model(self, net, test_metrics):

        # enter only if better accuracy occurs
        if test_metrics['roc_auc'] >= self.best_roc_auc:

            # delete the depreciated best model
            old_fname = 'AE_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_roc_auc = test_metrics['roc_auc']
            save_fname = 'AE_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net.state_dict(), save_pth)

    def summary(self):
        print('Training Completed! '
              'Best Roc_auc: {:.2f}%,'
              'at epoch {:d}'.format(100 * self.best_roc_auc,
                                     self.best_epoch_idx),
              flush=True)


class DSVDDRecorder:
    def __init__(self, config) -> None:
        self.config = config
        self.output_dir = config.output_dir
        self.best_roc_auc = 0.0
        self.best_epoch_idx = 0
        self.begin_time = time.time()

    def report(self, train_metrics, test_metrics):
        print('epoch [{}],time:{:5d}s,loss:{:.4f},roc_auc:{:.2f}'.format(
            train_metrics['epoch_idx'], int(time.time() - self.begin_time),
            train_metrics['epoch_loss'], test_metrics['roc_auc']))

    def save_model(self, net, test_metrics):

        # enter only if better accuracy occurs
        if test_metrics['roc_auc'] >= self.best_roc_auc:

            # delete the depreciated best model
            old_fname = 'DSVDD_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_roc_auc = test_metrics['roc_auc']
            save_fname = 'DSVDD_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net.state_dict(), save_pth)

    def summary(self):
        print('Training Completed! '
              'Best Roc_auc: {:.2f}%,'
              'at epoch {:d}'.format(100 * self.best_roc_auc,
                                     self.best_epoch_idx),
              flush=True)

openood/recorders/dsvdd_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch


class KdadRecorder:
    def __init__(self, config) -> None:
        self.config = config
        self.output_dir = config.output_dir
        self.best_roc_auc = 0.0
        self.best_epoch_idx = 0
        self.begin_time = time.time()

    def report(self, train_metrics, test_metrics):
        print('epoch [{}],time:{:5d}s,loss:{:.4f},roc_auc:{:.2f}'.format(
            train_metrics['epoch_idx'], int(time.time() - self.begin_time),
            train_metrics['epoch_loss'], test_metrics['roc_auc']))

    def save_model(self, net, test_metrics):
        if self.config.recorder.save_all_models:
            torch.save(
                net['model'].state_dict(),
                os.path.join(
                    self.output_dir,
                    'Clone_epoch{}.ckpt'.format(test_metrics['epoch_idx'])))

        # enter only if better accuracy occurs
        if test_metrics['roc_auc'] >= self.best_roc_auc:

            # delete the depreciated best model
            old_fname = 'Clone_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            old_pth = os.path.join(self.output_dir, old_fname)
            Path(old_pth).unlink(missing_ok=True)

            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_roc_auc = test_metrics['roc_auc']
            save_fname = 'Clone_best_epoch{}_roc_auc{}.pth'.format(
                self.best_epoch_idx, self.best_roc_auc)
            save_pth = os.path.join(self.output_dir, save_fname)
            torch.save(net['model'].state_dict(), save_pth)
        if test_metrics['epoch_idx'] == self.config['last_checkpoint']:
            torch.save(
                net['model'].state_dict(),
                '{}/Cloner_{}_epoch_{}.pth'.format(self.config['output_dir'],
                                                   self.config.normal_class,
                                                   test_metrics['epoch_idx']))

    def summary(self):
        print('Training Completed! '
              'Best Roc_auc: {:.2f}%,'
              'at epoch {:d}'.format(100 * self.best_roc_auc,
                                     self.best_epoch_idx),
              flush=True)

openood/recorders/kdad_recorder.py
kai422/SCALE
import copy
import os
import time

import torch

from .base_recorder import BaseRecorder


class OpenGanRecorder(BaseRecorder):
    def __init__(self, config) -> None:
        super().__init__(config)
        self.save_dir = self.config.output_dir
        self.best_val_auroc = 0
        self.best_epoch_idx = 0

    def report(self, train_metrics, val_metrics):
        print('Epoch [{:03d}/{:03d}] | Time {:5d}s | Loss_G: {:.4f} | '
              'Loss_D: {:.4f} | Val AUROC: {:.2f}\n'.format(
                  train_metrics['epoch_idx'], self.config.optimizer.num_epochs,
                  int(time.time() - self.begin_time),
                  train_metrics['G_losses'][-1], train_metrics['D_losses'][-1],
                  val_metrics['auroc']),
              flush=True)

    def save_model(self, net, val_metrics):
        netG = net['netG']
        netD = net['netD']
        epoch_idx = val_metrics['epoch_idx']

        try:
            netG_wts = copy.deepcopy(netG.module.state_dict())
            netD_wts = copy.deepcopy(netD.module.state_dict())
        except AttributeError:
            netG_wts = copy.deepcopy(netG.state_dict())
            netD_wts = copy.deepcopy(netD.state_dict())

        if self.config.recorder.save_all_models:
            save_pth = os.path.join(self.save_dir,
                                    'epoch-{}_GNet.ckpt'.format(epoch_idx))
            torch.save(netG_wts, save_pth)
            save_pth = os.path.join(self.save_dir,
                                    'epoch-{}_DNet.ckpt'.format(epoch_idx))
            torch.save(netD_wts, save_pth)

        if val_metrics['auroc'] >= self.best_val_auroc:
            self.best_epoch_idx = epoch_idx
            self.best_val_auroc = val_metrics['auroc']

            torch.save(netG_wts, os.path.join(self.output_dir,
                                              'best_GNet.ckpt'))
            torch.save(netD_wts, os.path.join(self.output_dir,
                                              'best_DNet.ckpt'))

    def summary(self):
        print('Training Completed! '
              'Best val AUROC on netD: {:.6f} '
              'at epoch {:d}'.format(self.best_val_auroc, self.best_epoch_idx),
              flush=True)

openood/recorders/opengan_recorder.py
kai422/SCALE
import os
import time
from pathlib import Path

import torch

from .base_recorder import BaseRecorder


class Rd4adRecorder(BaseRecorder):
    def __init__(self, config) -> None:
        super(Rd4adRecorder, self).__init__(config)

        self.best_epoch_idx = 0
        self.best_result = 0

        self.begin_time = time.time()

    def report(self, train_metrics, test_metrics):
        print('Epoch {:03d} | Time {:5d}s | Train Loss {:.4f} | '
              'Auroc {:.4f}\n'.format(train_metrics['epoch_idx'],
                                      int(time.time() - self.begin_time),
                                      train_metrics['loss'],
                                      100.0 * test_metrics['image_auroc']),
              flush=True)

    def save_model(self, net, test_metrics):
        if self.config.recorder.save_all_models:
            torch.save(
                {
                    'bn': net['bn'].state_dict(),
                    'decoder': net['decoder'].state_dict()
                },
                os.path.join(
                    self.output_dir,
                    'model_epoch{}.ckpt'.format(test_metrics['epoch_idx'])))

        # enter only if lower loss occurs
        if test_metrics['image_auroc'] >= self.best_result:

            # delete the depreciated best model
            old_fname1 = 'bn_best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)
            old_fname2 = 'decoder_best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)

            old_pth1 = os.path.join(self.output_dir, old_fname1)
            old_pth2 = os.path.join(self.output_dir, old_fname2)
            Path(old_pth1).unlink(missing_ok=True)
            Path(old_pth2).unlink(missing_ok=True)
            # update the best model
            self.best_epoch_idx = test_metrics['epoch_idx']
            self.best_result = test_metrics['image_auroc']
            torch.save({'bn': net['bn'].state_dict()},
                       os.path.join(self.output_dir, 'bn_best.ckpt'))
            torch.save({'decoder': net['decoder'].state_dict()},
                       os.path.join(self.output_dir, 'decoder_best.ckpt'))
            save_fname1 = 'bn_best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)
            save_pth1 = os.path.join(self.output_dir, save_fname1)
            save_fname2 = 'decoder_best_epoch{}_auroc{:.4f}.ckpt'.format(
                self.best_epoch_idx, self.best_result)
            save_pth2 = os.path.join(self.output_dir, save_fname2)
            torch.save({'bn': net['bn'].state_dict()}, save_pth1)
            torch.save({'decoder': net['decoder'].state_dict()}, save_pth2)

    def summary(self):
        print('Training Completed!\n '
              'Best Auroc: {:.4f} at epoch {:d}\n'.format(
                  100.0 * self.best_result, self.best_epoch_idx),
              flush=True)

openood/recorders/rd4ad_recorder.py
kai422/SCALE
from openood.utils import Config

from .ad_recorder import ADRecorder
from .arpl_recorder import ARPLRecorder
from .base_recorder import BaseRecorder
from .cider_recorder import CiderRecorder
from .cutpaste_recorder import CutpasteRecorder
from .draem_recorder import DRAEMRecorder
from .dsvdd_recorder import DCAERecorder, DSVDDRecorder
from .kdad_recorder import KdadRecorder
from .opengan_recorder import OpenGanRecorder
from .rd4ad_recorder import Rd4adRecorder


def get_recorder(config: Config):
    recorders = {
        'base': BaseRecorder,
        'cider': CiderRecorder,
        'draem': DRAEMRecorder,
        'opengan': OpenGanRecorder,
        'dcae': DCAERecorder,
        'dsvdd': DSVDDRecorder,
        'kdad': KdadRecorder,
        'arpl': ARPLRecorder,
        'cutpaste': CutpasteRecorder,
        'ad': ADRecorder,
        'rd4ad': Rd4adRecorder,
    }

    return recorders[config.recorder.name](config)

openood/recorders/utils.py
kai422/SCALE
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.utils import Config

from .lr_scheduler import cosine_annealing


class ARPLGANTrainer:
    def __init__(self, net: dict, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net['netF']
        self.netG = net['netG']
        self.netD = net['netD']
        self.train_loader = train_loader
        self.config = config
        self.criterion = net['criterion']

        self.fixed_noise = torch.FloatTensor(64, config.network.nz, 1,
                                             1).normal_(0, 1).cuda()
        self.criterionD = nn.BCELoss()

        params_list = [{
            'params': self.net.parameters()
        }, {
            'params': self.criterion.parameters()
        }]

        self.optimizer = torch.optim.SGD(
            params_list,
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

        self.optimizerD = torch.optim.Adam(self.netD.parameters(),
                                           lr=config.optimizer.gan_lr,
                                           betas=(0.5, 0.999))
        self.optimizerG = torch.optim.Adam(self.netG.parameters(),
                                           lr=config.optimizer.gan_lr,
                                           betas=(0.5, 0.999))

    def train_epoch(self, epoch_idx):
        self.net.train()
        self.netD.train()
        self.netG.train()

        loss_avg, lossG_avg, lossD_avg = 0.0, 0.0, 0.0
        train_dataiter = iter(self.train_loader)

        real_label, fake_label = 1, 0
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()
            gan_target = torch.FloatTensor(target.size()).fill_(0).cuda()

            noise = torch.FloatTensor(
                data.size(0), self.config.network.nz, self.config.network.ns,
                self.config.network.ns).normal_(0, 1).cuda()
            noise = noise.cuda()
            noise = Variable(noise)
            fake = self.netG(noise)

            ###########################
            # (1) Update D network    #
            ###########################
            # train with real
            gan_target.fill_(real_label)
            targetv = Variable(gan_target)
            self.optimizerD.zero_grad()
            output = self.netD(data)
            errD_real = self.criterionD(output, targetv)
            errD_real.backward()

            # train with fake
            targetv = Variable(gan_target.fill_(fake_label))
            output = self.netD(fake.detach())
            errD_fake = self.criterionD(output, targetv)
            errD_fake.backward()
            errD = errD_real + errD_fake
            self.optimizerD.step()

            ###########################
            # (2) Update G network    #
            ###########################
            self.optimizerG.zero_grad()
            # Original GAN loss
            targetv = Variable(gan_target.fill_(real_label))
            output = self.netD(fake)
            errG = self.criterionD(output, targetv)

            # minimize the true distribution
            _, feat = self.net(
                fake, True,
                1 * torch.ones(data.shape[0], dtype=torch.long).cuda())
            errG_F = self.criterion.fake_loss(feat).mean()
            generator_loss = errG + self.config.loss.beta * errG_F
            generator_loss.backward()
            self.optimizerG.step()

            ###########################
            # (3) Update classifier   #
            ###########################
            # cross entropy loss
            self.optimizer.zero_grad()
            _, feat = self.net(
                data, True,
                0 * torch.ones(data.shape[0], dtype=torch.long).cuda())
            _, loss = self.criterion(feat, target)

            # KL divergence
            noise = torch.FloatTensor(
                data.size(0), self.config.network.nz, self.config.network.ns,
                self.config.network.ns).normal_(0, 1).cuda()
            noise = Variable(noise)
            fake = self.netG(noise)
            _, feat = self.net(
                fake, True,
                1 * torch.ones(data.shape[0], dtype=torch.long).cuda())
            F_loss_fake = self.criterion.fake_loss(feat).mean()
            total_loss = loss + self.config.loss.beta * F_loss_fake
            total_loss.backward()
            self.optimizer.step()

            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(total_loss) * 0.2
                lossG_avg = lossG_avg * 0.8 + float(generator_loss) * 0.2
                lossD_avg = lossD_avg * 0.8 + float(errD) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg
        metrics['lossG'] = lossG_avg
        metrics['lossD'] = lossD_avg

        return {
            'netG': self.netG,
            'netD': self.netD,
            'netF': self.net,
            'criterion': self.criterion
        }, metrics

openood/trainers/arpl_gan_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class ARPLTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net['netF']
        self.train_loader = train_loader
        self.config = config
        self.criterion = net['criterion']

        params_list = [{
            'params': self.net.parameters()
        }, {
            'params': self.criterion.parameters()
        }]

        self.optimizer = torch.optim.SGD(
            params_list,
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            _, feat = self.net(data, return_feature=True)
            logits, loss = self.criterion(feat, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return {'netF': self.net, 'criterion': self.criterion}, metrics

openood/trainers/arpl_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class AugMixTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.lam = config.trainer.trainer_args.lam
        self.jsd = config.trainer.trainer_args.jsd

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            target = batch['label'].cuda()

            if self.jsd:
                orig_data = batch['data'].cuda()
                aug1_data = batch['data_aug1'].cuda()
                aug2_data = batch['data_aug2'].cuda()
                data = torch.cat([orig_data, aug1_data, aug2_data])

                # forward
                logits_all = self.net(data)
                logits_clean, logits_aug1, logits_aug2 = torch.split(
                    logits_all, orig_data.size(0))

                # Cross-entropy is only computed on clean images
                loss = F.cross_entropy(logits_clean, target)

                p_clean, p_aug1, p_aug2 = \
                    F.softmax(logits_clean, dim=1), \
                    F.softmax(logits_aug1, dim=1), \
                    F.softmax(logits_aug2, dim=1)

                # Clamp mixture distribution to avoid exploding KL divergence
                p_mixture = torch.clamp((p_clean + p_aug1 + p_aug2) / 3., 1e-7,
                                        1).log()
                loss += self.lam * (
                    F.kl_div(p_mixture, p_clean, reduction='batchmean') +
                    F.kl_div(p_mixture, p_aug1, reduction='batchmean') +
                    F.kl_div(p_mixture, p_aug2, reduction='batchmean')) / 3.
            else:
                data = batch['data'].cuda()

                # forward
                logits = self.net(data)
                loss = F.cross_entropy(logits, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced

openood/trainers/augmix_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class BaseTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            logits_classifier = self.net(data)
            loss = F.cross_entropy(logits_classifier, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced

openood/trainers/base_trainer.py
kai422/SCALE
import math
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config


class CIDERTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 val_loader: DataLoader, config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        if 'imagenet' in self.config.dataset.name:
            try:
                for name, p in self.net.backbone.named_parameters():
                    if not name.startswith('layer4'):
                        p.requires_grad = False
            except AttributeError:
                for name, p in self.net.module.backbone.named_parameters():
                    if not name.startswith('layer4'):
                        p.requires_grad = False

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        if config.dataset.train.batch_size \
                * config.num_gpus * config.num_machines > 256:
            config.optimizer.warm = True

        if config.optimizer.warm:
            self.warmup_from = 0.001
            self.warm_epochs = 10
            if config.optimizer.cosine:
                eta_min = config.optimizer.lr * \
                    (config.optimizer.lr_decay_rate**3)
                self.warmup_to = eta_min + (config.optimizer.lr - eta_min) * (
                    1 + math.cos(math.pi * self.warm_epochs /
                                 config.optimizer.num_epochs)) / 2
            else:
                self.warmup_to = config.optimizer.lr

        self.criterion_comp = CompLoss(
            config.dataset.num_classes,
            temperature=config.trainer.trainer_args.temp).cuda()
        # V2: EMA style prototypes
        self.criterion_dis = DisLoss(
            config.dataset.num_classes,
            config.network.feat_dim,
            config.trainer.trainer_args.proto_m,
            self.net,
            val_loader,
            temperature=config.trainer.trainer_args.temp).cuda()

    def train_epoch(self, epoch_idx):
        adjust_learning_rate(self.config, self.optimizer, epoch_idx - 1)

        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            warmup_learning_rate(self.config, self.warm_epochs,
                                 self.warmup_from,
                                 self.warmup_to, epoch_idx - 1, train_step,
                                 len(train_dataiter), self.optimizer)

            batch = next(train_dataiter)
            data = batch['data']
            target = batch['label']

            data = torch.cat([data[0], data[1]], dim=0).cuda()
            target = target.repeat(2).cuda()

            # forward
            features = self.net(data)
            dis_loss = self.criterion_dis(features, target)  # V2: EMA style
            comp_loss = self.criterion_comp(features,
                                            self.criterion_dis.prototypes,
                                            target)
            loss = self.config.trainer.trainer_args.w * comp_loss + dis_loss

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced


def adjust_learning_rate(config, optimizer, epoch):
    lr = config.optimizer.lr
    if config.optimizer.cosine:
        eta_min = lr * (config.optimizer.lr_decay_rate**3)
        lr = eta_min + (lr - eta_min) * (
            1 + math.cos(math.pi * epoch / config.optimizer.num_epochs)) / 2
    else:
        steps = np.sum(epoch > np.asarray(config.optimizer.lr_decay_epochs))
        if steps > 0:
            lr = lr * (config.optimizer.lr_decay_rate**steps)

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def warmup_learning_rate(config, warm_epochs, warmup_from, warmup_to, epoch,
                         batch_id, total_batches, optimizer):
    if config.optimizer.warm and epoch <= warm_epochs:
        p = (batch_id + (epoch - 1) * total_batches) / \
            (warm_epochs * total_batches)
        lr = warmup_from + p * (warmup_to - warmup_from)

        for param_group in optimizer.param_groups:
            param_group['lr'] = lr


class CompLoss(nn.Module):
    """Compactness Loss with class-conditional prototypes."""
    def __init__(self, n_cls, temperature=0.07, base_temperature=0.07):
        super(CompLoss, self).__init__()
        self.n_cls = n_cls
        self.temperature = temperature
        self.base_temperature = base_temperature

    def forward(self, features, prototypes, labels):
        prototypes = F.normalize(prototypes, dim=1)
        proxy_labels = torch.arange(0, self.n_cls).cuda()
        labels = labels.contiguous().view(-1, 1)
        mask = torch.eq(labels, proxy_labels.T).float().cuda()  # bz, cls

        # compute logits
        feat_dot_prototype = torch.div(torch.matmul(features, prototypes.T),
                                       self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(feat_dot_prototype, dim=1, keepdim=True)
        logits = feat_dot_prototype - logits_max.detach()

        # compute log_prob
        exp_logits = torch.exp(logits)
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1)

        # loss
        loss = -(self.temperature /
                 self.base_temperature) * mean_log_prob_pos.mean()
        return loss


class DisLoss(nn.Module):
    """Dispersion Loss with EMA prototypes."""
    def __init__(self,
                 n_cls,
                 feat_dim,
                 proto_m,
                 model,
                 loader,
                 temperature=0.1,
                 base_temperature=0.1):
        super(DisLoss, self).__init__()
        self.n_cls = n_cls
        self.feat_dim = feat_dim
        self.proto_m = proto_m

        self.temperature = temperature
        self.base_temperature = base_temperature
        self.register_buffer('prototypes',
                             torch.zeros(self.n_cls, self.feat_dim))
        self.model = model
        self.loader = loader
        self.init_class_prototypes()

    def forward(self, features, labels):
        prototypes = self.prototypes
        num_cls = self.n_cls
        for j in range(len(features)):
            prototypes[labels[j].item()] = F.normalize(
                prototypes[labels[j].item()] * self.proto_m + features[j] *
                (1 - self.proto_m),
                dim=0)
        self.prototypes = prototypes.detach()
        labels = torch.arange(0, num_cls).cuda()
        labels = labels.contiguous().view(-1, 1)
        labels = labels.contiguous().view(-1, 1)

        mask = (1 - torch.eq(labels, labels.T).float()).cuda()

        logits = torch.div(torch.matmul(prototypes, prototypes.T),
                           self.temperature)

        logits_mask = torch.scatter(torch.ones_like(mask), 1,
                                    torch.arange(num_cls).view(-1, 1).cuda(),
                                    0)
        mask = mask * logits_mask
        mean_prob_neg = torch.log(
            (mask * torch.exp(logits)).sum(1) / mask.sum(1))
        mean_prob_neg = mean_prob_neg[~torch.isnan(mean_prob_neg)]
        loss = self.temperature / self.base_temperature * mean_prob_neg.mean()
        return loss

    def init_class_prototypes(self):
        """Initialize class prototypes."""
        self.model.eval()
        start = time.time()
        prototype_counts = [0] * self.n_cls
        with torch.no_grad():
            prototypes = torch.zeros(self.n_cls, self.feat_dim).cuda()
            for i, batch in enumerate(self.loader):
                input = batch['data']
                target = batch['label']
                input, target = input.cuda(), target.cuda()
                features = self.model(input)
                for j, feature in enumerate(features):
                    prototypes[target[j].item()] += feature
                    prototype_counts[target[j].item()] += 1
            for cls in range(self.n_cls):
                prototypes[cls] /= prototype_counts[cls]
            # measure elapsed time
            duration = time.time() - start
            print(f'Time to initialize prototypes: {duration:.3f}')
            prototypes = F.normalize(prototypes, dim=1)
            self.prototypes = prototypes

openood/trainers/cider_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config
from .lr_scheduler import cosine_annealing


class ConfBranchTrainer:
    def __init__(self, net, train_loader, config: Config) -> None:
        self.train_loader = train_loader
        self.config = config
        self.net = net
        self.prediction_criterion = nn.NLLLoss().cuda()
        self.optimizer = torch.optim.SGD(
            net.parameters(),
            lr=config.optimizer['lr'],
            momentum=config.optimizer['momentum'],
            nesterov=config.optimizer['nesterov'],
            weight_decay=config.optimizer['weight_decay'])
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )
        self.lmbda = self.config.trainer['lmbda']

    def train_epoch(self, epoch_idx):
        self.net.train()
        correct_count = 0.
        total = 0.
        accuracy = 0.
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}'.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            images = Variable(batch['data']).cuda()
            labels = Variable(batch['label']).cuda()
            labels_onehot = Variable(
                encode_onehot(labels, self.config.num_classes))
            self.net.zero_grad()

            pred_original, confidence = self.net(images,
                                                 return_confidence=True)
            pred_original = F.softmax(pred_original, dim=-1)
            confidence = torch.sigmoid(confidence)
            eps = self.config.trainer['eps']
            pred_original = torch.clamp(pred_original, 0. + eps, 1. - eps)
            confidence = torch.clamp(confidence, 0. + eps, 1. - eps)

            if not self.config.baseline:
                # Randomly set half of the confidences to 1 (i.e. no hints)
                b = Variable(
                    torch.bernoulli(
                        torch.Tensor(confidence.size()).uniform_(0,
                                                                 1))).cuda()
                conf = confidence * b + (1 - b)
                pred_new = pred_original * conf.expand_as(
                    pred_original) + labels_onehot * (
                        1 - conf.expand_as(labels_onehot))
                pred_new = torch.log(pred_new)
            else:
                pred_new = torch.log(pred_original)

            xentropy_loss = self.prediction_criterion(pred_new, labels)
            confidence_loss = torch.mean(-torch.log(confidence))

            if self.config.baseline:
                total_loss = xentropy_loss
            else:
                total_loss = xentropy_loss + (self.lmbda * confidence_loss)

                if self.config.trainer['budget'] > confidence_loss.item():
                    self.lmbda = self.lmbda / 1.01
                elif self.config.trainer['budget'] <= confidence_loss.item():
                    self.lmbda = self.lmbda / 0.99

            total_loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            pred_idx = torch.max(pred_original.data, 1)[1]
            total += labels.size(0)
            correct_count += (pred_idx == labels.data).sum()
            accuracy = correct_count / total

        metrics = {}
        metrics['train_acc'] = accuracy
        metrics['loss'] = total_loss
        metrics['epoch_idx'] = epoch_idx
        return self.net, metrics


def encode_onehot(labels, n_classes):
    onehot = torch.FloatTensor(labels.size()[0],
                               n_classes)  # batchsize * num of class
    labels = labels.data
    if labels.is_cuda:
        onehot = onehot.cuda()
    onehot.zero_()
    onehot.scatter_(1, labels.view(-1, 1), 1)
    return onehot

openood/trainers/conf_branch_trainer.py
kai422/SCALE
import math
import numbers

import diffdist.functional as distops
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.autograd import Function
from torch.optim.lr_scheduler import ReduceLROnPlateau, _LRScheduler
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config


class CSITrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:
        self.net = net['backbone']
        self.train_loader = train_loader
        self.config = config
        self.mode = config.mode

        if self.config.num_gpus > 1:
            self.dummy_net = net['dummy_net'].module
        else:
            self.dummy_net = net['dummy_net']
        self.dummy_net.cpu()

        self.simclr_aug = get_simclr_augmentation(
            config, image_size=config.dataset.image_size).cuda()
        self.linear = net['linear']
        self.linear_optim = torch.optim.Adam(
            self.linear.parameters(),
            lr=1e-3,
            betas=(.9, .999),
            weight_decay=config.optimizer.weight_decay)
        self.criterion = nn.CrossEntropyLoss().cuda()
        self.hflip = HorizontalFlipLayer().cuda()

        self.simclr_layer = net['simclr_layer']
        self.rotation_linear = net['shift_cls_layer']
        self.joint_linear = net['joint_distribution_layer']

        if 'step1' in self.mode:
            self.optimizer = optim.SGD(
                list(self.net.parameters()) +
                list(self.simclr_layer.parameters()),
                lr=config.optimizer.lr,
                momentum=0.9,
                weight_decay=config.optimizer.weight_decay)
            self.scheduler = lr_scheduler.CosineAnnealingLR(
                self.optimizer, config.optimizer.num_epochs)
            self.scheduler_warmup = GradualWarmupScheduler(
                self.optimizer,
                multiplier=10.0,
                total_epoch=config.optimizer.warmup,
                after_scheduler=self.scheduler)
        else:
            milestones = [
                int(0.6 * config.optimizer.num_epochs),
                int(0.75 * config.optimizer.num_epochs),
                int(0.9 * config.optimizer.num_epochs)
            ]

            self.linear_optim = torch.optim.Adam(
                self.linear.parameters(),
                lr=1e-3,
                betas=(.9, .999),
                weight_decay=config.optimizer.weight_decay)
            self.linear_scheduler = lr_scheduler.MultiStepLR(
                self.linear_optim, gamma=0.1, milestones=milestones)

            self.rotation_linear_optim = torch.optim.SGD(
                self.rotation_linear.parameters(),
                lr=1e-1,
                weight_decay=config.optimizer.weight_decay)
            self.rot_scheduler = lr_scheduler.MultiStepLR(
                self.rotation_linear_optim, gamma=0.1, milestones=milestones)

            self.joint_linear_optim = torch.optim.SGD(
                self.joint_linear.parameters(),
                lr=1e-1,
                weight_decay=config.optimizer.weight_decay)
            self.joint_scheduler = lr_scheduler.MultiStepLR(
                self.joint_linear_optim, gamma=0.1, milestones=milestones)

    def train_epoch(self, epoch_idx):
        if 'step1' in self.mode:
            return self.train_sup_epoch(epoch_idx)
        else:
            return self.train_suplinear_epoch(epoch_idx)

    def train_sup_epoch(self, epoch_idx):
        self.net.train()
        train_dataiter = iter(self.train_loader)

        n = 0
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            images = batch['data'].cuda()
            labels = batch['label'].cuda()

            batch_size = images.size(0)
            images1, images2 = self.hflip(images.repeat(2, 1, 1,
                                                        1)).chunk(2)  # hflip
            images1 = torch.cat(
                [torch.rot90(images1, rot, (2, 3)) for rot in range(4)])  # 4B
            images2 = torch.cat(
                [torch.rot90(images2, rot, (2, 3)) for rot in range(4)])  # 4B
            images_pair = torch.cat([images1, images2], dim=0)  # 8B

            rot_sim_labels = torch.cat([
                labels + self.config.dataset.num_classes * i for i in range(4)
            ],
                                       dim=0)

            images_pair = self.simclr_aug(images_pair)  # simclr augment
            _, features = self.net(images_pair, return_feature=True)

            simclr_outputs = self.simclr_layer(features)
            simclr = normalize(simclr_outputs)  # normalize
            sim_matrix = get_similarity_matrix(
                simclr, multi_gpu=self.config.num_gpus > 1)
            loss_sim = Supervised_NT_xent(
                sim_matrix,
                labels=rot_sim_labels,
                temperature=self.config.temperature,
                multi_gpu=self.config.num_gpus > 1) * self.config.sim_lambda

            # total loss
            loss = loss_sim

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            self.scheduler.step(epoch_idx - 1 + n / len(self.train_loader))
            # lr = self.optimizer.param_groups[0]['lr']

            # Post-processing stuffs
            penul_1 = features[:batch_size]
            penul_2 = features[4 * batch_size:5 * batch_size]
            features = torch.cat([penul_1,
                                  penul_2])  # only use original rotation

            # Linear evaluation
            outputs_linear_eval = self.linear(features.detach())
            loss_linear = self.criterion(outputs_linear_eval, labels.repeat(2))

            self.linear_optim.zero_grad()
            loss_linear.backward()
            self.linear_optim.step()

            n = n + 1

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss

        if self.config.num_gpus > 1:
            self.dummy_net.backbone.load_state_dict(
                self.net.module.state_dict())
            self.dummy_net.linear.load_state_dict(
                self.linear.module.state_dict())
            self.dummy_net.simclr_layer.load_state_dict(
                self.simclr_layer.module.state_dict())
            self.dummy_net.joint_distribution_layer.load_state_dict(
                self.joint_distribution_layer.module.state_dict())
            self.dummy_net.shift_cls_layer.load_state_dict(
                self.shift_cls_layer.module.state_dict())
        else:
            self.dummy_net.backbone.load_state_dict(self.net.state_dict())
            self.dummy_net.linear.load_state_dict(self.linear.state_dict())
            self.dummy_net.simclr_layer.load_state_dict(
                self.simclr_layer.state_dict())
            self.dummy_net.joint_distribution_layer.load_state_dict(
                self.joint_distribution_layer.state_dict())
            self.dummy_net.shift_cls_layer.load_state_dict(
                self.shift_cls_layer.state_dict())

        return self.dummy_net, metrics

    def train_suplinear_epoch(self, epoch_idx):
        self.net.train()
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            self.net.eval()
            batch = next(train_dataiter)
            images = batch['data'].cuda()
            labels = batch['label'].cuda()

            batch_size = images.size(0)
            images = self.hflip(images)
            images = torch.cat(
                [torch.rot90(images, rot, (2, 3)) for rot in range(4)])  # 4B
            rot_labels = torch.cat(
                [torch.ones_like(labels) * k for k in range(4)], 0)  # B -> 4B
            joint_labels = torch.cat([
                labels + self.config.dataset.num_classes * i for i in range(4)
            ],
                                     dim=0)

            images = self.simclr_aug(images)  # simclr augmentation
            _, features = self.net(images, return_feature=True)
            penultimate = features.detach()

            outputs = self.linear(
                penultimate[0:batch_size]
            )  # only use 0 degree samples for linear eval
            outputs_rot = self.rotation_linear(penultimate)
            outputs_joint = self.joint_linear(penultimate)

            loss_ce = self.criterion(outputs, labels)
            loss_rot = self.criterion(outputs_rot, rot_labels)
            loss_joint = self.criterion(outputs_joint, joint_labels)

            # CE loss
            self.linear_optim.zero_grad()
            loss_ce.backward()
            self.linear_optim.step()

            # Rot loss
            self.rotation_linear_optim.zero_grad()
            loss_rot.backward()
            self.rotation_linear_optim.step()

            # Joint loss
            self.joint_linear_optim.zero_grad()
            loss_joint.backward()
            self.joint_linear_optim.step()

            # optimizer learning rate
            # lr = self.linear_optim.param_groups[0]['lr']

        self.linear_scheduler.step()
        self.rot_scheduler.step()
        self.joint_scheduler.step()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_ce + loss_rot + loss_joint

        if self.config.num_gpus > 1:
            self.dummy_net.backbone.load_state_dict(
                self.net.module.state_dict())
            self.dummy_net.linear.load_state_dict(
                self.linear.module.state_dict())
            self.dummy_net.simclr_layer.load_state_dict(
                self.simclr_layer.module.state_dict())
            self.dummy_net.joint_distribution_layer.load_state_dict(
                self.joint_distribution_layer.module.state_dict())
            self.dummy_net.shift_cls_layer.load_state_dict(
                self.shift_cls_layer.module.state_dict())
        else:
            self.dummy_net.backbone.load_state_dict(self.net.state_dict())
            self.dummy_net.linear.load_state_dict(self.linear.state_dict())
            self.dummy_net.simclr_layer.load_state_dict(
                self.simclr_layer.state_dict())
            self.dummy_net.joint_distribution_layer.load_state_dict(
                self.joint_distribution_layer.state_dict())
            self.dummy_net.shift_cls_layer.load_state_dict(
                self.shift_cls_layer.state_dict())

        return self.dummy_net, metrics


def get_similarity_matrix(outputs, chunk=2, multi_gpu=False):
    """Compute similarity matrix.

    - outputs: (B', d) tensor for B' = B * chunk
    - sim_matrix: (B', B') tensor
    """

    if multi_gpu:
        outputs_gathered = []
        for out in outputs.chunk(chunk):
            gather_t = [
                torch.empty_like(out) for _ in range(dist.get_world_size())
            ]
            gather_t = torch.cat(distops.all_gather(gather_t, out))
            outputs_gathered.append(gather_t)
        outputs = torch.cat(outputs_gathered)

    sim_matrix = torch.mm(outputs, outputs.t())  # (B', d), (d, B') -> (B', B')

    return sim_matrix


def Supervised_NT_xent(sim_matrix,
                       labels,
                       temperature=0.5,
                       chunk=2,
                       eps=1e-8,
                       multi_gpu=False):
    """Compute NT_xent loss.

    - sim_matrix: (B', B') tensor for B' = B * chunk (first 2B are pos samples)
    """

    device = sim_matrix.device

    if multi_gpu:
        gather_t = [
            torch.empty_like(labels) for _ in range(dist.get_world_size())
        ]
        labels = torch.cat(distops.all_gather(gather_t, labels))
    labels = labels.repeat(2)

    logits_max, _ = torch.max(sim_matrix, dim=1, keepdim=True)
    sim_matrix = sim_matrix - logits_max.detach()

    B = sim_matrix.size(0) // chunk  # B = B' / chunk

    eye = torch.eye(B * chunk).to(device)  # (B', B')
    sim_matrix = torch.exp(sim_matrix / temperature) * (1 - eye
                                                        )  # remove diagonal

    denom = torch.sum(sim_matrix, dim=1, keepdim=True)
    sim_matrix = -torch.log(sim_matrix / (denom + eps) + eps)  # loss matrix

    labels = labels.contiguous().view(-1, 1)
    Mask = torch.eq(labels, labels.t()).float().to(device)
    # Mask = eye * torch.stack([labels == labels[i]
    # for i in range(labels.size(0))]).float().to(device)
    Mask = Mask / (Mask.sum(dim=1, keepdim=True) + eps)

    loss = torch.sum(Mask * sim_matrix) / (2 * B)

    return loss


def normalize(x, dim=1, eps=1e-8):
    return x / (x.norm(dim=dim, keepdim=True) + eps)


def get_simclr_augmentation(config, image_size):

    # parameter for resizecrop
    resize_scale = (config.resize_factor, 1.0)  # resize scaling factor
    if config.resize_fix:  # if resize_fix is True, use same scale
        resize_scale = (config.resize_factor, config.resize_factor)

    # Align augmentation
    color_jitter = ColorJitterLayer(brightness=0.4,
                                    contrast=0.4,
                                    saturation=0.4,
                                    hue=0.1,
                                    p=0.8)
    color_gray = RandomColorGrayLayer(p=0.2)
    resize_crop = RandomResizedCropLayer(scale=resize_scale, size=image_size)

    # Transform define #
    if config.dataset.name == 'imagenet':
        # Using RandomResizedCrop at PIL transform
        transform = nn.Sequential(
            color_jitter,
            color_gray,
        )
    else:
        transform = nn.Sequential(
            color_jitter,
            color_gray,
            resize_crop,
        )

    return transform


# ----------Warmup Scheduler----------
class GradualWarmupScheduler(_LRScheduler):
    """Gradually warm-up(increasing) learning rate in optimizer. Proposed in
    'Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour'.

    Args:
        optimizer (Optimizer): Wrapped optimizer.
        multiplier: target learning rate = base lr * multiplier
        if multiplier > 1.0. if multiplier = 1.0,
        lr starts from 0 and ends up with the base_lr.
        total_epoch: target learning rate is reached at total_epoch, gradually
        after_scheduler: after target_epoch,
        use this scheduler (eg. ReduceLROnPlateau)
    """
    def __init__(self,
                 optimizer,
                 multiplier,
                 total_epoch,
                 after_scheduler=None):
        self.multiplier = multiplier
        if self.multiplier < 1.:
            raise ValueError(
                'multiplier should be greater than or equal to 1.')
        self.total_epoch = total_epoch
        self.after_scheduler = after_scheduler
        self.finished = False
        super(GradualWarmupScheduler, self).__init__(optimizer)

    def get_lr(self):
        if self.last_epoch > self.total_epoch:
            if self.after_scheduler:
                if not self.finished:
                    self.after_scheduler.base_lrs = [
                        base_lr * self.multiplier for base_lr in self.base_lrs
                    ]
                    self.finished = True
                return self.after_scheduler.get_lr()
            return [base_lr * self.multiplier for base_lr in self.base_lrs]

        if self.multiplier == 1.0:
            return [
                base_lr * (float(self.last_epoch) / self.total_epoch)
                for base_lr in self.base_lrs
            ]
        else:
            return [
                base_lr *
                ((self.multiplier - 1.) * self.last_epoch / self.total_epoch +
                 1.) for base_lr in self.base_lrs
            ]

    def step_ReduceLROnPlateau(self, metrics, epoch=None):
        if epoch is None:
            epoch = self.last_epoch + 1
        self.last_epoch = epoch if epoch != 0 else 1
        # ReduceLROnPlateau is called at the end of epoch,
        # whereas others are called at beginning
        if self.last_epoch <= self.total_epoch:
            warmup_lr = [
                base_lr *
                ((self.multiplier - 1.) * self.last_epoch / self.total_epoch +
                 1.) for base_lr in self.base_lrs
            ]
            for param_group, lr in zip(self.optimizer.param_groups, warmup_lr):
                param_group['lr'] = lr
        else:
            if epoch is None:
                self.after_scheduler.step(metrics, None)
            else:
                self.after_scheduler.step(metrics, epoch - self.total_epoch)

    def step(self, epoch=None, metrics=None):
        if type(self.after_scheduler) != ReduceLROnPlateau:
            if self.finished and self.after_scheduler:
                if epoch is None:
                    self.after_scheduler.step(None)
                else:
                    self.after_scheduler.step(epoch - self.total_epoch)
            else:
                return super(GradualWarmupScheduler, self).step(epoch)
        else:
            self.step_ReduceLROnPlateau(metrics, epoch)


# ----------transform layers----------
if torch.__version__ >= '1.4.0':
    kwargs = {'align_corners': False}
else:
    kwargs = {}


def rgb2hsv(rgb):
    """Convert a 4-d RGB tensor to the HSV counterpart.

    Here, we compute hue using atan2() based on the definition in [1],
    instead of using the common lookup table approach as in [2, 3].
    Those values agree when the angle is a multiple of 30°,
    otherwise they may differ at most ~1.2°.

    References
    [1] https://en.wikipedia.org/wiki/Hue
    [2] https://www.rapidtables.com/convert/color/rgb-to-hsv.html
    [3] https://github.com/scikit-image/scikit-image/
    blob/master/skimage/color/colorconv.py#L212
    """

    r, g, b = rgb[:, 0, :, :], rgb[:, 1, :, :], rgb[:, 2, :, :]

    Cmax = rgb.max(1)[0]
    Cmin = rgb.min(1)[0]
    delta = Cmax - Cmin

    hue = torch.atan2(math.sqrt(3) * (g - b), 2 * r - g - b)
    hue = (hue % (2 * math.pi)) / (2 * math.pi)
    saturate = delta / Cmax
    value = Cmax
    hsv = torch.stack([hue, saturate, value], dim=1)
    hsv[~torch.isfinite(hsv)] = 0.
    return hsv


def hsv2rgb(hsv):
    """Convert a 4-d HSV tensor to the RGB counterpart.

    >>> %timeit hsv2rgb(hsv)
    2.37 ms ± 13.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    >>> %timeit rgb2hsv_fast(rgb)
    298 µs ± 542 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
    >>> torch.allclose(hsv2rgb(hsv), hsv2rgb_fast(hsv), atol=1e-6)
    True

    References
    [1] https://en.wikipedia.org/wiki/HSL_and_HSV#HSV_to_RGB_alternative
    """
    h, s, v = hsv[:, [0]], hsv[:, [1]], hsv[:, [2]]
    c = v * s

    n = hsv.new_tensor([5, 3, 1]).view(3, 1, 1)
    k = (n + h * 6) % 6
    t = torch.min(k, 4 - k)
    t = torch.clamp(t, 0, 1)

    return v - c * t


class RandomResizedCropLayer(nn.Module):
    def __init__(self, size=None, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.)):
        """Inception Crop size (tuple): size of forwarding image (C, W, H)
        scale (tuple): range of size of the origin size cropped ratio (tuple):

        range of aspect ratio of the origin aspect ratio cropped.
        """
        super(RandomResizedCropLayer, self).__init__()

        _eye = torch.eye(2, 3)
        self.size = size
        self.register_buffer('_eye', _eye)
        self.scale = scale
        self.ratio = ratio

    def forward(self, inputs, whbias=None):
        _device = inputs.device
        N = inputs.size(0)
        _theta = self._eye.repeat(N, 1, 1)

        if whbias is None:
            whbias = self._sample_latent(inputs)

        _theta[:, 0, 0] = whbias[:, 0]
        _theta[:, 1, 1] = whbias[:, 1]
        _theta[:, 0, 2] = whbias[:, 2]
        _theta[:, 1, 2] = whbias[:, 3]

        grid = F.affine_grid(_theta, inputs.size(), **kwargs).to(_device)
        output = F.grid_sample(inputs,
                               grid,
                               padding_mode='reflection',
                               **kwargs)

        if self.size is not None:
            output = F.adaptive_avg_pool2d(output, self.size)

        return output

    def _clamp(self, whbias):

        w = whbias[:, 0]
        h = whbias[:, 1]
        w_bias = whbias[:, 2]
        h_bias = whbias[:, 3]

        # Clamp with scale
        w = torch.clamp(w, *self.scale)
        h = torch.clamp(h, *self.scale)

        # Clamp with ratio
        w = self.ratio[0] * h + torch.relu(w - self.ratio[0] * h)
        w = self.ratio[1] * h - torch.relu(self.ratio[1] * h - w)

        # Clamp with bias range: w_bias \in (w - 1, 1 - w),
        # h_bias \in (h - 1, 1 - h)
        w_bias = w - 1 + torch.relu(w_bias - w + 1)
        w_bias = 1 - w - torch.relu(1 - w - w_bias)

        h_bias = h - 1 + torch.relu(h_bias - h + 1)
        h_bias = 1 - h - torch.relu(1 - h - h_bias)

        whbias = torch.stack([w, h, w_bias, h_bias], dim=0).t()

        return whbias

    def _sample_latent(self, inputs):

        _device = inputs.device
        N, _, width, height = inputs.shape

        # N * 10 trial
        area = width * height
        target_area = np.random.uniform(*self.scale, N * 10) * area
        log_ratio = (math.log(self.ratio[0]), math.log(self.ratio[1]))
        aspect_ratio = np.exp(np.random.uniform(*log_ratio, N * 10))

        # If doesn't satisfy ratio condition, then do central crop
        w = np.round(np.sqrt(target_area * aspect_ratio))
        h = np.round(np.sqrt(target_area / aspect_ratio))
        cond = (0 < w) * (w <= width) * (0 < h) * (h <= height)
        w = w[cond]
        h = h[cond]
        cond_len = w.shape[0]
        if cond_len >= N:
            w = w[:N]
            h = h[:N]
        else:
            w = np.concatenate([w, np.ones(N - cond_len) * width])
            h = np.concatenate([h, np.ones(N - cond_len) * height])

        w_bias = np.random.randint(w - width, width - w + 1) / width
        h_bias = np.random.randint(h - height, height - h + 1) / height
        w = w / width
        h = h / height

        whbias = np.column_stack([w, h, w_bias, h_bias])
        whbias = torch.tensor(whbias, device=_device)

        return whbias


class HorizontalFlipRandomCrop(nn.Module):
    def __init__(self, max_range):
        super(HorizontalFlipRandomCrop, self).__init__()
        self.max_range = max_range
        _eye = torch.eye(2, 3)
        self.register_buffer('_eye', _eye)

    def forward(self, input, sign=None, bias=None, rotation=None):
        _device = input.device
        N = input.size(0)
        _theta = self._eye.repeat(N, 1, 1)

        if sign is None:
            sign = torch.bernoulli(torch.ones(N, device=_device) * 0.5) * 2 - 1
        if bias is None:
            bias = torch.empty(
                (N, 2), device=_device).uniform_(-self.max_range,
                                                 self.max_range)
        _theta[:, 0, 0] = sign
        _theta[:, :, 2] = bias

        if rotation is not None:
            _theta[:, 0:2, 0:2] = rotation

        grid = F.affine_grid(_theta, input.size(), **kwargs).to(_device)
        output = F.grid_sample(input,
                               grid,
                               padding_mode='reflection',
                               **kwargs)

        return output

    def _sample_latent(self, N, device=None):
        sign = torch.bernoulli(torch.ones(N, device=device) * 0.5) * 2 - 1
        bias = torch.empty(
            (N, 2), device=device).uniform_(-self.max_range, self.max_range)
        return sign, bias


class Rotation(nn.Module):
    def __init__(self, max_range=4):
        super(Rotation, self).__init__()
        self.max_range = max_range
        self.prob = 0.5

    def forward(self, input, aug_index=None):

        _, _, H, W = input.size()

        if aug_index is None:
            aug_index = np.random.randint(4)

            output = torch.rot90(input, aug_index, (2, 3))

            _prob = input.new_full((input.size(0), ), self.prob)
            _mask = torch.bernoulli(_prob).view(-1, 1, 1, 1)
            output = _mask * input + (1 - _mask) * output

        else:
            aug_index = aug_index % self.max_range
            output = torch.rot90(input, aug_index, (2, 3))

        return output


class CutPerm(nn.Module):
    def __init__(self, max_range=4):
        super(CutPerm, self).__init__()
        self.max_range = max_range
        self.prob = 0.5

    def forward(self, input, aug_index=None):

        _, _, H, W = input.size()

        if aug_index is None:
            aug_index = np.random.randint(4)

            output = self._cutperm(input, aug_index)

            _prob = input.new_full((input.size(0), ), self.prob)
            _mask = torch.bernoulli(_prob).view(-1, 1, 1, 1)
            output = _mask * input + (1 - _mask) * output

        else:
            aug_index = aug_index % self.max_range
            output = self._cutperm(input, aug_index)

        return output

    def _cutperm(self, inputs, aug_index):

        _, _, H, W = inputs.size()
        h_mid = int(H / 2)
        w_mid = int(W / 2)

        jigsaw_h = aug_index // 2
        jigsaw_v = aug_index % 2

        if jigsaw_h == 1:
            inputs = torch.cat(
                (inputs[:, :, h_mid:, :], inputs[:, :, 0:h_mid, :]), dim=2)
        if jigsaw_v == 1:
            inputs = torch.cat(
                (inputs[:, :, :, w_mid:], inputs[:, :, :, 0:w_mid]), dim=3)

        return inputs


class HorizontalFlipLayer(nn.Module):
    def __init__(self):
        """
        img_size : (int, int, int)
            Height and width must be powers of 2.  E.g. (32, 32, 1) or
            (64, 128, 3). Last number indicates number of channels, e.g. 1 for
            grayscale or 3 for RGB
        """
        super(HorizontalFlipLayer, self).__init__()

        _eye = torch.eye(2, 3)
        self.register_buffer('_eye', _eye)

    def forward(self, inputs):
        _device = inputs.device

        N = inputs.size(0)
        _theta = self._eye.repeat(N, 1, 1)
        r_sign = torch.bernoulli(torch.ones(N, device=_device) * 0.5) * 2 - 1
        _theta[:, 0, 0] = r_sign
        grid = F.affine_grid(_theta, inputs.size(), **kwargs).to(_device)
        inputs = F.grid_sample(inputs,
                               grid,
                               padding_mode='reflection',
                               **kwargs)

        return inputs


class RandomColorGrayLayer(nn.Module):
    def __init__(self, p):
        super(RandomColorGrayLayer, self).__init__()
        self.prob = p

        _weight = torch.tensor([[0.299, 0.587, 0.114]])
        self.register_buffer('_weight', _weight.view(1, 3, 1, 1))

    def forward(self, inputs, aug_index=None):

        if aug_index == 0:
            return inputs

        outputs = F.conv2d(inputs, self._weight)
        gray = torch.cat([outputs, outputs, outputs], dim=1)

        if aug_index is None:
            _prob = inputs.new_full((inputs.size(0), ), self.prob)
            _mask = torch.bernoulli(_prob).view(-1, 1, 1, 1)

            gray = inputs * (1 - _mask) + gray * _mask

        return gray


class ColorJitterLayer(nn.Module):
    def __init__(self, p, brightness, contrast, saturation, hue):
        super(ColorJitterLayer, self).__init__()
        self.prob = p
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue,
                                     'hue',
                                     center=0,
                                     bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)

    def _check_input(self,
                     value,
                     name,
                     center=1,
                     bound=(0, float('inf')),
                     clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError(
                    'If {} is a single number, it must be non negative.'.
                    format(name))
            value = [center - value, center + value]
            if clip_first_on_zero:
                value[0] = max(value[0], 0)
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError('{} values should be between {}'.format(
                    name, bound))
        else:
            raise TypeError(
                '{} should be a single number or a list/tuple with length 2.'.
                format(name))

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value

    def adjust_contrast(self, x):
        if self.contrast:
            factor = x.new_empty(x.size(0), 1, 1, 1).uniform_(*self.contrast)
            means = torch.mean(x, dim=[2, 3], keepdim=True)
            x = (x - means) * factor + means
        return torch.clamp(x, 0, 1)

    def adjust_hsv(self, x):
        f_h = x.new_zeros(x.size(0), 1, 1)
        f_s = x.new_ones(x.size(0), 1, 1)
        f_v = x.new_ones(x.size(0), 1, 1)

        if self.hue:
            f_h.uniform_(*self.hue)
        if self.saturation:
            f_s = f_s.uniform_(*self.saturation)
        if self.brightness:
            f_v = f_v.uniform_(*self.brightness)

        return RandomHSVFunction.apply(x, f_h, f_s, f_v)

    def transform(self, inputs):
        # Shuffle transform
        if np.random.rand() > 0.5:
            transforms = [self.adjust_contrast, self.adjust_hsv]
        else:
            transforms = [self.adjust_hsv, self.adjust_contrast]

        for t in transforms:
            inputs = t(inputs)

        return inputs

    def forward(self, inputs):
        _prob = inputs.new_full((inputs.size(0), ), self.prob)
        _mask = torch.bernoulli(_prob).view(-1, 1, 1, 1)
        return inputs * (1 - _mask) + self.transform(inputs) * _mask


class RandomHSVFunction(Function):
    @staticmethod
    def forward(ctx, x, f_h, f_s, f_v):
        # ctx is a context object that can be used to stash information
        # for backward computation
        x = rgb2hsv(x)
        h = x[:, 0, :, :]
        h += (f_h * 255. / 360.)
        h = (h % 1)
        x[:, 0, :, :] = h
        x[:, 1, :, :] = x[:, 1, :, :] * f_s
        x[:, 2, :, :] = x[:, 2, :, :] * f_v
        x = torch.clamp(x, 0, 1)
        x = hsv2rgb(x)
        return x

    @staticmethod
    def backward(ctx, grad_output):
        # We return as many input gradients as there were arguments.
        # Gradients of non-Tensor arguments to forward must be None.
        grad_input = None
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.clone()
        return grad_input, None, None, None


class NormalizeLayer(nn.Module):
    """In order to certify radii in original coordinates rather than
    standardized coordinates, we add the Gaussian noise _before_ standardizing,
    which is why we have standardization be the first layer of the classifier
    rather than as a part of preprocessing as is typical."""
    def __init__(self):
        super(NormalizeLayer, self).__init__()

    def forward(self, inputs):
        return (inputs - 0.5) / 0.5

openood/trainers/csi_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class CutMixTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.args = config.trainer.trainer_args

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # perform cutmix augmentation in a batch
            r = np.random.rand(1)
            if self.args.beta > 0 and r < self.args.cutmix_prob:
                # generate mixed sample
                lam = np.random.beta(self.args.beta, self.args.beta)
                rand_index = torch.randperm(data.size()[0]).cuda()
                target_a = target
                target_b = target[rand_index]
                bbx1, bby1, bbx2, bby2 = rand_bbox(data.size(), lam)
                data[:, :, bbx1:bbx2, bby1:bby2] = data[rand_index, :,
                                                        bbx1:bbx2, bby1:bby2]
                # adjust lambda to exactly match pixel ratio
                lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) /
                           (data.size()[-1] * data.size()[-2]))
                # forward
                logits_classifier = self.net(data)
                loss = F.cross_entropy(
                    logits_classifier, target_a) * lam + F.cross_entropy(
                        logits_classifier, target_b) * (1. - lam)
            else:
                # forward
                logits_classifier = self.net(data)
                loss = F.cross_entropy(logits_classifier, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics


def rand_bbox(size, lam):
    W = size[2]
    H = size[3]
    cut_rat = np.sqrt(1. - lam)
    cut_w = np.int(W * cut_rat)
    cut_h = np.int(H * cut_rat)

    # uniform
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    return bbx1, bby1, bbx2, bby2

openood/trainers/cutmix_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.utils import Config

from .lr_scheduler import cosine_annealing


class CutPasteTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        embeds = []

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            data = torch.cat(batch['data'], 0)
            data = data.cuda()
            y = torch.arange(2)
            y = y.repeat_interleave(len(batch['data'][0]))
            y = y.cuda()
            # forward
            embed, logits_classifier = self.net(data)
            loss = F.cross_entropy(logits_classifier, y)

            embeds.append(embed.cuda())

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        embeds = torch.cat(embeds)
        embeds = torch.nn.functional.normalize(embeds, p=2, dim=1)

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/cutpaste_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.utils import Config

from .lr_scheduler import cosine_annealing


class DropoutTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.p = config.trainer.dropout_p

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            logits_classifier = self.net.forward_with_dropout(data, self.p)
            loss = F.cross_entropy(logits_classifier, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/dropout_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.optim as optim
from tqdm import tqdm

from openood.utils import Config


class AETrainer:
    def __init__(self, net, train_loader, config: Config):
        self.config = config
        self.net = net
        self.train_loader = train_loader
        if config.optimizer.name == 'adam':
            self.optimizer = optim.Adam(
                net.parameters(),
                lr=config.lr,
                weight_decay=config.weight_decay,
                amsgrad=config.optimizer.name == 'amsgrad')
        self.scheduler = optim.lr_scheduler.MultiStepLR(
            self.optimizer, milestones=config.lr_milestones, gamma=0.1)

    def train_epoch(self, epoch_idx):

        self.net.train()
        epoch_loss = 0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d} '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            inputs = batch['data'].cuda()
            self.optimizer.zero_grad()
            outputs = self.net(inputs)
            scores = torch.sum((outputs - inputs)**2,
                               dim=tuple(range(1, outputs.dim())))
            loss = torch.mean(scores)
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()
            epoch_loss += loss.item()
        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = epoch_loss
        return self.net, metrics


class DSVDDTrainer:
    def __init__(self, net, train_loader, config: Config) -> None:
        self.config = config
        self.net = net
        self.train_loader = train_loader
        if config.optimizer.name == 'adam':
            self.optimizer = optim.Adam(
                net.parameters(),
                lr=config.lr,
                weight_decay=config.weight_decay,
                amsgrad=config.optimizer.name == 'amsgrad')
        self.scheduler = optim.lr_scheduler.MultiStepLR(
            self.optimizer, milestones=config.lr_milestones, gamma=0.1)

        if self.config.c == 'None' and self.config.network.name != 'dcae':
            self.config.c = init_center_c(train_loader, net)
        self.c = self.config.c

    def train_epoch(self, epoch_idx):
        self.net.train()
        epoch_loss = 0
        train_dataiter = iter(self.train_loader)
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}'.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            inputs = batch['data'].cuda()
            self.optimizer.zero_grad()
            outputs = self.net(inputs)
            if self.config.network.name != 'dcae':
                scores = torch.sum((outputs - self.c)**2, dim=1)

            # this is for pre-training the dcae network from the original paper
            elif self.config.network.name == 'dcae':
                scores = torch.sum((outputs - inputs)**2,
                                   dim=tuple(range(1, outputs.dim())))
            else:
                raise NotImplementedError
            loss = torch.mean(scores)
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()
            epoch_loss += loss.item()
        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = epoch_loss
        return self.net, metrics


def init_center_c(train_loader, net, eps=0.1):
    """Initialize hypersphere center c as the mean from an initial forward pass
    on the data."""
    n_samples = 0
    first_iter = True
    train_dataiter = iter(train_loader)
    net.eval()
    with torch.no_grad():
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Initialize center',
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            inputs = batch['data'].cuda()
            outputs = net(inputs)
            if first_iter:
                c = torch.zeros(outputs.shape[1]).cuda()
                first_iter = False
            n_samples += outputs.shape[0]
            c += torch.sum(outputs, dim=0)

    c /= n_samples

    # If c_i is too close to 0, set to +-eps.
    # Reason: a zero unit can be trivially matched with zero weights.
    c[(abs(c) < eps) & (c < 0)] = -eps
    c[(abs(c) < eps) & (c > 0)] = eps

    return c


def get_radius(dist: torch.Tensor, nu: float):
    """Optimally solve for radius R via the (1-nu)-quantile of distances."""
    return np.quantile(np.sqrt(dist.clone().data.cpu().numpy()), 1 - nu)

openood/trainers/dsvdd_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class GodinTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        parameters = []
        h_parameters = []
        for name, parameter in net.named_parameters():
            if name in [
                    'h.h.weight', 'h.h.bias', 'module.h.h.weight',
                    'module.h.h.bias'
            ]:
                h_parameters.append(parameter)
            else:
                parameters.append(parameter)

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            parameters,
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

        # no weight decaty
        self.h_optimizer = torch.optim.SGD(
            h_parameters,
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            nesterov=True,
        )

        # same as normal
        self.h_scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            logits_classifier = self.net(data)
            loss = F.cross_entropy(logits_classifier, target)

            # backward
            self.optimizer.zero_grad()
            self.h_optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.h_optimizer.step()
            self.scheduler.step()
            self.h_scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/godin_trainer.py
kai422/SCALE
import torch
from torch.autograd import Variable
from tqdm import tqdm

from openood.losses.kdad_losses import DirectionOnlyLoss, MseDirectionLoss
from openood.utils import Config


class KdadTrainer:
    def __init__(self, net, train_loader, config: Config):
        self.vgg = net['vgg']
        self.model = net['model']
        self.train_loader = train_loader
        self.config = config
        # choose loss type
        if self.config['direction_loss_only']:
            self.criterion = DirectionOnlyLoss()
        else:
            self.criterion = MseDirectionLoss(self.config['lamda'])
        self.optimizer = torch.optim.Adam(self.model.parameters(),
                                          lr=float(
                                              self.config['learning_rate']))

    def train_epoch(self, epoch_idx):

        self.model.train()
        epoch_loss = 0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}'.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            X = batch['data']
            if X.shape[1] == 1:
                X = X.repeat(1, 3, 1, 1)
            X = Variable(X).cuda()

            # compute respective output
            output_pred = self.model.forward(X)
            output_real = self.vgg(X)

            # compute loss
            total_loss = self.criterion(output_pred, output_real)

            # Add loss to the list
            epoch_loss += total_loss.item()

            # Clear the previous gradients
            self.optimizer.zero_grad()

            # Compute gradients
            total_loss.backward()

            # Adjust weights
            self.optimizer.step()
        net = {}
        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['epoch_loss'] = epoch_loss
        net['vgg'] = self.vgg
        net['model'] = self.model
        return net, metrics

openood/trainers/kdad_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class LogitNormTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

        self.loss_fn = LogitNormLoss(tau=config.trainer.trainer_args.tau)

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            logits_classifier = self.net(data)
            loss = self.loss_fn(logits_classifier, target)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced


class LogitNormLoss(nn.Module):
    def __init__(self, tau=0.04):
        super(LogitNormLoss, self).__init__()
        self.tau = tau

    def forward(self, x, target):
        norms = torch.norm(x, p=2, dim=-1, keepdim=True) + 1e-7
        logit_norm = torch.div(x, norms) / self.tau
        return F.cross_entropy(logit_norm, target)

openood/trainers/logitnorm_trainer.py
kai422/SCALE
import numpy as np


def cosine_annealing(step, total_steps, lr_max, lr_min):
    return lr_min + (lr_max - lr_min) * 0.5 * \
                (1 + np.cos(step / total_steps * np.pi))

openood/trainers/lr_scheduler.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .base_trainer import BaseTrainer


class MCDTrainer(BaseTrainer):
    def __init__(
        self,
        net: nn.Module,
        train_loader: DataLoader,
        train_unlabeled_loader: DataLoader,
        config: Config,
    ) -> None:
        super().__init__(net, train_loader, config)
        self.train_unlabeled_loader = train_unlabeled_loader
        self.lambda_oe = config.trainer.lambda_oe
        self.margin = config.trainer.margin
        self.epoch_ft = config.trainer.start_epoch_ft

    def train_epoch(self, epoch_idx):
        self.net.train()  # enter train mode

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        if self.train_unlabeled_loader:
            unlabeled_dataiter = iter(self.train_unlabeled_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)

            data = batch['data'].cuda()
            if epoch_idx < self.epoch_ft:
                logits1, logits2 = self.net(data, return_double=True)
                loss = F.cross_entropy(logits1, batch['label'].cuda()) \
                    + F.cross_entropy(logits2, batch['label'].cuda())

            elif self.train_unlabeled_loader and epoch_idx >= self.epoch_ft:
                try:
                    unlabeled_batch = next(unlabeled_dataiter)
                except StopIteration:
                    unlabeled_dataiter = iter(self.train_unlabeled_loader)
                    unlabeled_batch = next(unlabeled_dataiter)

                id_bs = data.size(0)

                unlabeled_data = unlabeled_batch['data'].cuda()
                all_data = torch.cat([data, unlabeled_data])
                logits1, logits2 = self.net(all_data, return_double=True)

                logits1_id, logits2_id = logits1[:id_bs], logits2[:id_bs]
                logits1_ood, logits2_ood = logits1[id_bs:], logits2[id_bs:]

                loss = F.cross_entropy(logits1_id, batch['label'].cuda()) \
                    + F.cross_entropy(logits2_id, batch['label'].cuda())

                ent = torch.mean(entropy(logits1_ood) - entropy(logits2_ood))
                loss_oe = F.relu(self.margin - ent)

                loss += self.lambda_oe * loss_oe

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics


def entropy(logits):
    score = torch.softmax(logits, dim=0)
    logscore = torch.log(score)
    entropy = torch.sum(-score * logscore, dim=0)
    return entropy

openood/trainers/mcd_trainer.py
kai422/SCALE
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .base_trainer import BaseTrainer


class MixOETrainer(BaseTrainer):
    def __init__(
        self,
        net: nn.Module,
        train_loader: DataLoader,
        train_unlabeled_loader: DataLoader,
        config: Config,
    ) -> None:
        super().__init__(net, train_loader, config)
        self.train_unlabeled_loader = train_unlabeled_loader
        self.lambda_oe = config.trainer.lambda_oe
        self.alpha = config.trainer.alpha
        self.beta = config.trainer.beta
        self.mix_op = config.trainer.mix_op
        self.num_classes = config.dataset.num_classes
        self.criterion = SoftCE()

    def train_epoch(self, epoch_idx):
        self.net.train()  # enter train mode

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        if self.train_unlabeled_loader:
            unlabeled_dataiter = iter(self.train_unlabeled_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            # manually drop last batch to avoid batch size mismatch
            if train_step == len(train_dataiter):
                continue

            batch = next(train_dataiter)

            try:
                unlabeled_batch = next(unlabeled_dataiter)
            except StopIteration:
                unlabeled_dataiter = iter(self.train_unlabeled_loader)
                unlabeled_batch = next(unlabeled_dataiter)

            if len(unlabeled_batch['data']) < len(batch['data']):
                unlabeled_dataiter = iter(self.train_unlabeled_loader)
                unlabeled_batch = next(unlabeled_dataiter)

            x, y = batch['data'].cuda(), batch['label'].cuda()
            oe_x = unlabeled_batch['data'].cuda()
            bs = x.size(0)
            one_hot_y = torch.zeros(bs, self.num_classes).cuda()
            one_hot_y.scatter_(1, y.view(-1, 1), 1)

            # ID loss
            logits = self.net(x)
            id_loss = F.cross_entropy(logits, y)

            # MixOE loss
            # build mixed samples
            lam = np.random.beta(self.alpha, self.beta)

            if self.mix_op == 'cutmix':
                mixed_x = x.clone().detach()
                bbx1, bby1, bbx2, bby2 = rand_bbox(x.size(), lam)
                # adjust lambda to exactly match pixel ratio
                lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) /
                           (x.size()[-1] * x.size()[-2]))
                # we empirically find that pasting outlier patch into ID data performs better
                # than pasting ID patch into outlier data
                mixed_x[:, :, bbx1:bbx2, bby1:bby2] = oe_x[:, :, bbx1:bbx2,
                                                           bby1:bby2]
            elif self.mix_op == 'mixup':
                mixed_x = lam * x + (1 - lam) * oe_x

            # construct soft labels and compute loss
            oe_y = torch.ones(oe_x.size(0),
                              self.num_classes).cuda() / self.num_classes
            soft_labels = lam * one_hot_y + (1 - lam) * oe_y
            mixed_loss = self.criterion(self.net(mixed_x), soft_labels)

            # Total loss
            loss = id_loss + self.lambda_oe * mixed_loss

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics


class SoftCE(nn.Module):
    def __init__(self, reduction='mean'):
        super(SoftCE, self).__init__()
        self.reduction = reduction

    def forward(self, logits, soft_targets):
        preds = logits.log_softmax(dim=-1)
        assert preds.shape == soft_targets.shape

        loss = torch.sum(-soft_targets * preds, dim=-1)

        if self.reduction == 'mean':
            return torch.mean(loss)
        elif self.reduction == 'sum':
            return torch.sum(loss)
        elif self.reduction == 'none':
            return loss
        else:
            raise ValueError("Reduction type '{:s}' is not supported!".format(
                self.reduction))


def rand_bbox(size, lam):
    W = size[2]
    H = size[3]
    cut_rat = np.sqrt(1. - lam)
    cut_w = int(W * cut_rat)
    cut_h = int(H * cut_rat)

    # uniform
    cx = np.random.randint(W)
    cy = np.random.randint(H)

    bbx1 = np.clip(cx - cut_w // 2, 0, W)
    bby1 = np.clip(cy - cut_h // 2, 0, H)
    bbx2 = np.clip(cx + cut_w // 2, 0, W)
    bby2 = np.clip(cy + cut_h // 2, 0, H)

    return bbx1, bby1, bbx2, bby2

openood/trainers/mixoe_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.losses import soft_cross_entropy
from openood.utils import Config

from .lr_scheduler import cosine_annealing


def prepare_mixup(batch, alpha=1.0, use_cuda=True):
    """Returns mixed inputs, pairs of targets, and lambda."""
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = batch['data'].size()[0]
    if use_cuda:
        index = torch.randperm(batch_size).cuda()
    else:
        index = torch.randperm(batch_size)

    return index, lam


def mixing(data, index, lam):
    return lam * data + (1 - lam) * data[index]


class MixupTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.alpha = self.config.trainer.trainer_args.alpha

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)

            # mixup operation
            index, lam = prepare_mixup(batch, self.alpha)
            data_mix = mixing(batch['data'].cuda(), index, lam)
            soft_label_mix = mixing(batch['soft_label'].cuda(), index, lam)

            # forward
            logits_classifier = self.net(data_mix)
            loss = soft_cross_entropy(logits_classifier, soft_label_mix)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/mixup_trainer.py
kai422/SCALE
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config


def get_mixup(dataset_size):
    return 0.0 if dataset_size < 20_000 else 0.1


def get_group_slices(classes_per_group):
    group_slices = []
    start = 0
    for num_cls in classes_per_group:
        end = start + num_cls + 1
        group_slices.append([start, end])
        start = end
    return torch.LongTensor(group_slices)


def get_schedule(dataset_size):
    if dataset_size < 20_000:
        return [100, 200, 300, 400, 500]
    elif dataset_size < 500_000:
        return [500, 3000, 6000, 9000, 10_000]
    else:
        return [500, 6000, 12_000, 18_000, 20_000]


def get_lr(step, dataset_size, base_lr=0.003):
    """Returns learning-rate for `step` or None at the end."""
    supports = get_schedule(dataset_size)
    # Linear warmup
    if step < supports[0]:
        return base_lr * step / supports[0]
    # End of training
    elif step >= supports[-1]:
        return None
    # Staircase decays by factor of 10
    else:
        for s in supports[1:]:
            if s < step:
                base_lr /= 10
        return base_lr


def mixup_data(x, y, lam):
    """Returns mixed inputs, pairs of targets, and lambda."""
    indices = torch.randperm(x.shape[0]).to(x.device)

    mixed_x = lam * x + (1 - lam) * x[indices]
    y_a, y_b = y, y[indices]
    return mixed_x, y_a, y_b


def mixup_criterion_group(criterion, pred, y_a, y_b, lam, group_slices):
    return lam * calc_group_softmax_loss(criterion, pred, y_a, group_slices) \
           + (1 - lam) * calc_group_softmax_loss(criterion,
                                               pred, y_b, group_slices)


def calc_group_softmax_loss(criterion, logits, labels, group_slices):
    num_groups = group_slices.shape[0]
    loss = 0
    for i in range(num_groups):
        group_logit = logits[:, group_slices[i][0]:group_slices[i][1]]
        group_label = labels[:, i]

        loss += criterion(group_logit, group_label)

    return loss


def calc_group_softmax_acc(logits, labels, group_slices):
    num_groups = group_slices.shape[0]
    loss = 0
    num_samples = logits.shape[0]

    all_group_max_score, all_group_max_class = [], []

    smax = torch.nn.Softmax(dim=-1).cuda()
    cri = torch.nn.CrossEntropyLoss(reduction='none').cuda()

    for i in range(num_groups):
        group_logit = logits[:, group_slices[i][0]:group_slices[i][1]]
        group_label = labels[:, i]
        loss += cri(group_logit, group_label)

        group_softmax = smax(group_logit)
        group_softmax = group_softmax[:, 1:]  # disregard others category
        group_max_score, group_max_class = torch.max(group_softmax, dim=1)
        group_max_class += 1  # shift the class index by 1

        all_group_max_score.append(group_max_score)
        all_group_max_class.append(group_max_class)

    all_group_max_score = torch.stack(all_group_max_score, dim=1)
    all_group_max_class = torch.stack(all_group_max_class, dim=1)

    final_max_score, max_group = torch.max(all_group_max_score, dim=1)

    pred_cls_within_group = all_group_max_class[torch.arange(num_samples),
                                                max_group]

    gt_class, gt_group = torch.max(labels, dim=1)

    selected_groups = (max_group == gt_group)

    pred_acc = torch.zeros(logits.shape[0]).bool().cuda()

    pred_acc[selected_groups] = (
        pred_cls_within_group[selected_groups] == gt_class[selected_groups])

    return loss, pred_acc


def topk(output, target, ks=(1, )):
    """Returns one boolean vector for each k, whether the target is within the
    output's top-k."""
    _, pred = output.topk(max(ks), 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))
    return [correct[:k].max(0)[0] for k in ks]


def run_eval(model, data_loader, step, group_slices, num_group):
    # switch to evaluate mode
    model.eval()

    all_c, all_top1 = [], []

    train_dataiter = iter(data_loader)
    for train_step in tqdm(range(1,
                                 len(train_dataiter) + 1),
                           desc='Test : ',
                           position=0,
                           leave=True,
                           disable=not comm.is_main_process()):
        batch = next(train_dataiter)
        data = batch['data'].cuda()
        group_label = batch['group_label'].cuda()
        class_label = batch['class_label'].cuda()
        labels = []
        for i in range(len(group_label)):
            label = torch.zeros(num_group, dtype=torch.int64)
            label[group_label[i]] = class_label[i] + 1
            labels.append(label.unsqueeze(0))
        labels = torch.cat(labels, dim=0).cuda()

        with torch.no_grad():
            x = data
            y = labels

            # compute output, measure accuracy and record loss.
            logits = model(x)
            if group_slices is not None:
                c, top1 = calc_group_softmax_acc(logits, y, group_slices)
            else:
                c = torch.nn.CrossEntropyLoss(reduction='none')(logits, y)
                top1 = topk(logits, y, ks=(1, ))[0]

            all_c.extend(c.cpu())  # Also ensures a sync point.
            all_top1.extend(top1.cpu())

    model.train()
    # print(f'Validation@{step} loss {np.mean(all_c):.5f}, '
    #       f'top1 {np.mean(all_top1):.2%}')

    # writer.add_scalar('Val/loss', np.mean(all_c), step)
    # writer.add_scalar('Val/top1', np.mean(all_top1), step)
    return all_c, all_top1


class MOSTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net.cuda()
        self.train_loader = train_loader
        self.config = config
        self.lr = config.optimizer.lr

        trainable_params = filter(lambda p: p.requires_grad, net.parameters())
        self.optim = torch.optim.SGD(trainable_params,
                                     lr=self.lr,
                                     momentum=0.9)
        self.optim.zero_grad()
        self.net.train()

        # train_set len
        self.train_set_len = config.dataset.train.batch_size * len(
            train_loader)
        self.mixup = get_mixup(self.train_set_len)
        self.cri = torch.nn.CrossEntropyLoss().cuda()

        self.accum_steps = 0
        self.mixup_l = np.random.beta(self.mixup,
                                      self.mixup) if self.mixup > 0 else 1

        # if specified group_config
        if (config.trainer.group_config.endswith('npy')):
            self.classes_per_group = np.load(config.trainer.group_config)
        elif (config.trainer.group_config.endswith('txt')):
            self.classes_per_group = np.loadtxt(config.trainer.group_config,
                                                dtype=int)
        else:
            self.cal_group_slices(self.train_loader)

        self.num_group = len(self.classes_per_group)
        self.group_slices = get_group_slices(self.classes_per_group)
        self.group_slices = self.group_slices.cuda()

        self.step = 0
        self.batch_split = 1

    def cal_group_slices(self, train_loader):
        # cal group config
        group = {}
        train_dataiter = iter(self.train_loader)
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='cal group_config',
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            group_label = deepcopy(batch['group_label'])
            class_label = deepcopy(batch['class_label'])

            for i in range(len(class_label)):
                gl = group_label[i].item()
                cl = class_label[i].item()

                try:
                    group[str(gl)]
                except:
                    group[str(gl)] = []

                if cl not in group[str(gl)]:
                    group[str(gl)].append(cl)

        self.classes_per_group = []
        for i in range(len(group)):
            self.classes_per_group.append(max(group[str(i)]) + 1)

    def train_epoch(self, epoch_idx):
        total_loss = 0

        train_dataiter = iter(self.train_loader)
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            group_label = batch['group_label'].cuda()
            class_label = batch['class_label'].cuda()

            labels = []
            for i in range(len(group_label)):
                label = torch.zeros(self.num_group, dtype=torch.int64)
                label[group_label[i]] = class_label[i] + 1
                labels.append(label.unsqueeze(0))
            labels = torch.cat(labels, dim=0).cuda()

            # Update learning-rate, including stop training if over.
            lr = get_lr(self.step, self.train_set_len, self.lr)
            if lr is None:
                break
            for param_group in self.optim.param_groups:
                param_group['lr'] = lr

            if self.mixup > 0.0:
                x, y_a, y_b = mixup_data(data, labels, self.mixup_l)

            logits = self.net(data)

            y_a = y_a.cuda()
            y_b = y_b.cuda()
            if self.mixup > 0.0:
                c = mixup_criterion_group(self.cri, logits, y_a, y_b,
                                          self.mixup_l, self.group_slices)
            else:
                c = calc_group_softmax_loss(self.cri, logits, labels,
                                            self.group_slices)

            c_num = float(c.data.cpu().numpy())  # Also ensures a sync point.
            #  # Accumulate grads
            (c / self.batch_split).backward()
            self.accum_steps += 1

            # accstep = f' ({self.accum_steps}/{self.batch_split})' \
            #     if self.batch_split > 1 else ''
            # print(
            #     f'[step {self.step}{accstep}]: loss={c_num:.5f} (lr={lr:.1e})')

            total_loss += c_num

            # Update params
            # if self.accum_steps == self.batch_split:
            self.optim.step()
            self.optim.zero_grad()

            self.step += 1
            self.accum_steps = 0
            # Sample new mixup ratio for next batch
            self.mixup_l = np.random.beta(self.mixup,
                                          self.mixup) if self.mixup > 0 else 1

        # torch.save(self.net.state_dict(),
        #            os.path.join(self.config.output_dir, 'mos_epoch_latest.ckpt'))

        # step, all_top1 = run_eval(self.net, self.train_loader, self.step, self.group_slices,
        #          self.num_group)

        loss_avg = total_loss / len(train_dataiter)

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg
        # metrics['acc'] = np.mean(all_top1) # the acc used in there is the top1 acc

        return self.net, metrics, self.num_group, self.group_slices

openood/trainers/mos_trainer.py
kai422/SCALE
import faiss.contrib.torch_utils
import math
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import MultivariateNormal
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config


class NPOSTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 val_loader: DataLoader, config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        # a bunch of constants or hyperparams
        self.n_cls = config.dataset.num_classes
        self.sample_number = config.trainer.trainer_args.sample_number
        self.sample_from = config.trainer.trainer_args.sample_from
        try:
            self.penultimate_dim = net.backbone.feature_size
        except AttributeError:
            self.penultimate_dim = net.backbone.module.feature_size
        self.start_epoch_KNN = config.trainer.trainer_args.start_epoch_KNN
        self.K = config.trainer.trainer_args.K
        self.select = config.trainer.trainer_args.select
        self.cov_mat = config.trainer.trainer_args.cov_mat
        self.pick_nums = config.trainer.trainer_args.pick_nums
        self.w_disp = config.trainer.trainer_args.w_disp
        self.w_comp = config.trainer.trainer_args.w_comp
        self.loss_weight = config.trainer.trainer_args.loss_weight
        self.temp = config.trainer.trainer_args.temp
        self.ID_points_num = config.trainer.trainer_args.ID_points_num

        res = faiss.StandardGpuResources()
        self.KNN_index = faiss.GpuIndexFlatL2(res, self.penultimate_dim)

        self.number_dict = {}
        for i in range(self.n_cls):
            self.number_dict[i] = 0

        if self.config.num_gpus > 1:
            params = [{
                'params': net.module.backbone.parameters()
            }, {
                'params': net.module.head.parameters()
            }, {
                'params':
                net.module.mlp.parameters(),
                'lr':
                config.optimizer.lr * config.optimizer.mlp_decay_rate
            }]
        else:
            params = [{
                'params': net.backbone.parameters()
            }, {
                'params': net.head.parameters()
            }, {
                'params':
                net.mlp.parameters(),
                'lr':
                config.optimizer.lr * config.optimizer.mlp_decay_rate
            }]

        self.optimizer = torch.optim.SGD(
            params,
            lr=config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        if config.dataset.train.batch_size \
                * config.num_gpus * config.num_machines > 256:
            config.optimizer.warm = True

        if config.optimizer.warm:
            self.warmup_from = 0.001
            self.warm_epochs = 10
            if config.optimizer.cosine:
                eta_min = config.optimizer.lr * \
                    (config.optimizer.lr_decay_rate**3)
                self.warmup_to = eta_min + (config.optimizer.lr - eta_min) * (
                    1 + math.cos(math.pi * self.warm_epochs /
                                 config.optimizer.num_epochs)) / 2
            else:
                self.warmup_to = config.optimizer.lr

        self.criterion_comp = CompLoss(self.n_cls,
                                       temperature=self.temp).cuda()
        # V2: EMA style prototypes
        self.criterion_disp = DispLoss(self.n_cls,
                                       config.network.feat_dim,
                                       config.trainer.trainer_args.proto_m,
                                       self.net,
                                       val_loader,
                                       temperature=self.temp).cuda()

    def train_epoch(self, epoch_idx):
        adjust_learning_rate(self.config, self.optimizer, epoch_idx - 1)

        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        data_dict = torch.zeros(self.n_cls, self.sample_number,
                                self.penultimate_dim).cuda()

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            warmup_learning_rate(self.config, self.warm_epochs,
                                 self.warmup_from,
                                 self.warmup_to, epoch_idx - 1, train_step,
                                 len(train_dataiter), self.optimizer)

            batch = next(train_dataiter)
            data = batch['data']
            target = batch['label']

            data = torch.cat([data[0], data[1]], dim=0).cuda()
            target = target.repeat(2).cuda()

            # forward
            penultimate = self.net.backbone(data)
            features = self.net.head(penultimate)

            sum_temp = 0
            for index in range(self.n_cls):
                sum_temp += self.number_dict[index]
            lr_reg_loss = torch.zeros(1).cuda()[0]

            if sum_temp == self.n_cls * self.sample_number \
                    and epoch_idx < self.start_epoch_KNN:
                # maintaining an ID data queue for each class.
                target_numpy = target.cpu().data.numpy()
                for index in range(len(target)):
                    dict_key = target_numpy[index]
                    data_dict[dict_key] = torch.cat(
                        (data_dict[dict_key][1:],
                         penultimate[index].detach().view(1, -1)), 0)
            elif sum_temp == self.n_cls * self.sample_number \
                    and epoch_idx >= self.start_epoch_KNN:
                target_numpy = target.cpu().data.numpy()
                for index in range(len(target)):
                    dict_key = target_numpy[index]
                    data_dict[dict_key] = torch.cat(
                        (data_dict[dict_key][1:],
                         penultimate[index].detach().view(1, -1)), 0)
                # Standard Gaussian distribution
                new_dis = MultivariateNormal(
                    torch.zeros(self.penultimate_dim).cuda(),
                    torch.eye(self.penultimate_dim).cuda())
                negative_samples = new_dis.rsample((self.sample_from, ))
                for index in range(self.n_cls):
                    ID = data_dict[index]
                    sample_point = generate_outliers(
                        ID,
                        input_index=self.KNN_index,
                        negative_samples=negative_samples,
                        ID_points_num=self.ID_points_num,
                        K=self.K,
                        select=self.select,
                        cov_mat=self.cov_mat,
                        sampling_ratio=1.0,
                        pic_nums=self.pick_nums,
                        depth=self.penultimate_dim)
                    if index == 0:
                        ood_samples = sample_point
                    else:
                        ood_samples = torch.cat((ood_samples, sample_point), 0)

                if len(ood_samples) != 0:
                    energy_score_for_fg = self.net.mlp(penultimate)
                    energy_score_for_bg = self.net.mlp(ood_samples)
                    input_for_lr = torch.cat(
                        (energy_score_for_fg, energy_score_for_bg),
                        0).squeeze()
                    labels_for_lr = torch.cat(
                        (torch.ones(len(energy_score_for_fg)).cuda(),
                         torch.zeros(len(energy_score_for_bg)).cuda()), -1)
                    criterion_BCE = torch.nn.BCEWithLogitsLoss()
                    lr_reg_loss = criterion_BCE(input_for_lr.view(-1),
                                                labels_for_lr)
            else:
                target_numpy = target.cpu().data.numpy()
                for index in range(len(target)):
                    dict_key = target_numpy[index]
                    if self.number_dict[dict_key] < self.sample_number:
                        data_dict[dict_key][self.number_dict[
                            dict_key]] = penultimate[index].detach()
                        self.number_dict[dict_key] += 1
            normed_features = F.normalize(features, dim=1)

            disp_loss = self.criterion_disp(normed_features, target)
            comp_loss = self.criterion_comp(normed_features,
                                            self.criterion_disp.prototypes,
                                            target)

            loss = self.w_disp * disp_loss + self.w_comp * comp_loss
            loss = self.loss_weight * lr_reg_loss + loss

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced


def adjust_learning_rate(config, optimizer, epoch):
    lr = config.optimizer.lr
    if config.optimizer.cosine:
        eta_min = lr * (config.optimizer.lr_decay_rate**3)
        lr = eta_min + (lr - eta_min) * (
            1 + math.cos(math.pi * epoch / config.optimizer.num_epochs)) / 2
    else:
        steps = np.sum(epoch > np.asarray(config.optimizer.lr_decay_epochs))
        if steps > 0:
            lr = lr * (config.optimizer.lr_decay_rate**steps)

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def warmup_learning_rate(config, warm_epochs, warmup_from, warmup_to, epoch,
                         batch_id, total_batches, optimizer):
    if config.optimizer.warm and epoch <= warm_epochs:
        p = (batch_id + (epoch - 1) * total_batches) / \
            (warm_epochs * total_batches)
        lr = warmup_from + p * (warmup_to - warmup_from)

        for param_group in optimizer.param_groups:
            param_group['lr'] = lr


class CompLoss(nn.Module):
    def __init__(self, n_cls, temperature=0.07, base_temperature=0.07):
        super(CompLoss, self).__init__()
        self.n_cls = n_cls
        self.temperature = temperature
        self.base_temperature = base_temperature

    def forward(self, features, prototypes, labels):
        device = torch.device('cuda')

        proxy_labels = torch.arange(0, self.n_cls).to(device)
        batch_size = features.shape[0]
        labels = labels.contiguous().view(-1, 1)
        if labels.shape[0] != batch_size:
            raise ValueError('Num of labels does not match num of features')
        mask = torch.eq(labels, proxy_labels.T).float().to(device)

        # compute logits
        anchor_feature = features
        contrast_feature = prototypes / prototypes.norm(dim=-1, keepdim=True)
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T), self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # compute log_prob
        exp_logits = torch.exp(logits)
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1)
        loss = -(self.temperature /
                 self.base_temperature) * mean_log_prob_pos.mean()
        return loss


class DispLoss(nn.Module):
    def __init__(self,
                 n_cls,
                 feat_dim,
                 proto_m,
                 model,
                 loader,
                 temperature=0.1,
                 base_temperature=0.1):
        super(DispLoss, self).__init__()
        self.n_cls = n_cls
        self.feat_dim = feat_dim
        self.proto_m = proto_m
        self.temperature = temperature
        self.base_temperature = base_temperature
        self.register_buffer('prototypes',
                             torch.zeros(self.n_cls, self.feat_dim))
        self.model = model
        self.loader = loader
        self.init_class_prototypes()

    def forward(self, features, labels):
        prototypes = self.prototypes
        num_cls = self.n_cls
        for j in range(len(features)):
            prototypes[labels[j].item()] = F.normalize(
                prototypes[labels[j].item()] * self.proto_m + features[j] *
                (1 - self.proto_m),
                dim=0)
        self.prototypes = prototypes.detach()
        labels = torch.arange(0, num_cls).cuda()
        labels = labels.contiguous().view(-1, 1)

        mask = (1 - torch.eq(labels, labels.T).float()).cuda()

        logits = torch.div(torch.matmul(prototypes, prototypes.T),
                           self.temperature)

        logits_mask = torch.scatter(torch.ones_like(mask), 1,
                                    torch.arange(num_cls).view(-1, 1).cuda(),
                                    0)
        mask = mask * logits_mask
        mean_prob_neg = torch.log(
            (mask * torch.exp(logits)).sum(1) / mask.sum(1))
        mean_prob_neg = mean_prob_neg[~torch.isnan(mean_prob_neg)]
        loss = self.temperature / self.base_temperature * mean_prob_neg.mean()
        return loss

    def init_class_prototypes(self):
        """Initialize class prototypes."""
        self.model.eval()
        start = time.time()
        prototype_counts = [0] * self.n_cls
        with torch.no_grad():
            prototypes = torch.zeros(self.n_cls, self.feat_dim).cuda()
            for i, batch in enumerate(self.loader):
                input = batch['data']
                target = batch['label']
                input, target = input.cuda(), target.cuda()
                features = self.model(input)
                for j, feature in enumerate(features):
                    prototypes[target[j].item()] += feature
                    prototype_counts[target[j].item()] += 1
            for cls in range(self.n_cls):
                prototypes[cls] /= prototype_counts[cls]
            # measure elapsed time
            duration = time.time() - start
            print(f'Time to initialize prototypes: {duration:.3f}')
            prototypes = F.normalize(prototypes, dim=1)
            self.prototypes = prototypes


def generate_outliers(ID,
                      input_index,
                      negative_samples,
                      ID_points_num=2,
                      K=20,
                      select=1,
                      cov_mat=0.1,
                      sampling_ratio=1.0,
                      pic_nums=30,
                      depth=342):
    length = negative_samples.shape[0]
    data_norm = torch.norm(ID, p=2, dim=1, keepdim=True)
    normed_data = ID / data_norm
    rand_ind = np.random.choice(normed_data.shape[0],
                                int(normed_data.shape[0] * sampling_ratio),
                                replace=False)
    index = input_index
    index.add(normed_data[rand_ind])
    minD_idx, k_th = KNN_dis_search_decrease(ID, index, K, select)
    minD_idx = minD_idx[np.random.choice(select, int(pic_nums), replace=False)]
    data_point_list = torch.cat(
        [ID[i:i + 1].repeat(length, 1) for i in minD_idx])
    negative_sample_cov = cov_mat * negative_samples.cuda().repeat(pic_nums, 1)
    negative_sample_list = negative_sample_cov + data_point_list
    point = KNN_dis_search_distance(negative_sample_list, index, K,
                                    ID_points_num, length, depth)

    index.reset()
    return point


def KNN_dis_search_distance(target,
                            index,
                            K=50,
                            num_points=10,
                            length=2000,
                            depth=342):
    '''
    data_point: Queue for searching k-th points
    target: the target of the search
    K
    '''
    # Normalize the features
    target_norm = torch.norm(target, p=2, dim=1, keepdim=True)
    normed_target = target / target_norm

    distance, output_index = index.search(normed_target, K)
    k_th_distance = distance[:, -1]
    k_th = k_th_distance.view(length, -1)
    # target_new = target.view(length, -1, depth)
    k_th_distance, minD_idx = torch.topk(k_th, num_points, dim=0)
    minD_idx = minD_idx.squeeze()
    point_list = []
    for i in range(minD_idx.shape[1]):
        point_list.append(i * length + minD_idx[:, i])
    return target[torch.cat(point_list)]


def KNN_dis_search_decrease(
    target,
    index,
    K=50,
    select=1,
):
    '''
    data_point: Queue for searching k-th points
    target: the target of the search
    K
    '''
    # Normalize the features
    target_norm = torch.norm(target, p=2, dim=1, keepdim=True)
    normed_target = target / target_norm

    distance, output_index = index.search(normed_target, K)
    k_th_distance = distance[:, -1]
    k_th_distance, minD_idx = torch.topk(k_th_distance, select)
    return minD_idx, k_th_distance

openood/trainers/npos_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .base_trainer import BaseTrainer


class OETrainer(BaseTrainer):
    def __init__(
        self,
        net: nn.Module,
        train_loader: DataLoader,
        train_unlabeled_loader: DataLoader,
        config: Config,
    ) -> None:
        super().__init__(net, train_loader, config)
        self.train_unlabeled_loader = train_unlabeled_loader
        self.lambda_oe = config.trainer.lambda_oe

    def train_epoch(self, epoch_idx):
        self.net.train()  # enter train mode

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        if self.train_unlabeled_loader:
            unlabeled_dataiter = iter(self.train_unlabeled_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)

            try:
                unlabeled_batch = next(unlabeled_dataiter)
            except StopIteration:
                unlabeled_dataiter = iter(self.train_unlabeled_loader)
                unlabeled_batch = next(unlabeled_dataiter)

            data = torch.cat((batch['data'], unlabeled_batch['data'])).cuda()
            batch_size = batch['data'].size(0)

            # forward
            logits_classifier = self.net(data)
            loss = F.cross_entropy(logits_classifier[:batch_size],
                                   batch['label'].cuda())

            loss_oe = -(
                logits_classifier[batch_size:].mean(1) -
                torch.logsumexp(logits_classifier[batch_size:], dim=1)).mean()
            loss += self.lambda_oe * loss_oe

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

openood/trainers/oe_trainer.py
kai422/SCALE
import random

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

import openood.utils.comm as comm


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


class OpenGanTrainer:
    def __init__(self, net, feat_loader, config) -> None:

        manualSeed = 999
        print('Random Seed: ', manualSeed)
        random.seed(manualSeed)
        torch.manual_seed(manualSeed)

        self.config = config
        self.netG = net['netG']
        self.netD = net['netD']
        self.netG.apply(weights_init)
        self.netD.apply(weights_init)
        self.feat_loader = feat_loader

        self.nz = self.config.network.nz

        self.real_label = 1
        self.fake_label = 0

        optimizer_config = self.config.optimizer
        self.optimizerD = optim.Adam(self.netD.parameters(),
                                     lr=optimizer_config.lr / 1.5,
                                     betas=(optimizer_config.beta1, 0.999))
        self.optimizerG = optim.Adam(self.netG.parameters(),
                                     lr=optimizer_config.lr,
                                     betas=(optimizer_config.beta1, 0.999))

        self.criterion = nn.BCELoss()

        self.G_losses = []
        self.D_losses = []

    def train_epoch(self, epoch_idx):

        feat_dataiter = iter(self.feat_loader)

        for train_step in tqdm(range(1,
                                     len(feat_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            data = next(feat_dataiter)['data']
            ############################
            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
            ###########################
            # Train with all-real batch
            self.netD.zero_grad()
            # Format batch
            loaded_data = data.cuda()
            b_size = loaded_data.size(0)
            label = torch.full((b_size, ), self.real_label).cuda()
            label = label.to(torch.float32)

            # Forward pass real batch through D
            output = self.netD(loaded_data).view(-1)
            # import pdb
            # pdb.set_trace()
            # Calculate loss on all-real batch
            errD_real = self.criterion(output, label)
            # Calculate gradients for D in backward pass
            errD_real.backward()
            D_x = output.mean().item()

            # Train with all-fake batch
            # Generate batch of latent vectors
            noise = torch.randn(b_size, self.nz, 1, 1).cuda()
            # Generate fake image batch with G
            fake = self.netG(noise)
            label.fill_(self.fake_label)
            # Classify all fake batch with D
            output = self.netD(fake.detach()).view(-1)
            # Calculate D's loss on the all-fake batch
            errD_fake = self.criterion(output, label)
            # Calculate the gradients for this batch
            errD_fake.backward()
            D_G_z1 = output.mean().item()
            # Add the gradients from the all-real and all-fake batches
            errD = errD_real + errD_fake
            # Update D
            self.optimizerD.step()

            ############################
            # (2) Update G network: maximize log(D(G(z)))
            ###########################
            self.netG.zero_grad()
            label.fill_(
                self.real_label)  # fake labels are real for generator cost
            # Since we just updated D,
            # perform another forward pass of all-fake batch through D
            output = self.netD(fake).view(-1)
            # Calculate G's loss based on this output
            errG = self.criterion(output, label)
            # Calculate gradients for G
            errG.backward()
            D_G_z2 = output.mean().item()
            # Update G
            self.optimizerG.step()

            # Save Losses for plotting later, if needed
            self.G_losses.append(errG.item())
            self.D_losses.append(errD.item())

        return {
            'netG': self.netG,
            'netD': self.netD
        }, {
            'G_losses': self.G_losses,
            'D_losses': self.D_losses,
            'epoch_idx': epoch_idx
        }

openood/trainers/opengan_trainer.py
kai422/SCALE
import torch

from torchvision.datasets import ImageFolder
import numpy as np
import random
import os
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
import argparse
from torch.nn import functional as F
from tqdm import tqdm
from openood.utils import Config
from openood.losses.rd4ad_loss import loss_function

class Rd4adTrainer:
    def __init__(self, net, train_loader, config: Config):
        self.config = config
        self.train_loader = train_loader
        self.encoder = net['encoder']
        self.bn = net['bn'] 
        self.decoder = net['decoder']
        if config.optimizer.name == 'adam':
            self.optimizer=torch.optim.Adam(list(self.decoder.parameters())+list(self.bn.parameters()), lr=config.optimizer.lr, betas=config.optimizer.betas)

    def train_epoch(self, epoch_idx):
        self.encoder.eval()
        self.bn.train()
        self.decoder.train()
        train_dataiter = iter(self.train_loader)
        epoch_loss = 0
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d} '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            img = batch['data'].cuda()
            feature_list = self.encoder.forward(img,return_feature_list=True)[1]
            inputs = feature_list[1:4]
            outputs = self.decoder(self.bn(inputs))
            loss = loss_function(inputs, outputs)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            epoch_loss += loss.item()
        metrics = {}
        net = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = epoch_loss
        net['encoder'] = self.encoder
        net['bn'] = self.bn
        net['decoder'] = self.decoder
        return net, metrics

openood/trainers/rd4ad_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config
from .lr_scheduler import cosine_annealing


# https://github.com/FrancescoPinto/RegMixup/blob/main/models/regmixup.py
def mixup_data(x, y, alpha=1.0):
    """Returns mixed inputs, pairs of targets, and lambda."""

    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = x.size()[0]
    index = torch.randperm(batch_size).cuda()
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam


def regmixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)


class RegMixupTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.alpha = self.config.trainer.trainer_args.alpha

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            x, y = batch['data'].cuda(), batch['label'].cuda()

            # mixup operation
            mixup_x, part_y_a, part_y_b, lam = mixup_data(x, y, self.alpha)
            targets_a = torch.cat([y, part_y_a])
            targets_b = torch.cat([y, part_y_b])
            x = torch.cat([x, mixup_x], dim=0)

            # forward
            logits = self.net(x)
            loss = regmixup_criterion(F.cross_entropy, logits, targets_a,
                                      targets_b, lam)

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/regmixup_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from .lr_scheduler import cosine_annealing


class RotPredTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            batch_size = len(data)
            x_90 = torch.rot90(data, 1, [2, 3])
            x_180 = torch.rot90(data, 2, [2, 3])
            x_270 = torch.rot90(data, 3, [2, 3])

            x_rot = torch.cat([data, x_90, x_180, x_270])
            y_rot = torch.cat([
                torch.zeros(batch_size),
                torch.ones(batch_size),
                2 * torch.ones(batch_size),
                3 * torch.ones(batch_size),
            ]).long().cuda()

            # forward
            logits, logits_rot = self.net(x_rot, return_rot_logits=True)
            loss_cls = F.cross_entropy(logits[:batch_size], target)
            loss_rot = F.cross_entropy(logits_rot, y_rot)
            loss = loss_cls + loss_rot

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        # comm.synchronize()

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = self.save_metrics(loss_avg)

        return self.net, metrics

    def save_metrics(self, loss_avg):
        all_loss = comm.gather(loss_avg)
        total_losses_reduced = np.mean([x for x in all_loss])

        return total_losses_reduced

openood/trainers/rotpred_trainer.py
kai422/SCALE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.utils import Config

from .lr_scheduler import cosine_annealing



class RTSTrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )



    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # forward
            logits_classifier, variance = self.net(data, return_var=True)
            epsilon = torch.randn_like(variance)
            temperature = torch.sum(variance * epsilon * epsilon, dim=1, keepdim=True) / (self.config.network.dof - 2)
            loss_kl = ((variance - torch.log(variance + 1e-8) - 1) * 0.5).mean()
            loss_head = F.cross_entropy(logits_classifier / temperature, target)
            loss = loss_head + self.config.network.kl_scale * loss_kl
            
            # backward
            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(parameters=self.net.parameters(), max_norm=2.5, norm_type=2)
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/rts_trainer.py
kai422/SCALE
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.mixture import GaussianMixture
from torch.utils.data import DataLoader
from tqdm import tqdm

from openood.losses import soft_cross_entropy
from openood.postprocessors.gmm_postprocessor import compute_single_GMM_score
from openood.postprocessors.mds_ensemble_postprocessor import (
    process_feature_type, reduce_feature_dim, tensor2list)
from openood.utils import Config

from .lr_scheduler import cosine_annealing
from .mixup_trainer import mixing, prepare_mixup


class SAETrainer:
    def __init__(self, net: nn.Module, train_loader: DataLoader,
                 config: Config) -> None:

        self.net = net
        self.train_loader = train_loader
        self.config = config
        self.trainer_args = self.config.trainer.trainer_args

        self.optimizer = torch.optim.SGD(
            net.parameters(),
            config.optimizer.lr,
            momentum=config.optimizer.momentum,
            weight_decay=config.optimizer.weight_decay,
            nesterov=True,
        )

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step,
                config.optimizer.num_epochs * len(train_loader),
                1,
                1e-6 / config.optimizer.lr,
            ),
        )

    @torch.no_grad()
    def setup(self):
        feature_all = None
        label_all = []
        # collect features
        for batch in tqdm(self.train_loader,
                          desc='Compute GMM Stats [Collecting]'):
            data = batch['data_aux'].cuda()
            label = batch['label']
            _, feature_list = self.net(data, return_feature_list=True)
            label_all.extend(tensor2list(label))
            feature_processed = process_feature_type(
                feature_list[0], self.trainer_args.feature_type)
            if isinstance(feature_all, type(None)):
                feature_all = tensor2list(feature_processed)
            else:
                feature_all.extend(tensor2list(feature_processed))
        label_all = np.array(label_all)

        # reduce feature dim and perform gmm estimation
        feature_all = np.array(feature_all)
        transform_matrix = reduce_feature_dim(feature_all, label_all,
                                              self.trainer_args.reduce_dim)
        feature_all = np.dot(feature_all, transform_matrix)
        # GMM estimation
        gm = GaussianMixture(n_components=self.trainer_args.num_clusters,
                             random_state=0,
                             covariance_type='tied').fit(feature_all)
        feature_mean = gm.means_
        feature_prec = gm.precisions_
        component_weight = gm.weights_

        self.feature_mean = torch.Tensor(feature_mean).cuda()
        self.feature_prec = torch.Tensor(feature_prec).cuda()
        self.component_weight = torch.Tensor(component_weight).cuda()
        self.transform_matrix = torch.Tensor(transform_matrix).cuda()

    def train_epoch(self, epoch_idx):
        self.net.train()

        loss_avg = 0.0
        train_dataiter = iter(self.train_loader)

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True):
            batch = next(train_dataiter)
            data = batch['data'].cuda()
            target = batch['label'].cuda()

            # mixup operation
            index, lam = prepare_mixup(batch, self.trainer_args.alpha)
            data_mix = mixing(batch['data'].cuda(), index, lam)
            soft_label_mix = mixing(batch['soft_label'].cuda(), index, lam)

            # classfication loss
            logits_cls = self.net(data)
            loss_clsstd = F.cross_entropy(logits_cls, target)  # standard cls
            logits_mix = self.net(data_mix)
            loss_clsmix = soft_cross_entropy(logits_mix, soft_label_mix)

            # source awareness enhancement
            prob_id = compute_single_GMM_score(self.net, data,
                                               self.feature_mean,
                                               self.feature_prec,
                                               self.component_weight,
                                               self.transform_matrix, 0,
                                               self.trainer_args.feature_type)
            prob_ood = compute_single_GMM_score(self.net, data_mix,
                                                self.feature_mean,
                                                self.feature_prec,
                                                self.component_weight,
                                                self.transform_matrix, 0,
                                                self.trainer_args.feature_type)
            loss_sae_id = 1 - torch.mean(prob_id)
            loss_sae_ood = torch.mean(prob_ood)

            # loss
            loss = self.trainer_args.loss_weight[0] * loss_clsstd \
                + self.trainer_args.loss_weight[1] * loss_clsmix \
                + self.trainer_args.loss_weight[2] * loss_sae_id \
                + self.trainer_args.loss_weight[3] * loss_sae_ood

            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            # exponential moving average, show smooth values
            with torch.no_grad():
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['loss'] = loss_avg

        return self.net, metrics

openood/trainers/sae_trainer.py
kai422/SCALE
import time

import faiss
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config

from ..losses import rew_ce, rew_sce
from .base_trainer import BaseTrainer


class UDGTrainer(BaseTrainer):
    def __init__(
        self,
        net: nn.Module,
        train_loader: DataLoader,
        train_unlabeled_loader: DataLoader,
        config: Config,
    ) -> None:
        super().__init__(net, train_loader, config)

        self.train_unlabeled_loader = train_unlabeled_loader

        self.num_clusters = config.trainer.num_clusters
        self.purity_ind_thresh = config.trainer.purity_ind_thresh
        self.purity_ood_thresh = config.trainer.purity_ood_thresh
        self.oe_enhance_ratio = config.trainer.oe_enhance_ratio
        self.lambda_oe = config.trainer.lambda_oe
        self.lambda_aux = config.trainer.lambda_aux

        # Init clustering algorithm
        self.k_means = KMeans(k=config.trainer.num_clusters,
                              pca_dim=config.trainer.pca_dim)

    def train_epoch(self, epoch_idx):
        self._run_clustering(epoch_idx)
        metrics = self._compute_loss(epoch_idx)

        return self.net, metrics

    def _compute_loss(self, epoch_idx):
        self.net.train()  # enter train mode

        loss_avg, loss_cls_avg, loss_oe_avg, loss_aux_avg = 0.0, 0.0, 0.0, 0.0
        train_dataiter = iter(self.train_loader)
        unlabeled_dataiter = iter(self.train_unlabeled_loader)
        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}: '.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            try:
                unlabeled_batch = next(unlabeled_dataiter)
            except StopIteration:
                unlabeled_dataiter = iter(self.train_unlabeled_loader)
                unlabeled_batch = next(unlabeled_dataiter)
            data = batch['data'].cuda()
            unlabeled_data = unlabeled_batch['data'].cuda()

            # concat labeled and unlabeled data
            logits_cls, logits_aux = self.net(data, return_aux=True)
            logits_oe_cls, logits_oe_aux = self.net(unlabeled_data,
                                                    return_aux=True)

            # classification loss
            concat_logits_cls = torch.cat([logits_cls, logits_oe_cls])
            concat_label = torch.cat([
                batch['label'],
                unlabeled_batch['pseudo_label'].type_as(batch['label']),
            ])
            loss_cls = F.cross_entropy(
                concat_logits_cls[concat_label != -1],
                concat_label[concat_label != -1].cuda(),
            )
            # oe loss
            concat_softlabel = torch.cat(
                [batch['soft_label'], unlabeled_batch['pseudo_softlabel']])
            concat_conf = torch.cat(
                [batch['ood_conf'], unlabeled_batch['ood_conf']])
            loss_oe = rew_sce(
                concat_logits_cls[concat_label == -1],
                concat_softlabel[concat_label == -1].cuda(),
                concat_conf[concat_label == -1].cuda(),
            )
            # aux loss
            concat_logits_aux = torch.cat([logits_aux, logits_oe_aux])
            concat_cluster_id = torch.cat(
                [batch['cluster_id'], unlabeled_batch['cluster_id']])
            concat_cluster_reweight = torch.cat([
                batch['cluster_reweight'], unlabeled_batch['cluster_reweight']
            ])
            loss_aux = rew_ce(
                concat_logits_aux,
                concat_cluster_id.cuda(),
                concat_cluster_reweight.cuda(),
            )

            # loss addition
            loss = loss_cls + self.lambda_oe * loss_oe \
                + self.lambda_aux * loss_aux
            # backward
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.scheduler.step()

            with torch.no_grad():
                # exponential moving average, show smooth values
                loss_cls_avg = loss_cls_avg * 0.8 + float(loss_cls) * 0.2
                loss_oe_avg = loss_oe_avg * 0.8 + float(
                    self.lambda_oe * loss_oe) * 0.2
                loss_aux_avg = (loss_aux_avg * 0.8 +
                                float(self.lambda_aux * loss_aux) * 0.2)
                loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['epoch_idx'] = epoch_idx
        metrics['train_cls_loss'] = loss_cls_avg
        metrics['train_oe_loss'] = loss_oe_avg
        metrics['train_aux_loss'] = loss_aux_avg
        metrics['loss'] = loss_avg

        return metrics

    def _run_clustering(self, epoch_idx):
        self.net.eval()

        start_time = time.time()
        # get data from train loader
        print('Clustering: starting gather training features...', flush=True)
        # gather train image feature
        train_idx_list, unlabeled_idx_list, feature_list, train_label_list = (
            [],
            [],
            [],
            [],
        )
        train_dataiter = iter(self.train_loader)
        for step in tqdm(range(1,
                               len(train_dataiter) + 1),
                         desc='Epoch {:03d} ID Clustering: '.format(epoch_idx),
                         position=0,
                         leave=True,
                         disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            index = batch['index']
            label = batch['label']
            # we use no augmented image for clustering
            data = batch['data_aux'].cuda()
            _, feature = self.net(data, return_feature=True)
            feature = feature.detach()
            # evaluation
            for idx in range(len(data)):
                train_idx_list.append(index[idx].tolist())
                train_label_list.append(label[idx].tolist())
                feature_list.append(feature[idx].cpu().tolist())
        num_train_data = len(feature_list)
        train_idx_list = np.array(train_idx_list, dtype=int)
        train_label_list = np.array(train_label_list, dtype=int)
        train_label_list = sort_array(train_label_list, train_idx_list)
        # in-distribution samples always have pseudo labels == actual labels
        self.train_loader.dataset.pseudo_label = train_label_list

        torch.cuda.empty_cache()

        # gather unlabeled image feature in order
        unlabeled_conf_list, unlabeled_pseudo_list = [], []
        unlabeled_dataiter = iter(self.train_unlabeled_loader)
        for step in tqdm(range(1,
                               len(unlabeled_dataiter) + 1),
                         desc='Epoch {:03d} OE Clustering: '.format(epoch_idx),
                         position=0,
                         leave=True,
                         disable=not comm.is_main_process()):
            batch = next(unlabeled_dataiter)
            index = batch['index']
            # we use no augmented image for clustering
            data = batch['data_aux'].cuda()
            logit, feature = self.net(data, return_feature=True)
            feature = feature.detach()
            logit = logit.detach()
            score = torch.softmax(logit, dim=1)
            conf, pseudo = torch.max(score, dim=1)
            # evaluation
            for idx in range(len(data)):
                unlabeled_idx_list.append(index[idx].tolist())
                feature_list.append(feature[idx].cpu().tolist())
                unlabeled_conf_list.append(conf[idx].cpu().tolist())
                unlabeled_pseudo_list.append(pseudo[idx].cpu().tolist())
        feature_list = np.array(feature_list)
        unlabeled_idx_list = np.array(unlabeled_idx_list, dtype=int)
        unlabeled_conf_list = np.array(unlabeled_conf_list)
        unlabeled_pseudo_list = np.array(unlabeled_pseudo_list)
        unlabeled_conf_list = sort_array(unlabeled_conf_list,
                                         unlabeled_idx_list)
        unlabeled_pseudo_list = sort_array(unlabeled_pseudo_list,
                                           unlabeled_idx_list)
        torch.cuda.empty_cache()

        print('\nAssigning Cluster Labels...', flush=True)
        cluster_id = self.k_means.cluster(feature_list)
        train_cluster_id = cluster_id[:num_train_data]
        unlabeled_cluster_id = cluster_id[num_train_data:]
        # assign cluster id to samples. Sorted by shuffle-recording index.
        train_cluster_id = sort_array(train_cluster_id, train_idx_list)
        unlabeled_cluster_id = sort_array(unlabeled_cluster_id,
                                          unlabeled_idx_list)
        self.train_loader.dataset.cluster_id = train_cluster_id
        self.train_unlabeled_loader.dataset.cluster_id = unlabeled_cluster_id
        cluster_id = np.concatenate([train_cluster_id, unlabeled_cluster_id])
        # reweighting based on samples in clusters
        cluster_stat = np.zeros(self.num_clusters)
        cluster_id_list, cluster_id_counts = np.unique(cluster_id,
                                                       return_counts=True)
        for cluster_idx, counts in zip(cluster_id_list, cluster_id_counts):
            cluster_stat[cluster_idx] = counts
        inv_class_freq = 1 / (cluster_stat + 1e-10)
        sample_weight = np.power(inv_class_freq, 0.5)
        sample_weight *= 1 / sample_weight.mean()
        sample_weight_list = np.array([sample_weight[i] for i in cluster_id])
        self.train_loader.dataset.cluster_reweight \
            = sample_weight_list[:num_train_data]
        self.train_unlabeled_loader.dataset.cluster_reweight \
            = sample_weight_list[num_train_data:]

        print('In-Distribution Filtering (with OOD Enhancement)...',
              flush=True)
        old_train_pseudo_label \
            = self.train_loader.dataset.pseudo_label
        old_unlabeled_pseudo_label \
            = self.train_unlabeled_loader.dataset.pseudo_label
        old_pseudo_label = np.append(old_train_pseudo_label,
                                     old_unlabeled_pseudo_label).astype(int)
        new_pseudo_label = (-1 * np.ones_like(old_pseudo_label)).astype(int)
        # process ood confidence for oe loss enhancement (ole)
        new_ood_conf = np.ones_like(old_pseudo_label).astype(float)

        total_num_to_filter = 0
        purity_ind_thresh = self.purity_ind_thresh
        purity_ood_thresh = self.purity_ood_thresh
        # pick out clusters with purity over threshold
        for cluster_idx in range(self.num_clusters):
            label_in_cluster, label_counts = np.unique(
                old_pseudo_label[cluster_id == cluster_idx],
                return_counts=True)
            cluster_size = len(old_pseudo_label[cluster_id == cluster_idx])
            purity = label_counts / cluster_size  # purity list for each label
            # idf
            if np.any(purity > purity_ind_thresh):
                majority_label = label_in_cluster[purity > purity_ind_thresh][
                    0]  # first element in the list
                new_pseudo_label[cluster_id == cluster_idx] = majority_label
                # this might also change some ID but nvm
                if majority_label > 0:  # ID cluster
                    num_to_filter = len(label_in_cluster == -1)
                    total_num_to_filter += num_to_filter
            # ole
            elif np.any(purity > purity_ood_thresh):
                majority_label = label_in_cluster[
                    purity > purity_ood_thresh][0]
                if majority_label == -1:
                    new_ood_conf[cluster_id ==
                                 cluster_idx] = self.oe_enhance_ratio
        print(f'{total_num_to_filter} sample(s) filtered!', flush=True)

        self.train_unlabeled_loader.dataset.pseudo_label = new_pseudo_label[
            num_train_data:]
        self.train_unlabeled_loader.dataset.ood_conf = new_ood_conf[
            num_train_data:]

        print('Randomize Auxiliary Head...', flush=True)
        if hasattr(self.net, 'fc_aux'):
            # reset auxiliary branch
            self.net.fc_aux.weight.data.normal_(mean=0.0, std=0.01)
            self.net.fc_aux.bias.data.zero_()
        else:
            # reset fc for unsupervised learning (baseline)
            self.net.fc.weight.data.normal_(mean=0.0, std=0.01)
            self.net.fc.bias.data.zero_()

        print(
            '# Online Clustering Completed! Duration: {:.2f}s #'.format(
                time.time() - start_time),
            flush=True,
        )


def preprocess_features(npdata, pca=256):
    """Preprocess an array of features.
    Args:
        npdata (np.array N * ndim): features to preprocess
        pca (int): dim of output
    Returns:
        np.array of dim N * pca: data PCA-reduced, whitened and L2-normalized
    """
    _, ndim = npdata.shape
    npdata = npdata.astype('float32')

    # Apply PCA-whitening with Faiss
    mat = faiss.PCAMatrix(ndim, pca, eigen_power=-0.5)
    mat.train(npdata)
    assert mat.is_trained
    npdata = mat.apply_py(npdata)

    # L2 normalization
    row_sums = np.linalg.norm(npdata, axis=1)
    npdata = npdata / row_sums[:, np.newaxis]

    return npdata


def run_kmeans(x, nmb_clusters, verbose=False):
    """Runs kmeans on 1 GPU.

    Args:
        x: data
        nmb_clusters (int): number of clusters
    Returns:
        list: ids of data in each cluster
    """
    n_data, d = x.shape

    # faiss implementation of k-means
    clus = faiss.Clustering(d, nmb_clusters)

    # Change faiss seed at each k-means so that the randomly picked
    # initialization centroids do not correspond to the same feature ids
    # from an epoch to another.
    clus.seed = np.random.randint(1234)

    clus.niter = 20
    clus.max_points_per_centroid = 10000000
    res = faiss.StandardGpuResources()
    flat_config = faiss.GpuIndexFlatConfig()
    flat_config.useFloat16 = False
    flat_config.device = 0
    index = faiss.GpuIndexFlatL2(res, d, flat_config)

    # perform the training
    clus.train(x, index)
    _, I = index.search(x, 1)
    return I.reshape(-1, )


class KMeans(object):
    def __init__(self, k, pca_dim):
        self.k = k
        self.pca_dim = pca_dim

    def cluster(self, data, verbose=True):
        """Performs k-means clustering.
        Args:
            x_data (np.array N * dim): data to cluster
        """
        # PCA-reducing, whitening and L2-normalization
        xb = preprocess_features(data, pca=self.pca_dim)

        if np.isnan(xb).any():
            row_sums = np.linalg.norm(data, axis=1)
            data_norm = data / row_sums[:, np.newaxis]
            if np.isnan(data_norm).any():
                I = run_kmeans(data_norm, self.k, verbose)
            else:
                I = run_kmeans(data, self.k, verbose)
        else:
            # cluster the data
            I = run_kmeans(xb, self.k, verbose)
        return I


def sort_array(old_array, index_array):
    sorted_array = np.ones_like(old_array)
    sorted_array[index_array] = old_array
    return sorted_array

openood/trainers/udg_trainer.py
kai422/SCALE
from torch.utils.data import DataLoader

from openood.utils import Config

from .base_trainer import BaseTrainer


def get_trainer(net, train_loader: DataLoader, val_loader: DataLoader,
                config: Config):
    if type(train_loader) is DataLoader:
        trainers = {
            'base': BaseTrainer,
        }
        return trainers[config.trainer.name](net, train_loader, config)

openood/trainers/utils.py
kai422/SCALE
import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions.multivariate_normal import MultivariateNormal
from tqdm import tqdm

import openood.utils.comm as comm
from openood.utils import Config


def cosine_annealing(step, total_steps, lr_max, lr_min):
    return lr_min + (lr_max -
                     lr_min) * 0.5 * (1 + np.cos(step / total_steps * np.pi))


class VOSTrainer:
    def __init__(self, net, train_loader, config: Config):
        self.train_loader = train_loader
        self.config = config
        self.net = net
        weight_energy = torch.nn.Linear(config.num_classes, 1).cuda()
        torch.nn.init.uniform_(weight_energy.weight)
        self.logistic_regression = torch.nn.Linear(1, 2).cuda()
        self.optimizer = torch.optim.SGD(
            list(net.parameters()) + list(weight_energy.parameters()) +
            list(self.logistic_regression.parameters()),
            config.optimizer['lr'],
            momentum=config.optimizer['momentum'],
            weight_decay=config.optimizer['weight_decay'],
            nesterov=True)
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: cosine_annealing(
                step, config.optimizer['num_epochs'] * len(train_loader), 1,
                1e-6 / config.optimizer['lr']))
        self.number_dict = {}
        for i in range(self.config['num_classes']):
            self.number_dict[i] = 0
        self.data_dict = torch.zeros(self.config['num_classes'],
                                     self.config['sample_number'],
                                     self.config['feature_dim']).cuda()

    def train_epoch(self, epoch_idx):
        self.net.train()
        loss_avg = 0.0
        sample_number = self.config['sample_number']
        num_classes = self.config['num_classes']
        train_dataiter = iter(self.train_loader)
        eye_matrix = torch.eye(self.config['feature_dim'], device='cuda')

        for train_step in tqdm(range(1,
                                     len(train_dataiter) + 1),
                               desc='Epoch {:03d}'.format(epoch_idx),
                               position=0,
                               leave=True,
                               disable=not comm.is_main_process()):
            batch = next(train_dataiter)
            images = batch['data'].cuda()
            labels = batch['label'].cuda()

            x, output = self.net.forward(images, return_feature=True)

            sum_temp = 0
            for index in range(num_classes):
                sum_temp += self.number_dict[index]
            lr_reg_loss = torch.zeros(1).cuda()[0]
            if (sum_temp == num_classes * sample_number
                    and epoch_idx < self.config['start_epoch']):
                target_numpy = labels.cpu().data.numpy()
                for index in range(len(labels)):
                    dict_key = target_numpy[index]
                    self.data_dict[dict_key] = torch.cat(
                        (self.data_dict[dict_key][1:],
                         output[index].detach().view(1, -1)), 0)
            elif (sum_temp == num_classes * sample_number
                  and epoch_idx >= self.config['start_epoch']):
                target_numpy = labels.cpu().data.numpy()
                for index in range(len(labels)):
                    dict_key = target_numpy[index]
                    self.data_dict[dict_key] = torch.cat(
                        (self.data_dict[dict_key][1:],
                         output[index].detach().view(1, -1)), 0)
                for index in range(num_classes):
                    if index == 0:
                        X = self.data_dict[index] - self.data_dict[index].mean(
                            0)
                        mean_embed_id = self.data_dict[index].mean(0).view(
                            1, -1)
                    else:
                        X = torch.cat((X, self.data_dict[index] -
                                       self.data_dict[index].mean(0)), 0)
                        mean_embed_id = torch.cat(
                            (mean_embed_id, self.data_dict[index].mean(0).view(
                                1, -1)), 0)

                temp_precision = torch.mm(X.t(), X) / len(X)
                temp_precision += 0.0001 * eye_matrix
                for index in range(num_classes):
                    new_dis = MultivariateNormal(
                        loc=mean_embed_id[index],
                        covariance_matrix=temp_precision)
                    negative_samples = new_dis.rsample(
                        (self.config['sample_from'], ))
                    prob_density = new_dis.log_prob(negative_samples)
                    cur_samples, index_prob = torch.topk(
                        -prob_density, self.config['select'])
                    if index == 0:
                        ood_samples = negative_samples[index_prob]
                    else:
                        ood_samples = torch.cat(
                            (ood_samples, negative_samples[index_prob]), 0)
                if len(ood_samples) != 0:

                    energy_score_for_fg = log_sum_exp(x,
                                                      num_classes=num_classes,
                                                      dim=1)
                    try:
                        predictions_ood = self.net.fc(ood_samples)
                    except AttributeError:
                        predictions_ood = self.net.module.fc(ood_samples)

                    energy_score_for_bg = log_sum_exp(predictions_ood,
                                                      num_classes=num_classes,
                                                      dim=1)

                    input_for_lr = torch.cat(
                        (energy_score_for_fg, energy_score_for_bg), -1)
                    labels_for_lr = torch.cat(
                        (torch.ones(len(output)).cuda(),
                         torch.zeros(len(ood_samples)).cuda()), -1)

                    output1 = self.logistic_regression(input_for_lr.view(
                        -1, 1))

                    lr_reg_loss = F.cross_entropy(output1,
                                                  labels_for_lr.long())
            else:
                target_numpy = labels.cpu().data.numpy()
                for index in range(len(labels)):
                    dict_key = target_numpy[index]

                    if self.number_dict[dict_key] < sample_number:
                        self.data_dict[dict_key][self.number_dict[
                            dict_key]] = output[index].detach()
                        self.number_dict[dict_key] += 1
            self.optimizer.zero_grad()
            loss = F.cross_entropy(x, labels)
            loss += self.config.trainer['loss_weight'] * lr_reg_loss
            loss.backward()

            self.optimizer.step()
            self.scheduler.step()

            loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        metrics = {}
        metrics['loss'] = loss_avg
        metrics['epoch_idx'] = epoch_idx
        return self.net, metrics


def log_sum_exp(value, num_classes=10, dim=None, keepdim=False):
    """Numerically stable implementation of the operation."""
    value.exp().sum(dim, keepdim).log()

    # TODO: torch.max(value, dim=None) threw an error at time of writing
    weight_energy = torch.nn.Linear(num_classes, 1).cuda()
    if dim is not None:
        m, _ = torch.max(value, dim=dim, keepdim=True)
        value0 = value - m
        if keepdim is False:
            m = m.squeeze(dim)

        output = m + torch.log(
            torch.sum(F.relu(weight_energy.weight) * torch.exp(value0),
                      dim=dim,
                      keepdim=keepdim))
        # set lower bound
        out_list = output.cpu().detach().numpy().tolist()
        for i in range(len(out_list)):
            if out_list[i] < -1:
                out_list[i] = -1
            else:
                continue
        output = torch.Tensor(out_list).cuda()
        return output
    else:
        m = torch.max(value)
        sum_exp = torch.sum(torch.exp(value - m))
        # if isinstance(sum_exp, Number):
        #     return m + math.log(sum_exp)
        # else:
        return m + torch.log(sum_exp)

openood/trainers/vos_trainer.py
kai422/SCALE
from .config import Config, setup_config
from .launch import launch
from .logger import setup_logger

openood/utils/__init__.py
kai422/SCALE
import argparse
import os
import re

import yaml


def setup_config(config_process_order=('merge', 'parse_args', 'parse_refs')):
    """Parsing configuration files and command line augments.

    This method reads the command line to
        1. extract and stack YAML config files,
        2. collect modification in command line arguments,
    so that the finalized configuration file is generated.

    Note:
        The default arguments allow the following equivalent code:
            config = merge_configs(*config)
                --> merge multiple YAML config files
            config.parse_args(unknown_args)
                --> use command line arguments to overwrite default settings
            config.parse_refs()
                --> replace '@{xxx.yyy}'-like values with referenced values
        It is recommended to merge before parse_args so that the latter configs
        can re-use references in the previous configs.
        For example, if
            config1.key1 = jkyang
            config1.key2 = '@{key1}'
            config2.key1 = yzang
            config3 = merge_configs(config1, config2)
            config3.parse_refs()
        then
            config3.key2 will be yzang rather than jkyang

    Return:
        An object of <class 'openood.utils.config.Config'>.
        Can be understanded as a dictionary.
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--config', dest='config', nargs='+', required=True)
    opt, unknown_args = parser.parse_known_args()
    config = [Config(path) for path in opt.config]

    for process in config_process_order:
        if process == 'merge':
            config = merge_configs(*config)
        elif process == 'parse_args':
            if isinstance(config, Config):
                config.parse_args(unknown_args)
            else:
                for cfg in config:
                    cfg.parse_args(unknown_args)
        elif process == 'parse_refs':
            if isinstance(config, Config):
                config.parse_refs()
            else:
                for cfg in config:
                    cfg.parse_refs()
        else:
            raise ValueError('unknown config process name: {}'.format(process))

    # manually modify 'output_dir'
    config.output_dir = os.path.join(config.output_dir, config.exp_name)

    return config


def parse_config(config):
    config_process_order = ('merge', 'parse_refs')
    for process in config_process_order:
        if process == 'merge':
            config = merge_configs(*config)
        elif process == 'parse_refs':
            if isinstance(config, Config):
                config.parse_refs()
            else:
                for cfg in config:
                    cfg.parse_refs()
        else:
            raise ValueError('unknown config process name: {}'.format(process))
    # manually modify 'output_dir'
    config.output_dir = os.path.join(config.output_dir, config.exp_name)

    return config


class Config(dict):
    def __init__(self, *args, **kwargs):
        super(Config, self).__init__()
        for arg in args:
            if arg == ' ':
                continue  # hard code remove white space in config file list
            if isinstance(arg, str):
                if arg.endswith('.yml'):
                    with open(arg, 'r') as f:
                        raw_dict = yaml.safe_load(f)
                else:
                    raise Exception('unknown file format %s' % arg)
                init_assign(self, raw_dict)
            elif isinstance(arg, dict):
                init_assign(self, arg)
            else:
                raise TypeError('arg should be an instance of <str> or <dict>')
        if kwargs:
            init_assign(self, kwargs)

    def __call__(self, *args, **kwargs):
        return Config(self, *args, **kwargs)

    def __repr__(self, indent=4, prefix=''):
        r = []
        for key, value in sorted(self.items()):
            if isinstance(value, Config):
                r.append('{}{}:'.format(prefix, key))
                r.append(value.__repr__(indent, prefix + ' ' * indent))
            else:
                r.append('{}{}: {}'.format(prefix, key, value))
        return '\n'.join(r)

    def __setstate__(self, state):
        init_assign(self, state)

    def __getstate__(self):
        d = dict()
        for key, value in self.items():
            if type(value) is Config:
                value = value.__getstate__()
            d[key] = value
        return d

    # access by '.' -> access by '[]'
    def __getattr__(self, key):
        return self[key]

    def __setattr__(self, key, value):
        self[key] = value

    def __delattr__(self, key):
        del self[key]

    # access by '[]'
    def __getitem__(self, key):
        sub_cfg, sub_key = consume_dots(self, key, create_default=False)
        return dict.__getitem__(sub_cfg, sub_key)

    def __setitem__(self, key, value):
        sub_cfg, sub_key = consume_dots(self, key, create_default=True)
        if sub_cfg.__contains__(sub_key) and value == '_DELETE_CONFIG_':
            dict.__delitem__(sub_cfg, sub_key)
        else:
            dict.__setitem__(sub_cfg, sub_key, value)

    def __delitem__(self, key):
        sub_cfg, sub_key = consume_dots(self, key, create_default=False)
        dict.__delitem__(sub_cfg, sub_key)

    # access by 'in'
    def __contains__(self, key):
        try:
            sub_cfg, sub_key = consume_dots(self, key, create_default=False)
        except KeyError:
            return False
        return dict.__contains__(sub_cfg, sub_key)

    # traverse keys / values/ items
    def all_keys(self, only_leaf=True):
        for key in traverse_dfs(self,
                                'key',
                                continue_type=Config,
                                only_leaf=only_leaf):
            yield key

    def all_values(self, only_leaf=True):
        for value in traverse_dfs(self,
                                  'value',
                                  continue_type=Config,
                                  only_leaf=only_leaf):
            yield value

    def all_items(self, only_leaf=True):
        for key, value in traverse_dfs(self,
                                       'item',
                                       continue_type=Config,
                                       only_leaf=only_leaf):
            yield key, value

    # for command line arguments
    def parse_args(self, cmd_args=None, strict=True):
        unknown_args = []
        if cmd_args is None:
            import sys
            cmd_args = sys.argv[1:]
        index = 0
        while index < len(cmd_args):
            arg = cmd_args[index]
            err_msg = 'invalid command line argument pattern: %s' % arg
            assert arg.startswith('--'), err_msg
            assert len(arg) > 2, err_msg
            assert arg[2] != '-', err_msg

            arg = arg[2:]
            if '=' in arg:
                key, full_value_str = arg.split('=')
                index += 1
            else:
                assert len(
                    cmd_args) > index + 1, 'incomplete command line arguments'
                key = arg
                full_value_str = cmd_args[index + 1]
                index += 2
            if ':' in full_value_str:
                value_str, value_type_str = full_value_str.split(':')
                value_type = eval(value_type_str)
            else:
                value_str = full_value_str
                value_type = None

            if key not in self:
                if strict:
                    raise KeyError(key)
                else:
                    unknown_args.extend(['--' + key, full_value_str])
                    continue

            if value_type is None:
                value_type = type(self[key])

            if value_type is bool:
                self[key] = {
                    'true': True,
                    'True': True,
                    '1': True,
                    'false': False,
                    'False': False,
                    '0': False,
                }[value_str]
            else:
                self[key] = value_type(value_str)

        return unknown_args

    # for key reference
    def parse_refs(self, subconf=None, stack_depth=1, max_stack_depth=10):
        if stack_depth > max_stack_depth:
            raise Exception(
                ('Recursively calling `parse_refs` too many times'
                 'with stack depth > {}. '
                 'A circular reference may exists in your config.\n'
                 'If deeper calling stack is really needed,'
                 'please call `parse_refs` with extra argument like: '
                 '`parse_refs(max_stack_depth=9999)`').format(max_stack_depth))
        if subconf is None:
            subconf = self
        for key in subconf.keys():
            value = subconf[key]
            if type(value) is str and '@' in value:
                if value.count('@') == 1 and value.startswith(
                        '@{') and value.endswith('}'):
                    # pure reference
                    ref_key = value[2:-1]
                    ref_value = self[ref_key]
                    subconf[key] = ref_value
                else:
                    # compositional references
                    ref_key_list = re.findall("'@{(.+?)}'", value)
                    ref_key_list = list(set(ref_key_list))
                    ref_value_list = [
                        self[ref_key] for ref_key in ref_key_list
                    ]
                    origin_ref_key_list = [
                        "'@{" + ref_key + "}'" for ref_key in ref_key_list
                    ]
                    for origin_ref_key, ref_value in zip(
                            origin_ref_key_list, ref_value_list):
                        value = value.replace(origin_ref_key, str(ref_value))
                    subconf[key] = value
        for key in subconf.keys():
            value = subconf[key]
            if type(value) is Config:
                self.parse_refs(value, stack_depth + 1)


def merge_configs(*configs):
    final_config = Config()
    for i in range(len(configs)):
        config = configs[i]
        if not isinstance(config, Config):
            raise TypeError(
                'config.merge_configs expect `Config` type inputs, '
                'but got `{}`.\n'
                'Correct usage: merge_configs(config1, config2, ...)\n'
                'Incorrect usage: merge_configs([configs1, configs2, ...])'.
                format(type(config)))
        final_config = final_config(dict(config.all_items()))
    return final_config


def consume_dots(config, key, create_default):
    sub_keys = key.split('.', 1)
    sub_key = sub_keys[0]

    if sub_key in Config.__dict__:
        raise KeyError(
            '"{}" is a preserved API name, '
            'which should not be used as normal dictionary key'.format(
                sub_key))

    if not dict.__contains__(config, sub_key) and len(sub_keys) == 2:
        if create_default:
            dict.__setitem__(config, sub_key, Config())
        else:
            raise KeyError(key)

    if len(sub_keys) == 1:
        return config, sub_key
    else:
        sub_config = dict.__getitem__(config, sub_key)
        if type(sub_config) != Config:
            if create_default:
                sub_config = Config()
                dict.__setitem__(config, sub_key, sub_config)
            else:
                raise KeyError(key)
        return consume_dots(sub_config, sub_keys[1], create_default)


def traverse_dfs(root, mode, continue_type, only_leaf, key_prefix=''):
    for key, value in root.items():
        full_key = '.'.join([key_prefix, key]).strip('.')
        child_kvs = []
        if type(value) == continue_type:
            for kv in traverse_dfs(value, mode, continue_type, only_leaf,
                                   full_key):
                child_kvs.append(kv)
        # equivalent:
        # if not (len(child_kvs) > 0 and
        # type(value) == continue_type and
        # only_leaf)
        if len(child_kvs
               ) == 0 or type(value) != continue_type or not only_leaf:
            yield {
                'key': full_key,
                'value': value,
                'item': (full_key, value)
            }[mode]
        for kv in child_kvs:
            yield kv


def init_assign(config, d):
    for full_key, value in traverse_dfs(d,
                                        'item',
                                        continue_type=dict,
                                        only_leaf=True):
        sub_cfg, sub_key = consume_dots(config, full_key, create_default=True)
        sub_cfg[sub_key] = value

openood/utils/config.py
kai422/SCALE
# Copyright (c) Facebook, Inc. and its affiliates.
import logging
from datetime import timedelta

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from openood.utils import comm

__all__ = ['DEFAULT_TIMEOUT', 'launch']

DEFAULT_TIMEOUT = timedelta(minutes=30)


def _find_free_port():
    import socket

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # Binding to port 0 will cause the OS to find an available port for us
    sock.bind(('', 0))
    port = sock.getsockname()[1]
    sock.close()
    # NOTE: there is still a chance the port could be taken by other processes.
    return port


def launch(
        main_func,
        num_gpus_per_machine,
        num_machines=1,
        machine_rank=0,
        dist_url=None,
        args=(),
        timeout=DEFAULT_TIMEOUT,
):
    """Launch multi-gpu or distributed training. This function must be called
    on all machines involved in the training. It will spawn child processes
    (defined by ``num_gpus_per_machine``) on each machine.

    Args:
        main_func: a function that will be called by `main_func(*args)`
        num_gpus_per_machine (int): number of GPUs per machine
        num_machines (int): the total number of machines
        machine_rank (int): the rank of this machine
        dist_url (str): url to connect to for distributed jobs,
        including protocol e.g. "tcp://127.0.0.1:8686".
        Can be set to "auto" to automatically select a free port on localhost
        timeout (timedelta): timeout of the distributed workers
        args (tuple): arguments passed to main_func
    """
    world_size = num_machines * num_gpus_per_machine
    if world_size > 1:
        # https://github.com/pytorch/pytorch/pull/14391
        # TODO prctl in spawned processes

        if dist_url == 'auto':
            assert num_machines == 1, \
            'dist_url=auto not supported in multi-machine jobs.'
            port = _find_free_port()
            dist_url = f'tcp://127.0.0.1:{port}'
        if num_machines > 1 and dist_url.startswith('file://'):
            logger = logging.getLogger(__name__)
            logger.warning(
                'file:// is not a reliable init_method in multi-machine jobs.'
                'Prefer tcp://'
            )

        mp.spawn(
            _distributed_worker,
            nprocs=num_gpus_per_machine,
            args=(
                main_func,
                world_size,
                num_gpus_per_machine,
                machine_rank,
                dist_url,
                args,
                timeout,
            ),
            daemon=False,
        )
    else:
        main_func(*args)


def _distributed_worker(
    local_rank,
    main_func,
    world_size,
    num_gpus_per_machine,
    machine_rank,
    dist_url,
    args,
    timeout=DEFAULT_TIMEOUT,
):
    assert torch.cuda.is_available(
    ), 'cuda is not available. Please check your installation.'
    global_rank = machine_rank * num_gpus_per_machine + local_rank
    try:
        dist.init_process_group(
            backend='NCCL',
            init_method=dist_url,
            world_size=world_size,
            rank=global_rank,
            timeout=timeout,
        )
    except Exception as e:
        logger = logging.getLogger(__name__)
        logger.error('Process group URL: {}'.format(dist_url))
        raise e

    # Setup the local process group
    # (which contains ranks within the same machine)
    assert comm._LOCAL_PROCESS_GROUP is None
    num_machines = world_size // num_gpus_per_machine
    for i in range(num_machines):
        ranks_on_i = list(
            range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine))
        pg = dist.new_group(ranks_on_i)
        if i == machine_rank:
            comm._LOCAL_PROCESS_GROUP = pg

    assert num_gpus_per_machine <= torch.cuda.device_count()
    torch.cuda.set_device(local_rank)

    # synchronize is needed here to prevent a possible timeout
    # after calling init_process_group
    # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
    comm.synchronize()

    main_func(*args)

openood/utils/launch.py
kai422/SCALE
import errno
import os
import os.path as osp
import sys

import yaml

import openood.utils.comm as comm


def mkdir_if_missing(dirname):
    """Create dirname if it is missing."""
    if not osp.exists(dirname):
        try:
            os.makedirs(dirname)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise


class Logger:
    """Write console output to external text file.

    Imported from
    `<https://github.com/Cysu/open-reid/blob/master/reid/utils/logging.py>`

    Args:
        fpath (str): directory to save logging file.

    Examples:
       >>> import sys
       >>> import os.path as osp
       >>> save_dir = 'output/experiment-1'
       >>> log_name = 'train.log'
       >>> sys.stdout = Logger(osp.join(save_dir, log_name))
    """
    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            mkdir_if_missing(osp.dirname(fpath))
            self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()


def setup_logger(config):
    """generate exp directory to save configs, logger, checkpoints, etc.

    Args:
        config: all configs of the experiment
    """
    print('------------------ Config --------------------------', flush=True)
    print(config, flush=True)
    print(u'\u2500' * 70, flush=True)

    output = config.output_dir

    if config.save_output and comm.is_main_process():
        print('Output dir: {}'.format(output), flush=True)
        if osp.isdir(output):
            if config.merge_option == 'default':
                ans = input('Exp dir already exists, merge it? (y/n)')
                if ans in ['yes', 'Yes', 'YES', 'y', 'Y', 'can']:
                    save_logger(config, output)
                elif ans in ['no', 'No', 'NO', 'n', 'N']:
                    print('Quitting the process...', flush=True)
                    quit()
                else:
                    raise ValueError('Unexpected Input.')
            elif config.merge_option == 'merge':
                save_logger(config, output)
            elif config.merge_option == 'pass':
                if os.path.exists(os.path.join(config.save_output, 'ood.csv')):
                    print('Exp dir already exists, quitting the process...',
                          flush=True)
                    quit()
                else:
                    save_logger(config, output)
        else:
            save_logger(config, output)
    else:
        print('No output directory.', flush=True)

    comm.synchronize()


def save_logger(config, output):
    print('Output directory path: {}'.format(output), flush=True)
    os.makedirs(output, exist_ok=True)
    # Save config
    # FIXME: saved config file is not beautified.
    config_save_path = osp.join(output, 'config.yml')
    with open(config_save_path, 'w') as f:
        yaml.dump(config,
                  f,
                  default_flow_style=False,
                  sort_keys=False,
                  indent=2)
    # save log file
    fpath = osp.join(output, 'log.txt')
    sys.stdout = Logger(fpath)

openood/utils/logger.py
kai422/SCALE
import argparse
import os
import zipfile

import gdown

benchmarks_dict = {
    'imagenet-1k': [
        'imagenet_1k', 'ssb_hard', 'ninco', 'inaturalist', 'texture',
        'openimage_o', 
    ],
}

dir_dict = {
    'images_classic/': [ 'texture'],
    'images_largescale/': [
        'imagenet_1k',
        'species_sub',
        'ssb_hard',
        'ninco',
        'inaturalist',
        'openimage_o',
    ],
}

download_id_dict = {
    'imagenet_res50_v1.5': '15PdDMNRfnJ7f2oxW6lI-Ge4QJJH3Z0Fy',
    'benchmark_imglist': '1XKzBdWCqg3vPoj-D32YixJyJJ0hL63gP',
    'texture': '1OSz1m3hHfVWbRdmMwKbUzoU8Hg9UKcam',
    'imagenet_1k': '1i1ipLDFARR-JZ9argXd2-0a6DXwVhXEj',
    'ssb_hard': '1PzkA-WGG8Z18h0ooL_pDdz9cO-DCIouE',
    'ninco': '1Z82cmvIB0eghTehxOGP5VTdLt7OD3nk6',
    'openimage_o': '1VUFXnB_z70uHfdgJG2E_pjYOcEgqM7tE',
    'inaturalist': '1zfLfMvoUD0CUlKNnkk7LgxZZBnTBipdj',
}


def require_download(filename, path):
    for item in os.listdir(path):
        if item.startswith(filename) or filename.startswith(
                item) or path.endswith(filename):
            return False

    else:
        print(filename + ' needs download:')
        return True


def download_dataset(dataset, args):
    for key in dir_dict.keys():
        if dataset in dir_dict[key]:
            store_path = os.path.join(args.save_dir[0], key, dataset)
            if not os.path.exists(store_path):
                os.makedirs(store_path)
            break
    else:
        print('Invalid dataset detected {}'.format(dataset))
        return

    if require_download(dataset, store_path):
        print(store_path)
        if not store_path.endswith('/'):
            store_path = store_path + '/'
        gdown.download(id=download_id_dict[dataset], output=store_path)

        file_path = os.path.join(store_path, dataset + '.zip')
        with zipfile.ZipFile(file_path, 'r') as zip_file:
            zip_file.extractall(store_path)
        os.remove(file_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Download datasets and checkpoints')
    parser.add_argument('--contents',
                        nargs='+',
                        default=['datasets', 'checkpoints'])
    parser.add_argument('--datasets', nargs='+', default=['default'])
    parser.add_argument('--checkpoints', nargs='+', default=['all'])
    parser.add_argument('--save_dir',
                        nargs='+',
                        default=['./data', './results'])
    parser.add_argument('--dataset_mode', default='default')
    args = parser.parse_args()

    args.datasets = ['imagenet-1k']
    args.checkpoints = ['imagenet_res50_v1.5']

    for content in args.contents:
        if content == 'datasets':

            store_path = args.save_dir[0]
            if not store_path.endswith('/'):
                store_path = store_path + '/'
            if not os.path.exists(os.path.join(store_path,
                                               'benchmark_imglist')):
                gdown.download(id=download_id_dict['benchmark_imglist'],
                               output=store_path)
                file_path = os.path.join(args.save_dir[0],
                                         'benchmark_imglist.zip')
                with zipfile.ZipFile(file_path, 'r') as zip_file:
                    zip_file.extractall(store_path)
                os.remove(file_path)

            if args.dataset_mode == 'default' or \
                    args.dataset_mode == 'benchmark':
                for benchmark in args.datasets:
                    for dataset in benchmarks_dict[benchmark]:
                        download_dataset(dataset, args)

            if args.dataset_mode == 'dataset':
                for dataset in args.datasets:
                    download_dataset(dataset, args)

        elif content == 'checkpoints':
            if 'v1.5' in args.checkpoints[0]:
                store_path = args.save_dir[1]
            else:
                store_path = os.path.join(args.save_dir[1], 'checkpoints/')
            if not os.path.exists(store_path):
                os.makedirs(store_path)

            if not store_path.endswith('/'):
                store_path = store_path + '/'

            for checkpoint in args.checkpoints:
                if require_download(checkpoint, store_path):
                    gdown.download(id=download_id_dict[checkpoint],
                                   output=store_path)
                    file_path = os.path.join(store_path, checkpoint + '.zip')
                    with zipfile.ZipFile(file_path, 'r') as zip_file:
                        zip_file.extractall(store_path)
                    os.remove(file_path)
scripts/download.py
kai422/SCALE
import collections
import os, sys
ROOT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
sys.path.append(ROOT_DIR)
import numpy as np
import pandas as pd
import argparse
import pickle

import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision.models import ResNet50_Weights, Swin_T_Weights, ViT_B_16_Weights
from torchvision import transforms as trn
from torch.hub import load_state_dict_from_url

from openood.evaluation_api import Evaluator

from openood.networks import ResNet50


def update(d, u):
    for k, v in u.items():
        if isinstance(v, collections.abc.Mapping):
            d[k] = update(d.get(k, {}), v)
        else:
            d[k] = v
    return d


parser = argparse.ArgumentParser()
parser.add_argument('--arch',
                    default='resnet50',
                    choices=['resnet50', 'swin-t', 'vit-b-16'])
parser.add_argument('--tvs-version', default=1, choices=[1, 2])
parser.add_argument('--ckpt-path', default=None)
parser.add_argument('--tvs-pretrained', action='store_true')
parser.add_argument('--postprocessor', default='msp')
parser.add_argument('--save-csv', action='store_true')
parser.add_argument('--save-score', action='store_true')
parser.add_argument('--fsood', action='store_true')
parser.add_argument('--batch-size', default=200, type=int)
args = parser.parse_args()

if not args.tvs_pretrained:
    assert args.ckpt_path is not None
    root = '/'.join(args.ckpt_path.split('/')[:-1])
else:
    root = os.path.join(
        ROOT_DIR, 'results',
        f'imagenet_{args.arch}_tvsv{args.tvs_version}_base_default')
    if not os.path.exists(root):
        os.makedirs(root)

# specify an implemented postprocessor
# 'openmax', 'msp', 'temp_scaling', 'odin'...
postprocessor_name = args.postprocessor
# load pre-setup postprocessor if exists
if os.path.isfile(
        os.path.join(root, 'postprocessors', f'{postprocessor_name}.pkl')):
    with open(
            os.path.join(root, 'postprocessors', f'{postprocessor_name}.pkl'),
            'rb') as f:
        postprocessor = pickle.load(f)
else:
    postprocessor = None

# assuming the model is either
# 1) torchvision pre-trained; or
# 2) a specified checkpoint
if args.tvs_pretrained:
    if args.arch == 'resnet50':
        net = ResNet50()
        weights = eval(f'ResNet50_Weights.IMAGENET1K_V{args.tvs_version}')
        net.load_state_dict(load_state_dict_from_url(weights.url))
        preprocessor = weights.transforms()
    elif args.arch == 'swin-t':
        net = Swin_T()
        weights = eval(f'Swin_T_Weights.IMAGENET1K_V{args.tvs_version}')
        net.load_state_dict(load_state_dict_from_url(weights.url))
        preprocessor = weights.transforms()
    elif args.arch == 'vit-b-16':
        net = ViT_B_16()
        weights = eval(f'ViT_B_16_Weights.IMAGENET1K_V{args.tvs_version}')
        net.load_state_dict(load_state_dict_from_url(weights.url))
        preprocessor = weights.transforms()
    else:
        raise NotImplementedError
else:
    if args.arch == 'resnet50':
        if postprocessor_name == 'conf_branch':
            net = ConfBranchNet(backbone=ResNet50(), num_classes=1000)
        elif postprocessor_name == 'godin':
            backbone = ResNet50()
            net = GodinNet(backbone=backbone,
                           feature_size=backbone.feature_size,
                           num_classes=1000)
        elif postprocessor_name == 'rotpred':
            net = RotNet(backbone=ResNet50(), num_classes=1000)
        elif postprocessor_name == 'cider':
            net = CIDERNet(backbone=ResNet50(),
                           head='mlp',
                           feat_dim=128,
                           num_classes=1000)
        else:
            net = ResNet50()

        ckpt = torch.load(args.ckpt_path, map_location='cpu')
        net.load_state_dict(ckpt)
        preprocessor = trn.Compose([
            trn.Resize(256),
            trn.CenterCrop(224),
            trn.ToTensor(),
            trn.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225])
        ])
    else:
        raise NotImplementedError

net.cuda()
net.eval()

# a unified evaluator
evaluator = Evaluator(
    net,
    id_name='imagenet',  # the target ID dataset
    data_root=os.path.join(ROOT_DIR, 'data'),
    config_root=os.path.join(ROOT_DIR, 'configs'),
    preprocessor=preprocessor,  # default preprocessing
    postprocessor_name=postprocessor_name,
    postprocessor=postprocessor,
    batch_size=args.
    batch_size,  # for certain methods the results can be slightly affected by batch size
    shuffle=False,
    num_workers=8)

# load pre-computed scores if exists
if os.path.isfile(os.path.join(root, 'scores', f'{postprocessor_name}.pkl')):
    with open(os.path.join(root, 'scores', f'{postprocessor_name}.pkl'),
              'rb') as f:
        scores = pickle.load(f)
    update(evaluator.scores, scores)
    print('Loaded pre-computed scores from file.')

# save postprocessor for future reuse
if hasattr(evaluator.postprocessor, 'setup_flag'
           ) or evaluator.postprocessor.hyperparam_search_done is True:
    pp_save_root = os.path.join(root, 'postprocessors')
    if not os.path.exists(pp_save_root):
        os.makedirs(pp_save_root)

    if not os.path.isfile(
            os.path.join(pp_save_root, f'{postprocessor_name}.pkl')):
        with open(os.path.join(pp_save_root, f'{postprocessor_name}.pkl'),
                  'wb') as f:
            pickle.dump(evaluator.postprocessor, f, pickle.HIGHEST_PROTOCOL)

# the metrics is a dataframe
metrics = evaluator.eval_ood(fsood=args.fsood)

# saving and recording
if args.save_csv:
    saving_root = os.path.join(root, 'ood' if not args.fsood else 'fsood')
    if not os.path.exists(saving_root):
        os.makedirs(saving_root)

    if not os.path.isfile(
            os.path.join(saving_root, f'{postprocessor_name}.csv')):
        metrics.to_csv(os.path.join(saving_root, f'{postprocessor_name}.csv'),
                       float_format='{:.2f}'.format)

if args.save_score:
    score_save_root = os.path.join(root, 'scores')
    if not os.path.exists(score_save_root):
        os.makedirs(score_save_root)
    with open(os.path.join(score_save_root, f'{postprocessor_name}.pkl'),
              'wb') as f:
        pickle.dump(evaluator.scores, f, pickle.HIGHEST_PROTOCOL)
scripts/eval_ood_imagenet.py
kai422/SCALE
import setuptools

with open('README.md', 'r', encoding='utf-8') as fh:
    long_description = fh.read()

setuptools.setup(
    name='openood',
    version='1.5',
    author='openood dev team',
    description=
    'This package provides a unified test platform for Out-of-Distribution detection.',
    long_description=long_description,
    long_description_content_type='text/markdown',
    url='https://github.com/Jingkang50/OpenOOD',
    packages=setuptools.find_packages(),
    install_requires=[
        'torch>=1.13.1', 'torchvision>=0.13', 'scikit-learn', 'json5',
        'matplotlib', 'scipy', 'tqdm', 'pyyaml>=5.4.1', 'pre-commit',
        'opencv-python>=4.4.0.46', 'imgaug>=0.4.0', 'pandas', 'diffdist>=0.1',
        'Cython>=0.29.30', 'faiss-gpu>=1.7.2', 'gdown>=4.7.1', 'libmr>=0.1.9'
    ],
    classifiers=[
        'Programming Language :: Python :: 3',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: OS Independent',
    ],
    include_package_data=True,
)

setup.py
kai422/SCALE