Official code for TEMI: 
from .augs import *

augs/__init__.py
HHU-MMBS/TEMI-official-BMVC2023
from enum import Enum
from typing import Union, Tuple, Callable, List, Optional

import PIL
import torch
from torchvision import transforms
import torchvision.transforms.functional as TF


class InputType(Enum):
    PIL_IMAGE = 1
    TENSOR_IMAGE = 2
    EMBEDDING = 3


class AugmentationBaseClass:

    def __init__(self, trafo, num_augs=1, input_type : InputType = InputType.PIL_IMAGE):
        super().__init__()
        self.input_type = input_type
        self.trafo = trafo
        self.num_augs = num_augs

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

    def __call__(self, x):
        if self.input_type == InputType.TENSOR_IMAGE:
            x = TF.to_tensor(x)
        aug_x = [self.transform(x) for _ in range(self.num_augs)]
        if self.input_type == InputType.PIL_IMAGE:
            aug_x = [TF.to_tensor(xi) for xi in aug_x]
        return aug_x

    @staticmethod
    def init_from_trafo(trafo, input_type=InputType.PIL_IMAGE):
        def init_fn(num_augs):
            return AugmentationBaseClass(trafo, num_augs=num_augs, input_type=input_type)
        return init_fn

    @staticmethod
    def init_from_trafo_init(trafo_init, input_type=InputType.PIL_IMAGE, defaults=None):
        if defaults is None:
            defaults = {}
        defaults = defaults.copy()

        def init_fn(num_augs, **kwargs):
            kwargs = {**defaults, **kwargs}
            return AugmentationBaseClass(trafo_init(**kwargs), num_augs=num_augs,
                                         input_type=input_type)
        return init_fn


aug_type = Callable[[PIL.Image.Image], List[torch.Tensor]]
img_size_type = Union[int, Tuple[int, int]]


class AugWrapper:

    def __init__(self,
                 global_augs: aug_type,
                 vit_image_size: Optional[img_size_type] = None,
                 normalize: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
                 local_augs: Optional[aug_type] = None,
                 aug_image_size: Optional[img_size_type] = None,
                 image_size: Optional[img_size_type] = None,
                 train=True):
        if local_augs is not None:
            raise NotImplementedError('Local augs not implemented')
        self.normalize = normalize
        self.global_augs = global_augs
        if vit_image_size is not None:
            image_size = image_size if image_size is not None else vit_image_size
            aug_image_size = aug_image_size if aug_image_size is not None else vit_image_size
            self.aux_resize = transforms.Resize(image_size)
            self.aug_resize = transforms.Resize(aug_image_size)
            if train:
                self.crop = transforms.RandomCrop(vit_image_size)
            else:
                self.crop = transforms.CenterCrop(vit_image_size)
        else:
            self.crop = None
            self.aug_resize = None
            self.aux_resize = None

    def __call__(self, image1, image2, image3=None):
        if self.aux_resize is not None and image3 is not None:
            image3 = self.aux_resize(image3)
        imgs = [img for img in [image1, image2, image3] if img is not None]
        augs =[]
        if self.aug_resize is not None:
            imgs = [self.aug_resize(img) for img in imgs]
        for img in imgs:
            augs.extend(self.global_augs(img))
        if self.crop is not None:
            augs = [self.crop(img) for img in augs]
        if self.normalize is not None:
            augs = [self.normalize(img) for img in augs]
        return augs


def _auto_aug_init(policy):
    policy = getattr(transforms.AutoAugmentPolicy, policy)
    return transforms.AutoAugment(policy=policy)


def _gaussian_noise(std):
    def trafo(x):
        return x + torch.randn_like(x) * std
    return trafo


IMAGE_AUGMENTATIONS = {
    'randaug': AugmentationBaseClass.init_from_trafo_init(transforms.RandAugment,
                                                          defaults={'num_ops': 9, 'magnitude': 1}),
    'autoaug': AugmentationBaseClass.init_from_trafo_init(_auto_aug_init, defaults={'policy': 'IMAGENET'}),
    'mild': AugmentationBaseClass.init_from_trafo(transforms.Compose([
        transforms.ColorJitter(
            brightness=0.1,
            contrast=0.1,
            saturation=0.1,
            hue=0.1),
        transforms.RandomAffine(
            degrees=10,
            translate=(0.05, 0.05),
            scale=(0.95, 1.05),
            shear=5,
            interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.RandomHorizontalFlip()
    ])),
    'gaussian': AugmentationBaseClass.init_from_trafo_init(
        _gaussian_noise, defaults={'std': 0.01}, input_type=InputType.TENSOR_IMAGE),
    'none': lambda *_, **__: lambda *x: [TF.to_tensor(xi) for xi in x]
}

EMBED_AUGMENTATIONS = {
    'gaussian': AugmentationBaseClass.init_from_trafo_init(
        _gaussian_noise, defaults={'std': 0.1}, input_type=InputType.EMBEDDING),
    'none': lambda *_, **__: lambda *x: x
}

augs/augs.py
HHU-MMBS/TEMI-official-BMVC2023
import os 
os.environ['OMP_NUM_THREADS'] = '50' 
from utils import compute_metrics
from model_builders import load_embeds, available_models
from sklearn.cluster import KMeans
from tqdm import tqdm
import torch
import pandas as pd
from pathlib import Path
from argparse import ArgumentParser
import pickle

import model_builders
from gen_embeds import main as gen_embeds

def load_data(args):
    dataset = args.dataset
    arch = args.arch
    try:
        emb, targets_train = load_embeds(arch=arch, dataset=dataset, test=False, with_label=True)
        emb_test, targets_test = load_embeds(arch=arch, dataset=dataset, test=True, with_label=True)
    except:
        print("Generating embeds for ", arch, dataset)
        gen_embeds(args)
        emb, targets_train = load_embeds(arch=arch, dataset=dataset, test=False, with_label=True)
        emb_test, targets_test = load_embeds(arch=arch, dataset=dataset, test=True, with_label=True)
    if args.normalize:
        mean, std = emb.mean(dim=0), emb.std(dim=0)
        emb = (emb - mean) / std
        emb_test = (emb_test - mean) / std
    return emb, targets_train, emb_test, targets_test


def kmeans_baseline(args, dir_path, num=1):
    emb, targets_train, emb_test, targets_test = load_data(args)
    
    rows = []
    n_clusters = {
        'CIFAR100': 100,
        'CIFAR10': 10,
        'STL10': 10,
        'IN1K': 1000,
        'IN100': 100,
        'IN50': 50,
        'IN200': 200,
    }[args.dataset]
    for _ in tqdm(range(num), leave=False):
        arch_name = args.arch.replace("/", "_")
        kmeans_save_path = dir_path / f"{arch_name}_kmeans.pkl"
        if not os.path.exists(kmeans_save_path):
            # fit based on train set
            print('Fitting K-Means classifier..')
            kmeans = KMeans(n_clusters=n_clusters, n_init=2, verbose=True).fit(emb)
            # save kmeans model
            print('Saving K-Means classifier..')
            pickle.dump(kmeans, open(kmeans_save_path, "wb"))
        else:
            print("kmeans model already exists.")
            kmeans = pickle.load(open(kmeans_save_path, "rb"))

        preds = torch.tensor(kmeans.predict(emb_test))
        data = compute_metrics(targets_test, preds)
        data = dict(zip(["Accuracy", "NMI", "Adjusted NMI", "Adjusted Rand-Index"], data))
        rows.append(data)
    return pd.DataFrame(rows)


def true_label_means_baseline(args, num=1):
    emb, targets_train, emb_test, targets_test = load_data(args)
    rows = []
    n_clusters = {
        'CIFAR100': 100,
        'CIFAR10': 10,
        'STL10': 10,
        'IN1K': 1000,
        'IN100': 100,
        'IN50': 50,
        'IN200': 200,
    }[args.dataset]
    centers = torch.stack([emb[targets_train == i].mean(dim=0) for i in range(n_clusters)])
    for _ in tqdm(range(num), leave=False):
        preds = torch.cdist(centers, emb_test).argmin(dim=0)
        data = compute_metrics(targets_test, preds)
        data = dict(zip(["Accuracy", "NMI", "Adjusted NMI", "Adjusted Rand-Index"], data))
        rows.append(data)
    return pd.DataFrame(rows)


def agg_stats(df, arch):
    name = pd.DataFrame({"arch":[arch]})
    means = df.mean().round(2).to_frame().T
    means.rename(columns = {'Accuracy':'Mean Acc', "NMI":"Mean NMI","Adjusted NMI":"Mean ANMI","Adjusted Rand-Index":"Mean ARI"}, inplace = True)
    stds = df.std().round(2).to_frame().T
    stds.rename(columns = {'Accuracy':'Std Acc', "NMI":"Std NMI","Adjusted NMI":"Std ANMI","Adjusted Rand-Index":"Std ARI"}, inplace = True)
    return pd.concat([name, means, stds], axis=1)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--dataset', choices=['CIFAR100', 'CIFAR10', "STL10", \
                                                                "CIFAR20", "IN1K", "IN50", 'IN100', "IN200", "IN1K"], type=str)
    parser.add_argument('--archs', nargs='+', default=available_models())
    parser.add_argument('--outpath', type=Path, default=Path('data'))
    parser.add_argument('--batch_size', type=int, default=256)
    parser.add_argument('--datapath', default='./data', type=str)
    parser.add_argument('--normalize', action='store_true', default=False)
    args = parser.parse_args()
    file_prefix = 'kmeans_normalized' if args.normalize else 'kmeans'
    dsets = [args.dataset]
    list_df_kmeans = []
    list_df_true_means = []
    for arch in tqdm(args.archs):
        arch_name = arch.replace("/", "_")
        for dset in dsets:
            dir_path = Path(f"experiments/clustering/kmeans_baseline/{dset}").expanduser().resolve()
            dir_path.mkdir(parents=True, exist_ok=True)
            args.arch = arch
            args.dataset = dset
            if os.path.isfile(dir_path / f"{file_prefix}_{arch_name}_{dset}.csv") and os.path.isfile(dir_path / f"labeled_centers_{arch_name}_{dset}.csv"):
                list_df_kmeans.append(pd.read_csv(dir_path / f"{file_prefix}_{arch_name}_{dset}.csv"))
                list_df_true_means.append(pd.read_csv(dir_path / f"{file_prefix}_{arch_name}_{dset}.csv"))
                continue
            res_kmeans = kmeans_baseline(args, dir_path, num=1)
            agg_df = agg_stats(res_kmeans, arch)
            list_df_kmeans.append(agg_df)
            agg_df.to_csv(dir_path / f"{file_prefix}_{arch_name}_{dset}.csv")

            agg_true_means = true_label_means_baseline(args, num=1)
            agg_df_true_means = agg_stats(agg_true_means, arch)
            list_df_true_means.append(agg_df_true_means)
            agg_df_true_means.to_csv(dir_path / f"labeled_centers_{arch_name}_{dset}.csv")
        
    df_all_km = pd.concat(list_df_kmeans, axis=0, ignore_index=True)
    df_all_km.to_csv(dir_path / f"{file_prefix}_ALL.csv")

    df_all_tm = pd.concat(list_df_true_means, axis=0, ignore_index=True )
    df_all_tm.to_csv(dir_path / f"labeled_centers_ALL.csv")
        

baseline_kmeans.py
HHU-MMBS/TEMI-official-BMVC2023
"""
Eval ood performance cifar100 -> cifar10
"""
import argparse
import json
import os
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from sklearn.metrics import mutual_info_score
from sklearn.metrics import normalized_mutual_info_score
from sklearn.metrics import roc_auc_score
from torch import linalg
from torch import nn
from torch.utils.data import Dataset
from tqdm import tqdm

import utils
from linear_evaluation import ModelEval
from loaders import get_dataset
from model_builders import load_model


def get_occupied_classes(train_features, plot=True , path='./', return_class_indices=False,
                                name='fig_occupied_classes_box_plot.png'):
    class_weights = train_features.mean(dim=0)
    # avg class probability
    threshhold = class_weights.mean()

    occupied_classes_idx = (class_weights> (threshhold))
    occupied_classes = occupied_classes_idx.sum()
    if plot:
        sorted_np_array = np.sort(class_weights.cpu().view(-1).numpy())[-occupied_classes:]
        ids = [i for i in range(sorted_np_array.shape[0])]
        _ = plt.bar(ids,sorted_np_array, label=name)
        plt.legend()
        #plt.savefig(os.path.join(path, name))

    return occupied_classes, occupied_classes_idx


def pk_statistics(features, config):
    pkx = (features / config.teacher_temp).softmax(dim=-1)
    labels = pkx.argmax(dim=-1)
    occupied_classes = len(torch.unique(labels))
    entropy = -(pkx * pkx.log()).sum(dim=-1).mean()
    pk = pkx.mean(dim=0)
    return pk, entropy, occupied_classes


def _kl_div(p, q):
    return (p * (p / q).log()).sum()


def _jsd(p, q):
    m = 0.5 * (p + q)
    return 0.5 * _kl_div(p, m) + 0.5 * _kl_div(q, m)


def jsd_to_train(pk_train, pk_test, pk_test_ood):
    return _jsd(pk_train, pk_test), _jsd(pk_train, pk_test_ood)


def compute_spmi(train_features, test_features, config):
    pkx = (train_features / config.teacher_temp).softmax(dim=-1)
    pk = pkx.mean(dim=0)
    pkx_test = (test_features / config.teacher_temp).softmax(dim=-1)
    spmi = (pkx**2 / pk).sum(dim=-1).mean()
    spmi_test = (pkx_test**2 / pk).sum(dim=-1).mean()
    return spmi.item(), spmi_test.item()


@torch.no_grad()
def knn_classifier(train_features, train_labels, test_features, test_labels, k, T, num_classes):
    if isinstance(train_labels, np.ndarray):
        train_labels = torch.from_numpy(train_labels).cuda()
        test_labels = torch.from_numpy(test_labels).cuda()

    train_features = nn.functional.normalize(train_features, dim=1, p=2)
    test_features = nn.functional.normalize(test_features, dim=1, p=2)

    top1, top5, total = 0.0, 0.0, 0
    train_features = train_features.t()
    num_test_images, num_chunks = test_labels.shape[0], 100
    imgs_per_chunk = num_test_images // num_chunks
    retrieval_one_hot = torch.zeros(k, num_classes).to(train_features.device)
    for idx in range(0, num_test_images, imgs_per_chunk):
        # get the features for test images
        features = test_features[
            idx : min((idx + imgs_per_chunk), num_test_images), :]
        targets = test_labels[idx : min((idx + imgs_per_chunk), num_test_images)]
        batch_size = targets.shape[0]

        # calculate the dot product and compute top-k neighbors
        similarity = torch.mm(features, train_features)
        distances, indices = similarity.topk(k, largest=True, sorted=True)

        candidates = train_labels.view(1, -1).expand(batch_size, -1)
        retrieved_neighbors = torch.gather(candidates, 1, indices)

        retrieval_one_hot.resize_(batch_size * k, num_classes).zero_()
        retrieval_one_hot.scatter_(1, retrieved_neighbors.view(-1, 1), 1)
        distances_transform = distances.clone().div_(T).exp_()

        temp = torch.mul(
                retrieval_one_hot.view(batch_size, -1, num_classes),
                distances_transform.view(batch_size, -1, 1),
            )

        probs = torch.sum(temp,1)
        _, predictions = probs.sort(1, True)

        # find the predictions that match the target
        correct = predictions.eq(targets.data.view(-1, 1))
        top1 = top1 + correct.narrow(1, 0, 1).sum().item()
        top5 = top5 + correct.narrow(1, 0, min(5, k)).sum().item()  # top5 does not make sense if k < 5
        total += targets.size(0)
    top1 = top1 * 100.0 / total
    top5 = top5 * 100.0 / total
    return top1, top5


class FeatureExtractionPipeline:
    def __init__(self, args, cache_backbone=False, datapath='./data'):
        if not args.head and cache_backbone:
            raise ValueError("head must be True if cache_backbone is True")
        self.args = args
        self.cache_backbone = cache_backbone
        self.embeds = None

        self.model, transform = load_model(self.args, head=args.head)
        self.model.cuda().eval()
        if not self.cache_backbone:
            self.model = nn.DataParallel(self.model)
        if args.lin_eval:
            backbone = self.model
            embed_dim = utils.embed_dim(self.args, self.model) 
            num_classes = 100 if args.dataset=='CIFAR100' else 10
            self.model = ModelEval(backbone, embed_dim, args.hidden_dim, args.bottleneck_dim,
                                    num_classes,args.nlayers,args.linear_only)
                
        precompute_arch = args.arch if args.precomputed else None
        dataset_train = get_dataset(args.dataset, datapath=datapath,
                                    train=True,
                                    download=True, transform=transform,
                                    precompute_arch=precompute_arch)
        dataset_val = get_dataset(args.dataset, datapath=datapath,
                                  train=False,
                                  download=True, transform=transform,
                                  precompute_arch=precompute_arch)

        for ds_name, dataset in zip(["dataset_labels", "val_labels", "test_labels"],
                                    [dataset_train, dataset_val]):
            if dataset is not None:
                try:
                    setattr(self, ds_name, np.array(dataset.targets, dtype=np.int64))
                except AttributeError:
                    setattr(self, ds_name, np.array(dataset.labels, dtype=np.int64))
            else:
                setattr(self, ds_name, None)

        batch_size = args.batch_size_per_gpu or args.batch_size
        self.data_loader_train = torch.utils.data.DataLoader(
            dataset_train,
            shuffle=False,
            batch_size=batch_size,
            num_workers=args.num_workers,
            pin_memory=True,
            drop_last=False,
        )
        self.data_loader_val = torch.utils.data.DataLoader(
            dataset_val,
            shuffle=False,
            batch_size=batch_size,
            num_workers=args.num_workers,
            pin_memory=True,
            drop_last=False,
        )


    @property
    def cached(self):
        return self.embeds is not None

    @torch.no_grad()
    def get_features(self, pretrained_weights):
        if isinstance(self.model, nn.DataParallel):
            module = self.model.module
        else:
            module = self.model
        if self.args.lin_eval:
            utils.load_pretrained_weights(module,
                                          pretrained_weights,
                                          self.args.checkpoint_key,
                                          head=True,
                                          head_only=False)
        else:
            utils.load_pretrained_weights(module,
                                          pretrained_weights,
                                          self.args.checkpoint_key,
                                          head=self.args.head,
                                          head_only=self.cached)
        self.model.eval()

        if not self.cache_backbone:
            train_features = extract_features(self.model, self.data_loader_train, self.args.head)
            test_features = extract_features(self.model, self.data_loader_val, self.args.head)
        else:
            embeds = self.get_embeds()
            train_features = self.head_features(embeds['train'])
            test_features = self.head_features(embeds['test'])

        print(f"Train feats {train_features.shape}")
        return train_features, test_features, self.dataset_labels, self.val_labels

    @torch.no_grad()
    def head_features(self, embed):
        features = []
        for samples in embed:
            samples = samples.cuda()
            features.append(self.model.head(samples).cpu())
        return torch.cat(features)

    @torch.no_grad()
    def get_embeds(self):
        if self.embeds is not None:
            print('Using cached embedding')
            return self.embeds
        print('Compute embeddings')
        embeds = defaultdict(list)
        for k, loader in zip(['train', 'test'],
                             [self.data_loader_train, self.data_loader_val]):
            if loader is None:
                continue
            for samples, _ in tqdm(loader):
                if not isinstance(samples, list):
                    samples = [samples]
                samples = torch.cat([im.cuda(non_blocking=True) for im in samples])
                output = self.model.backbone_embed(samples)
                embeds[k].append(output.cpu())
        self.embeds = embeds
        return embeds


@torch.no_grad()
def extract_features(model, data_loader, head):
    features = []
    for (samples, _) in tqdm(data_loader):
        if not isinstance(samples, list):
            samples = [samples]
        samples = torch.cat([im.cuda(non_blocking=True) for im in samples])
        if head:
            try:
                feats, _ = model(samples)
            except Exception:
                feats = model(samples)
        else:
            feats = model(samples)
        features.append(feats.cpu())
    return torch.cat(features, dim=0)


def norm_by_name(x, norm):
    if norm == "softmax":
        return x.softmax(dim=-1)
    elif norm == "l2":
        return torch.nn.functional.normalize(x, dim=-1, p=2)
    elif norm == "l1":
        return torch.nn.functional.normalize(x, dim=-1, p=1)
    return x


def norm_feats(*features, norm="softmax"):
    return [norm_by_name(x, norm) for x in features]

@torch.no_grad()
def calc_maha_distance(embeds, means_c, inv_cov_c):
    diff = embeds - means_c
    #dist = np.matmul(np.matmul(diff, inv_cov_c),diff.T)
    dist = np.matmul(diff,inv_cov_c)*diff
    dist = np.sum(dist,axis=1)
    return dist

@torch.no_grad()
def OOD_classifier_maha(train_embeds_in, train_labels_in, test_embeds_in, test_embeds_outs, num_classes,
                        relative=False, std_all=False):
    class_covs = []
    class_means = []
    used_classes = 0
    # calculate class-wise means and covariances
    for c in range(num_classes):
        train_embeds_c = train_embeds_in[np.where(train_labels_in == c)]
        if len(train_embeds_c)>1:
            class_mean = np.mean(train_embeds_c, axis = 0)
            cov = np.cov((train_embeds_c - (class_mean.reshape([1,-1]))).T )
            class_covs.append(cov)
            class_means.append(class_mean)
            used_classes += 1


    # class-wise std estimation
    if not std_all:
        cov_invs = np.linalg.inv(np.mean(np.stack(class_covs, axis=0),axis=0))
    else:
        # estimating the global std from train data
        avg_train_mean = np.mean(train_embeds_in, axis=0)
        cov_invs =  np.linalg.inv(np.cov((train_embeds_in-avg_train_mean.reshape([1,-1])).T))

    scores_in_dist = [calc_maha_distance(test_embeds_in, class_means[c], cov_invs) for c in range(used_classes)]
    scores_out_dist = [calc_maha_distance(test_embeds_outs, class_means[c], cov_invs) for c in range(used_classes)]

    # classes X num_data
    scores_in_dist = np.stack(scores_in_dist)
    scores_out_dist = np.stack(scores_out_dist)

    if relative == True:
        avg_train_mean = np.mean(train_embeds_in, axis=0)
        avg_train_inv_cov =  np.linalg.inv(np.cov((train_embeds_in-avg_train_mean.reshape([1,-1])).T))
        avg_train_score_in = calc_maha_distance(test_embeds_in, avg_train_mean , avg_train_inv_cov )
        avg_train_score_out = calc_maha_distance(test_embeds_outs, avg_train_mean , avg_train_inv_cov)
        scores_in_dist -= avg_train_score_in
        scores_out_dist -= avg_train_score_out

    # Get OOD score for each datapoint
    scores_in_dist = -np.min(scores_in_dist, axis=0)
    scores_out_dist = -np.min(scores_out_dist, axis=0)

    scores = np.concatenate([scores_in_dist, scores_out_dist])
    return scores


def get_eval_args(notebook=False):
    parser = argparse.ArgumentParser('Evaluation')
    parser.add_argument('--batch_size_per_gpu', default=512, type=int, help='Per-GPU batch-size')
    parser.add_argument('--temperature', default=0.02, type=float,
        help='Temperature used in the voting coefficient')
    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
    parser.add_argument('--out_dim', default=100, type=int)
    parser.add_argument('--head', default=True, type=utils.bool_flag,
        help="Whether to load the DINO head")
    parser.add_argument('--vit_image_size', type=int, default=224, help="""image size that enters vit; 
        must match with patch_size: num_patches = (vit_image_size/patch_size)**2""")
    parser.add_argument('--datapath', default='./data', type=str)

    parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag,
        help="Whether to use batch normalizations in projection head (Default: False)")
    parser.add_argument('--norm_last_layer', default=False, type=utils.bool_flag,
        help="""Whether or not to weight normalize the last layer of the DINO head.
        Not normalizing leads to better performance but can make the training unstable.
        In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""")

    parser.add_argument('--dataset', default='CIFAR100', choices=['CIFAR100', 'CIFAR10', "STL10", \
                                                                "CIFAR20", "IN1K", "IN50", 'IN100', "IN200"], type=str)
    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag,
        help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM")
    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
    parser.add_argument("--checkpoint_key", default="teacher", type=str,
        help='Key to use in the checkpoint (example: "teacher")')
    parser.add_argument('--num_workers', default=8, type=int, help='Number of data loading workers per GPU.')
    parser.add_argument('--image_size', type=int, default=32, help="""image size that enters vit; 
        must match with patch_size: num_patches = (vit_image_size/patch_size)**2""")
    parser.add_argument('--crops_scale', type=float, nargs='+', default=(0.8, 1.),
        help="""Scale range of the cropped image before resizing, relatively to the origin image.
        Used for large global view cropping.""")
    parser.add_argument('--crops_number', type=int, default=1, help="""Number of
        local views to generate. Set this parameter to 0 to disable multi-crop training.""")

    parser.add_argument('--nlayers', default=2, type=int, help='Head layers')
    parser.add_argument('--hidden_dim', default=512, type=int, help="Head's hidden dim")
    parser.add_argument('--bottleneck_dim', default=256, type=int, help="Head's bottleneck dim")
    parser.add_argument('--l2_norm', default=False, help="Whether to apply L2 norm after backbone") 
    parser.add_argument('--ckpt_folder', type=str) 
    parser.add_argument('--no_cache', action='store_true', default=False, help='Whether to cache backbone results')
    parser.add_argument('--ignore_hp_file', action='store_true', default=False, help='Whether to ignore hp.json')
    parser.add_argument('--lin_eval', default=False, help='True if the model has a backbone and linear layer.')
    parser.add_argument('--linear_only', default=False, type=utils.bool_flag, help='True if head is only a linear layer.')
    return parser.parse_args() if notebook is False else parser.parse_args("") 

eval_cluster_utils.py
HHU-MMBS/TEMI-official-BMVC2023
import glob
from functools import partial
import pandas as pd
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
from eval_cluster_utils import *


def plot_scatter(x_axis, values, outdir, fname, xlab="epoch",
                 ylab="AUROC", title="OOD AUROC & score CIFAR100 -> CIFAR10"):
    plt.figure(figsize=(10,10))
    plt.plot(x_axis, values, "-o")
    plt.title(title)
    plt.xlabel(xlab)
    plt.ylabel(ylab)
    plt.savefig(Path(outdir) / fname)


def _eval_setting_to_str(s):
    if not isinstance(s, tuple):
        return str(s)
    s = [str(x) for x in s]
    return '-'.join(s)


def print_results(d):
    for k, d_inner in d.items():
        print(k)
        for k_inner, v in d_inner.items():
            s = f'{_eval_setting_to_str(k_inner)}:'
            print(f'\t{s:<22} {v[-1]:.2f}')


def load_tensorboard_loss(path):
    tag = 'Train loss epoch'
    event_acc = EventAccumulator(str(next(path.glob('event*'))))
    event_acc.Reload()
    if tag in event_acc.Tags()['scalars']:
        return pd.DataFrame([{'Epoch': ev.step, 'loss': ev.value}
                             for ev in event_acc.Scalars(tag)]).set_index('Epoch')
    # Multihead case
    dfs = []
    for p in path.rglob('Train loss*/event*'):
        event_acc = EventAccumulator(str(p))
        event_acc.Reload()
        dfs.append(pd.DataFrame([{'Epoch': ev.step, 'loss': ev.value}
                                 for ev in event_acc.Scalars(tag)]).set_index('Epoch'))
    df = pd.concat(dfs)
    return df.groupby('Epoch').min()


def main():
    args = get_eval_args()
    cudnn.deterministic = True

    auroc_results = defaultdict(partial(defaultdict, list))
    cluster_results = {"cluster_acc": [], "nmi": [], "anmi": [], "ari": [],
                       "cluster_acc-train": [], "nmi-train": [], "anmi-train": [], "ari-train": []}
    loss_results = {"train_loss": []}

    checkpoint_list = glob.glob(os.path.join(args.ckpt_folder, "*.pth"))
    outdir = Path(args.ckpt_folder).expanduser().resolve()

    # Read hparams
    with open(outdir / 'hp.json', 'r') as f:
        hparams = json.load(f)
    if not args.ignore_hp_file:
        args.__dict__.update({k: v for k, v in hparams.items() if v is not None})

    # Load loss history
    losses_df = load_tensorboard_loss(outdir)

    # replace last saved checkpoint name to be last
    checkpoint_list = list(map(lambda st: str.replace(st, "checkpoint.pth", "checkpoint9999.pth"), checkpoint_list))
    checkpoint_list = sorted(checkpoint_list)
    checkpoint_list = list(map(lambda st: str.replace(st, "checkpoint9999.pth", "checkpoint.pth", ), checkpoint_list))
    epochs = []

    print(f"dataset: {args.dataset} \n Checkpoints found {len(checkpoint_list)}  \n {checkpoint_list} ")
    assert len(checkpoint_list) >= 1
    args.datapath = './data' if  args.dataset in ["CIFAR10", "CIFAR100", "STL10", "CIFAR20"] else args.datapath
    extractor = None
    for ckpt in checkpoint_list:
        print(ckpt)
        # Epoch number for next epoch is saved in the checkpoint
        epoch = torch.load(ckpt, map_location='cpu')['epoch'] - 1
        epochs.append(epoch)
        if extractor is None or args.no_cache:
            extractor = FeatureExtractionPipeline(args, cache_backbone=not args.no_cache, datapath=args.datapath)
        train_features, test_features, train_labels, val_labels = \
            extractor.get_features(ckpt)

        # Cluster performance test
        ( _ , max_indices) = torch.max(test_features, dim=1)
        max_indices = max_indices.cpu().numpy()
        cluster_acc, nmi, anmi, ari = utils.compute_metrics(val_labels, max_indices, min_samples_per_class=5)

        cluster_results["cluster_acc"].append(cluster_acc)
        cluster_results["nmi"].append(nmi)
        cluster_results["anmi"].append(anmi)
        cluster_results["ari"].append(ari)

        # Cluster performance train
        ( _ , max_indices) = torch.max(train_features, dim=1)
        max_indices = max_indices.cpu().numpy()
        cluster_acc, nmi, anmi, ari = utils.compute_metrics(train_labels, max_indices, min_samples_per_class=5)

        cluster_results["cluster_acc-train"].append(cluster_acc)
        cluster_results["nmi-train"].append(nmi)
        cluster_results["anmi-train"].append(anmi)
        cluster_results["ari-train"].append(ari)

        # Loss
        if epoch in losses_df.index:
            train_loss = losses_df.loc[epoch].item()
            loss_results["train_loss"].append(train_loss)
        else:
            loss_results["train_loss"].append(np.nan)

        print('\n', '-'*100, '\n')

        dict_data = {
                    "cluster_val_acc" : np.max(cluster_results["cluster_acc"]),
                    "NMI" : np.max(cluster_results["nmi"]),
                    "ARI" : np.max(cluster_results["ari"]),
                    "ckpt-best-cluster-acc": checkpoint_list[np.argmax(cluster_results["cluster_acc"])],
                    }
        cluster_results.update(loss_results)
        df = pd.DataFrame(cluster_results, index=epochs)
        df.index.name = "Epoch"
        print(df[["cluster_acc",
                "nmi",
                "ari"
                ]])
    with open(outdir / "best-results.json", 'w') as f:
        json.dump(dict_data, f, indent=4)
    
    df.to_csv(outdir / "checkpoint_metrics.csv")


if __name__ == '__main__':
    main()

eval_experiment.py
HHU-MMBS/TEMI-official-BMVC2023
"""
Precomputes embeddings for a given model and dataset.
"""
import json
from argparse import ArgumentParser
from pathlib import Path

import torch
from torch.backends import cudnn
from tqdm import tqdm

from eval_cluster_utils import knn_classifier
from loaders import get_dataset
from model_builders import load_model


@torch.no_grad()
def compute_embedding(model, loader):
    embeds = []
    labels = []
    for images, label in tqdm(loader):
        images = images.cuda()
        image_features = model(images).float()
        embeds.append(image_features.cpu())
        labels.append(label)
    return torch.cat(embeds), torch.cat(labels)

@torch.no_grad()
def compute_neighbors(embedding, k):
    embedding = embedding / embedding.norm(p=2, dim=-1, keepdim=True)
    num_embeds = embedding.shape[0]
    if num_embeds <= 8*1e4:
        dists = embedding @ embedding.permute(1, 0)
        # exclude self-similarity
        dists.fill_diagonal_(-torch.inf)
        return dists.topk(k, dim=-1)   
    else:
        topk_knn_ids = []
        topk_knn_dists = []
        print("Chunk-wise implementation of k-nn in GPU")
        # num_chunks = 12000 
        step_size = 64 # num_embeds // num_chunks
        embedding = embedding.cuda()
        for idx in tqdm(range(0, num_embeds, step_size)):
            idx_next_chunk = min((idx + step_size), num_embeds)
            features = embedding[idx : idx_next_chunk, :]
            # calculate the dot product dist
            dists_chunk = torch.mm(features, embedding.T).cpu()
            dists_chunk.fill_diagonal_(-torch.inf)
            max_dists, indices = dists_chunk.topk(k, dim=-1)
            topk_knn_ids.append(indices)
            topk_knn_dists.append(max_dists)
        return torch.cat(topk_knn_dists), torch.cat(topk_knn_ids)
    
        
def get_outpath(arch, dataset, datapath='data'):
    datapath = Path(datapath).expanduser().resolve()
    arch = arch.replace('/', '_')
    dataset = dataset.replace('/', '_')
    return datapath / 'embeddings' / f'{dataset}-{arch}'


def get_nn(args, preprocess, model, test=False):
    datapath = './data' if  args.dataset in ["CIFAR10", "CIFAR100", "STL10", "CIFAR20"] else args.datapath
    dset = get_dataset(args.dataset, datapath=datapath, train=not test, transform=preprocess, download=True)
    dataloader = torch.utils.data.DataLoader(dset, batch_size=args.batch_size, shuffle=False, drop_last=False, pin_memory=True, num_workers=16)
    embeddings, label = compute_embedding(model, dataloader)    
    embeddings = embeddings.squeeze()
    k = args.k or len(dset) // len(dset.classes)
    nn_dists, neighbors = compute_neighbors(embeddings, k)
    return embeddings, label, nn_dists, neighbors, len(dset.classes)

def compute_stats(outpath):
    for test in True, False:
        test_str = '-test' if test else ''
        embeddings = torch.load(outpath / f'embeddings{test_str}.pt', map_location='cpu')
        torch.save(embeddings.mean(dim=0), outpath / f'mean{test_str}.pt')
        torch.save(embeddings.std(dim=0), outpath / f'std{test_str}.pt')

def main(args):
    cudnn.benchmark = True
    cudnn.deterministic = True
    modelname = args.arch

    outpath = get_outpath(modelname, args.dataset)
    if args.stats_only:
        compute_stats(outpath)
        return

    model, preprocess = load_model(args, head=False)
    model = model.cuda()
    model.eval()

    outpath.mkdir(parents=True, exist_ok=True)

    embs = {}
    labels = {}
    
    for test in True, False:
        print('Computing', 'test' if test else 'train', 'dataset embedding')
        embeddings, label, nn_dists, neighbors, num_classes = get_nn(args, preprocess, model, test)
        embeddings, label, nn_dists, neighbors = embeddings.cpu(), label.cpu(), nn_dists.cpu(), neighbors.cpu()

        embs[test] = embeddings
        labels[test] = label
        test_str = '-test' if test else ''
        torch.save(embeddings, outpath / f'embeddings{test_str}.pt')
        torch.save(label, outpath / f'label{test_str}.pt')
        torch.save(neighbors, outpath / f'knn{test_str}.pt')
        torch.save(nn_dists, outpath / f'knn_dists{test_str}.pt')
        torch.save(embeddings.mean(dim=0), outpath / f'mean{test_str}.pt')
        torch.save(embeddings.std(dim=0), outpath / f'std{test_str}.pt')
    
    if not args.no_eval_knn:
        print('Computing KNN accuracy')
        top1, top5 = knn_classifier(
            train_features=embs[False],
            train_labels=labels[False],
            test_features=embs[True],
            test_labels=labels[True],
            k=args.classifier_k,
            T=args.temperature,
            num_classes=num_classes
        )
        print(f'Top-1 accuracy: {top1}, Top-5 accuracy: {top5}')
        with open(outpath / 'accuracy.json', 'w') as f:
            json.dump({'top1': top1, 'top5': top5}, f)
    # empty gpu memory
    model = model.cpu()
    del model


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--dataset', default='CIFAR100', choices=['CIFAR100', 'CIFAR10', "STL10", \
                                                                "CIFAR20", "IN1K", "IN50", 'IN100', "IN200", "IN1K"], type=str)
    parser.add_argument('--arch', default='clip_ViT-B/32')
    parser.add_argument('--outpath', type=Path, default=Path('data'))
    parser.add_argument('--temperature', default=0.02, type=float,
                        help='Temperature used in the voting coefficient')
    parser.add_argument('--classifier-k', default=20, type=int, help='Numbers of neighbors to use in the classifier')
    parser.add_argument('-k', type=int, default=None, help='total NNs to compute. Default: num images / num classes')
    parser.add_argument('--vit_image_size', type=int, default=224)
    parser.add_argument('--batch_size', type=int, default=512)
    parser.add_argument('--datapath', default='./data', type=str)
    parser.add_argument('--no_eval_knn', action='store_true', help='Do not evaluate k-nn accuracy', default=False)
    parser.add_argument('--stats_only', action='store_true',
                        help='Only compute the mean and std of the dataset for precomputed embeddings')

    main(parser.parse_args())

gen_embeds.py
HHU-MMBS/TEMI-official-BMVC2023
import os
import numpy as np
from pathlib import Path
import json

import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import DataLoader
import argparse

from model_builders import load_model
import utils
from loaders import get_dataset
from utils import embed_dim
from utils import trunc_normal_
        
class ModelEval(nn.Module):
    def __init__(self, backbone, in_dim, hidden_dim, 
                        bottleneck_dim, num_classes, num_layers, linear_only,
                        train_backbone=False, l2_norm=False):
        super(ModelEval, self).__init__()
        self.train_backbone = train_backbone
        self.linear_only = linear_only
        self.out_dim = num_classes
        self.l2_norm = l2_norm
       
        if linear_only:
            self.mlp = nn.Linear(in_dim, num_classes)
            self.mlp.weight.data.normal_(mean=0.0, std=0.01)
            self.mlp.bias.data.zero_()
        else:
            self.mlp = self.init_mlp(in_dim, hidden_dim, num_layers, bottleneck_dim)
        self.backbone = backbone
        for param in self.backbone.parameters():
            param.requires_grad = False
        
    def init_mlp(self, in_dim, hidden_dim, num_layers, bottleneck_dim):
        if num_layers == 1:
            layers = [nn.Linear(in_dim, bottleneck_dim)]
        else:
            layers = [nn.Linear(in_dim, hidden_dim)]
            layers.append(nn.GELU())
            for _ in range(num_layers - 2):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                layers.append(nn.GELU())
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
        
        for layer in layers:
            layer.apply(self._init_weights)
        last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, self.out_dim, bias=False))
        last_layer.weight_g.data.fill_(1)
        layers.append(last_layer)
        mlp = nn.Sequential(*layers)
        return mlp

    
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
 
    def forward(self, x):
        x = self.backbone(x).detach() if not self.train_backbone else self.backbone(x)
        if self.l2_norm:
            x = F.normalize(x, dim=1, p=2)
        return self.mlp(x) 
        
def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_iters=0, start_warmup_value=0):
    warmup_schedule = np.array([])
    warmup_iters = warmup_iters
    if warmup_iters > 0:
        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)

    iters = np.arange(epochs * niter_per_ep - warmup_iters)
    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))

    schedule = np.concatenate((warmup_schedule, schedule))
    assert len(schedule) == epochs * niter_per_ep
    return schedule

def validate(model, val_loader, device, targets):
    correct = 0
    total = 0
    loss_step = []
    max_indx_pred = []
    with torch.no_grad():
        for inp_data, labels in val_loader:
            labels = labels.to(device)
            inp_data = inp_data.to(device)
            outputs = model(inp_data)
            predicted = torch.max(outputs, 1)[1]
            total += labels.size(0)
            correct += (predicted == labels).sum()
            max_indx_pred.append(predicted.cpu())
    val_acc = (100 * correct / total).cpu().numpy()
    preds = torch.cat(max_indx_pred).numpy()
    cluster_acc, nmi, anmi, ari = utils.compute_metrics(targets, preds, min_samples_per_class=5)
    return val_acc, cluster_acc, nmi, anmi, ari

def apply_color_distortion(s=0.5):  
    color_jitter = transforms.ColorJitter(0.8*s, 0.8*s, 0.8*s, 0.2*s)
    rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)
    rnd_gray = transforms.RandomGrayscale(p=0.2)
    color_distort = transforms.Compose([rnd_color_jitter, rnd_gray])
    return color_distort

def get_data_loaders(args, normalize):
    datapath = './data' if  args.dataset in ["CIFAR10", "CIFAR100", "STL10", "CIFAR20"] else args.datapath
    
    if not args.weak_augs:
        transform_train = transforms.Compose([
                         transforms.Resize(224, interpolation=3),
                         normalize,
                         ])
    else:
        transform_train = transforms.Compose([
                         transforms.RandomResizedCrop(224, scale=(0.5, 1), interpolation=3),
                         transforms.Resize(224, interpolation=3),
                         transforms.RandomHorizontalFlip(),
                         apply_color_distortion(s=0.2),
                         normalize,
                         ])
    
    transform_test = transforms.Compose([
                     transforms.Resize(224,interpolation=3),
                     normalize,
                     ])
    
    dataset_train = get_dataset(args.dataset, datapath=datapath,
                        train=True,
                        download=True, transform=transform_train)
    dataset_val = get_dataset(args.dataset, datapath=datapath,
                        train=False,
                        download=True, transform=transform_test)

    train_loader = DataLoader(dataset_train,
                              batch_size=args.batch_size,
                              shuffle=True,
                              drop_last=True,
                              num_workers=4)
    
    val_loader = DataLoader(dataset_val,
                             batch_size=args.batch_size,
                             shuffle=False,
                             drop_last=False,
                             num_workers=4)   
    try:
        val_labels = np.array(dataset_val.targets,dtype=np.int64)
    except:
        val_labels = np.array(dataset_val.labels,dtype=np.int64)
    
    num_classes = len(np.unique(val_labels))
    
    return train_loader, val_loader, num_classes, val_labels
    
        
def train_one_epoch(model, train_loader, optimizer, scheduler, device):
    losses = []
    criterion = nn.CrossEntropyLoss()

    for it, data in enumerate(train_loader):
        optimizer.zero_grad()
        images, labels = data
        labels = labels.to(device)
        images = images.to(device)
        it = len(train_loader) * ep + it 
        for i, param_group in enumerate(optimizer.param_groups):
            param_group["lr"] = scheduler[it]
        logits = model(images) 
        loss = criterion(logits, labels)        
        loss.backward()
        optimizer.step()

        losses.append(loss.item())
            
    avg_loss = torch.tensor(losses).mean().numpy()
    return avg_loss


if __name__ == '__main__':
    parser = argparse.ArgumentParser('Training MLP on top of frozen backbone (supervised)')
    parser.add_argument('--load_path', default='', type=str, help='Path to pretrained weights.')
    parser.add_argument('--checkpoint_key', default='teacher', type=str, help='Key to use in the checkpoint.')
    parser.add_argument('--pretrained_weights', default=None, type=str, help='Path to pretrained model weights. ')
    parser.add_argument('--save_path', default='./experiments/finetune/', type=str, help='Path to save model checkpoint.')
    parser.add_argument('--dataset', default='CIFAR100', choices=['CIFAR100', 'CIFAR10', "STL10", \
                                                                "CIFAR20", "IN1K", "IN50", 'IN100', "IN200"], type=str)
    parser.add_argument('--datapath', default='./data', type=str)
    parser.add_argument('--batch_size', type=int, default=256, help="""Value for batch size.""")
    parser.add_argument('--lr', type=float, default=5e-3, help="""Value for learning rate.""")
    parser.add_argument('--wd', type=float, default=1e-2, help="""Value for weight decay.""")
    parser.add_argument('--num_epochs', type=int, default=100, help="""Number of training epochs.""")
    parser.add_argument('--arch', default='dino_vitb16', help="""Chosen architecture for backbone.""")
    parser.add_argument('--vit_image_size', type=int, default=224, help="""Size of images for VIT.""")
    parser.add_argument('--hidden_dim', type=int, default=512, help="""Hidden dimension in MLP.""")
    parser.add_argument('--bottleneck_dim', type=int, default=256, help="""Dimension of bottleneck in MLP.""")
    parser.add_argument('--num_layers', type=int, default=2, help="""Number of layers in MLP.""")
    parser.add_argument('--linear_head', type=bool, default=True, help="""True if head should only be a linear layer instead of MLP.""")
    parser.add_argument('--train_backbone', type=bool, default=False, help="""True if also the backbone should be trained.""")
    parser.add_argument('--l2_norm', type=bool, default=False, help="""Whether to apply L2 normalization to the output of the backbone.""")
    parser.add_argument('--weak_augs', default=False, help="""Whether to apply augmentations or not.""")
    
    args = parser.parse_args()
    args.save_path = os.path.join(args.save_path, args.dataset, f'exp_v002_adam_MLP_rrc05_{str(np.random.randint(1000)).zfill(4)}')
    output_dir = Path(args.save_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    with open(output_dir / "hp.json", 'wt') as f:
        json.dump(vars(args), f, indent=4, default=str)
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    backbone, _ , normalize = load_model(args, head=False, split_preprocess=True)
    train_loader, val_loader, num_classes, val_labels = get_data_loaders(args, normalize)
    
    if args.pretrained_weights != None: 
        utils.load_pretrained_weights(backbone,
                                      args.pretrained_weights,
                                      args.checkpoint_key,
                                      head=True,
                                      head_only=False)

    backbone = backbone.to(device)
    backbone_dim = embed_dim(args, backbone)
    
    model = ModelEval(backbone, backbone_dim, args.hidden_dim, 
                args.bottleneck_dim, num_classes, args.num_layers, args.linear_head,
                args.train_backbone)
    model = model.to(device)
    if not args.train_backbone:
        print('Training the head only \n\n')
        model.backbone.eval()
        # TODO try simple SGD here
        optimizer = torch.optim.Adam(model.mlp.parameters(), 
                        lr = args.lr,
                        weight_decay=args.wd)
    else:
        print('Training the whole model \n\n')
        optimizer = torch.optim.Adam(model.parameters(), 
                        lr = args.lr,
                        weight_decay=args.wd)
    
    lr_schedule = cosine_scheduler(args.lr, 0,args.num_epochs,len(train_loader), warmup_iters=0, start_warmup_value=0)
    
    max_val_acc = 0
    for ep in range(args.num_epochs):
        avg_loss = train_one_epoch(model= model,
                                   train_loader=train_loader, 
                                   optimizer=optimizer, 
                                   scheduler = lr_schedule,
                                   device = device)
        
        with torch.no_grad():
            val_acc, cluster_acc, nmi, anmi, ari = validate(model, val_loader, device, val_labels)
            if val_acc > max_val_acc:
                torch.save(model.state_dict(), output_dir / 'best_model.pth')
                max_val_acc = val_acc
                best_dict_data = {
                    "val_acc" : float(val_acc),
                    "cluster_acc" : float(cluster_acc),
                    "NMI" : nmi,
                    "ARI" : anmi,
                    "ANMI" : ari,
                    "epoch": ep}
                with open(output_dir / "best-results.json", 'w') as f:
                    json.dump(best_dict_data, f, indent=4)
            
        print(f'Epoch {ep} Average training loss: {avg_loss:.3}, Validation accuracy {val_acc:.3}, \
            Cluster Acc {cluster_acc:.2} Maximum Val acc: {max_val_acc:.2}')
    
    # Compute clustering metrics and save them to a file
    with open(output_dir / "best-results.json", 'w') as f:
        json.dump(best_dict_data, f, indent=4)

linear_evaluation.py
HHU-MMBS/TEMI-official-BMVC2023
from .embedNN import EmbedNN, get_dataset, TruePosNN
from .cifar20 import CIFAR20
from .imagenet import *

loaders/__init__.py
HHU-MMBS/TEMI-official-BMVC2023
from torchvision.datasets import CIFAR10, CIFAR100


# almost cp from https://github.com/wvangansbeke/Unsupervised-Classification/blob/master/data/cifar.py
class CIFAR20(CIFAR100):
    """CIFAR20 Dataset.
    This is a subclass of the `CIFAR10` Dataset.
    """
    base_folder = 'cifar-100-python'
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
    train_list = [
                    ['train', '16019d7e3df5f24257cddd939b257f8d'],
                 ]

    test_list = [
                    ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
                ]
    meta = {
        'filename': 'meta',
        'key': 'fine_label_names',
        'md5': '7973b15100ade9c7d40fb424638fde48',
    }
    def __init__(self, root='./data', train=True, transform=None, 
                    download=False):
        super(CIFAR20, self).__init__(root, train=train,transform=transform,
                                        download=download)
        # Remap classes from cifar-100 to cifar-20
        new_ = self.targets
        for idx, target in enumerate(self.targets):
            new_[idx] = _cifar100_to_cifar20(target)
        self.targets = new_
        self.classes = ['aquatic mammals', 'fish', 'flowers', 'food containers', 'fruit and vegetables', 'household electrical devices', 'househould furniture', 'insects', 'large carnivores', 'large man-made outdoor things', 'large natural outdoor scenes', 'large omnivores and herbivores', 'medium-sized mammals', 'non-insect invertebrates', 'people', 'reptiles', 'small mammals', 'trees', 'vehicles 1', 'vehicles 2']


def _cifar100_to_cifar20(target):
  _dict = \
    {0: 4,
     1: 1,
     2: 14,
     3: 8,
     4: 0,
     5: 6,
     6: 7,
     7: 7,
     8: 18,
     9: 3,
     10: 3,
     11: 14,
     12: 9,
     13: 18,
     14: 7,
     15: 11,
     16: 3,
     17: 9,
     18: 7,
     19: 11,
     20: 6,
     21: 11,
     22: 5,
     23: 10,
     24: 7,
     25: 6,
     26: 13,
     27: 15,
     28: 3,
     29: 15,
     30: 0,
     31: 11,
     32: 1,
     33: 10,
     34: 12,
     35: 14,
     36: 16,
     37: 9,
     38: 11,
     39: 5,
     40: 5,
     41: 19,
     42: 8,
     43: 8,
     44: 15,
     45: 13,
     46: 14,
     47: 17,
     48: 18,
     49: 10,
     50: 16,
     51: 4,
     52: 17,
     53: 4,
     54: 2,
     55: 0,
     56: 17,
     57: 4,
     58: 18,
     59: 17,
     60: 10,
     61: 3,
     62: 2,
     63: 12,
     64: 12,
     65: 16,
     66: 12,
     67: 1,
     68: 9,
     69: 19,
     70: 2,
     71: 10,
     72: 0,
     73: 1,
     74: 16,
     75: 12,
     76: 9,
     77: 13,
     78: 15,
     79: 13,
     80: 16,
     81: 19,
     82: 2,
     83: 4,
     84: 6,
     85: 19,
     86: 5,
     87: 5,
     88: 8,
     89: 19,
     90: 18,
     91: 1,
     92: 2,
     93: 15,
     94: 6,
     95: 0,
     96: 17,
     97: 8,
     98: 14,
     99: 13}

  return _dict[target]
loaders/cifar20.py
HHU-MMBS/TEMI-official-BMVC2023
import random

import numpy as np
from torch.utils.data import Dataset
import torch
import torchvision.datasets as tds
from pathlib import Path

import loaders
import model_builders


class PrecomputedEmbeddingDataset(Dataset):

    def __init__(self, dataset, arch, train, datapath):
        super().__init__()
        self.emb, self.targets = model_builders.load_embeds(
            arch=arch,
            dataset=dataset,
            datapath=datapath,
            with_label=True,
            test=not train)

    def __getitem__(self, index):
        return self.emb[index], self.targets[index]

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


def get_dataset(dataset, datapath='./data', train=True, transform=None, download=True, precompute_arch=None):
    if precompute_arch:
        return PrecomputedEmbeddingDataset(
            dataset=dataset,
            arch=precompute_arch,
            datapath="data", # assumes embeddings are saved in the ./data folder
            train=train)
    
    load_obj = tds if dataset in ["CIFAR10","CIFAR100", "STL10"] else loaders
    if dataset == "STL10":
        split = 'train' if train else 'test'
        return getattr(load_obj, dataset)(root=datapath,
                        split=split,
                        download=download, transform=transform)
    elif "CIFAR" in dataset:
        return getattr(load_obj, dataset)(root=datapath,
                        train=train,
                        download=download, transform=transform)
    else:
        # imagenet subsets
        # TODO i dont know if val and val_structured are the same
        if "ILSVRC" in datapath and train is False:
            datapath = datapath.replace("train","val")
        return getattr(load_obj, dataset)(root=datapath,
                         transform=transform)


class EmbedNN(Dataset):
    def __init__(self,
                 knn_path,
                 transform,
                 k=10,
                 dataset="CIFAR100",
                 datapath='./data',
                 precompute_arch=None):
        super().__init__()
        self.transform = transform
        self.complete_neighbors = torch.load(knn_path)
        if k < 0:
            k = self.complete_neighbors.size(1)
        self.k = k
        self.neighbors = self.complete_neighbors[:, :k]
        self.datapath = './data' if 'IN' not in dataset else datapath

        self.dataset = get_dataset(
            dataset,
            datapath=datapath,
            transform=None,
            train=True,
            download=True,
            precompute_arch=precompute_arch)

    def get_transformed_imgs(self, idx, *idcs):
        img, label = self.dataset[idx]
        rest_imgs = (self.dataset[i][0] for i in idcs)
        return self.transform(img, *rest_imgs), label

    def __getitem__(self, idx):
        # KNN pair
        pair_idx = np.random.choice(self.neighbors[idx], 1)[0]

        return self.get_transformed_imgs(idx, pair_idx)

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


class TruePosNN(EmbedNN):

    def __init__(self, knn_path, *args, **kwargs):
        super().__init__(knn_path, *args, **kwargs)
        p = Path(knn_path).parent
        nn_p = p / 'hard_pos_nn.pt'
        if nn_p.is_file():
            self.complete_neighbors = torch.load(nn_p)
        else:
            emb = torch.load(p / 'embeddings.pt')
            emb /= emb.norm(dim=-1, keepdim=True)
            d = emb @ emb.T
            labels = torch.tensor(self.dataset.targets)
            same_label = labels.view(1, -1) == labels.view(-1, 1)
            # Find minimum number of images per class
            k_max = same_label.sum(dim=1).min()
            d.fill_diagonal_(-2)
            d[torch.logical_not(same_label)] = -torch.inf
            self.complete_neighbors = d.topk(k_max, dim=-1)[1]
            torch.save(self.complete_neighbors, nn_p)
        self.neighbors = self.complete_neighbors[:, :self.k]

loaders/embedNN.py
HHU-MMBS/TEMI-official-BMVC2023
"""
Author: Wouter Van Gansbeke, Simon Vandenhende
Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)
copied and adjusted from https://raw.githubusercontent.com/wvangansbeke/Unsupervised-Classification/master/data/imagenet.py
"""
import os
import torch
import torchvision.datasets as datasets
import torch.utils.data as data
from PIL import Image
from torchvision import transforms as tf
from glob import glob


class ImageNet(datasets.ImageFolder):
    def __init__(self, root='', split='train', transform=None):
        super(ImageNet, self).__init__(root=root,
                                         transform=None)
        self.transform = transform 
        self.split = split
        self.resize = tf.Resize((256,256))
    
    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, index):
        path, target = self.imgs[index]
        with open(path, 'rb') as f:
            img = Image.open(f).convert('RGB')
        im_size = img.size
        img = self.resize(img)

        if self.transform is not None:
            img = self.transform(img)

        return img, target

    def get_image(self, index):
        path, target = self.imgs[index]
        with open(path, 'rb') as f:
            img = Image.open(f).convert('RGB')
        img = self.resize(img) 
        return img


class ImageNetSubset(data.Dataset):
    def __init__(self, subset_file, root='/home/shared/DataSets/ILSVRC/Data/CLS-LOC/train', split='train', 
                    transform=None):
        super(ImageNetSubset, self).__init__()

        self.root = root
        self.transform = transform
        self.split = split

        # Read the subset of classes to include (sorted)
        with open(subset_file, 'r') as f:
            result = f.read().splitlines()
        subdirs, class_names = [], []
        for line in result:
            subdir, class_name = line.split(' ', 1)
            subdirs.append(subdir)
            class_names.append(class_name)

        # Gather the files (sorted)
        imgs = []
        for i, subdir in enumerate(subdirs):
            subdir_path = os.path.join(self.root, subdir)
            files = sorted(glob(os.path.join(self.root, subdir, '*.JPEG')))
            for f in files:
                imgs.append((f, i)) 
        self.imgs = imgs 
        self.classes = class_names
        self.targets = list(zip(*self.imgs))[-1]
	    # Resize
        self.resize = tf.Resize((256,256))

    def get_image(self, index):
        path, target = self.imgs[index]
        with open(path, 'rb') as f:
            img = Image.open(f).convert('RGB')
        img = self.resize(img) 
        return img

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

    def __getitem__(self, index):
        path, target = self.imgs[index]
        with open(path, 'rb') as f:
            img = Image.open(f).convert('RGB')
        im_size = img.size
        img = self.resize(img) 
        class_name = self.classes[target]

        if self.transform is not None:
            img = self.transform(img)

        #out = {'image': img, 'target': target, 'meta': {'im_size': im_size, 'index': index, 'class_name': class_name}}
        return img,target



def IN50(**kwargs):
    return ImageNetSubset(subset_file="loaders/imagenet_subsets/imagenet_50.txt",**kwargs)

def IN100(**kwargs):
    return ImageNetSubset(subset_file="loaders/imagenet_subsets/imagenet_100.txt",**kwargs)

def IN200(**kwargs):
    return ImageNetSubset(subset_file="loaders/imagenet_subsets/imagenet_200.txt",**kwargs)


def IN1K(**kwargs):
    return datasets.ImageFolder(**kwargs)
loaders/imagenet.py
HHU-MMBS/TEMI-official-BMVC2023
from .mi import *
from .loss_utils import *
from .multihead_losses import *

losses/__init__.py
HHU-MMBS/TEMI-official-BMVC2023
import torch

@torch.jit.script
def sim_weight(p1, p2, gamma: float = 1.):
    return (p1 * p2).pow(gamma).sum(dim=-1)

@torch.jit.script
def beta_mi(p1, p2, pk, beta: float = 1., clip_min: float = -torch.inf):
    beta_emi = (((p1 * p2)**beta) / pk).sum(dim=-1)
    beta_pmi = beta_emi.log().clamp(min=clip_min)
    return -beta_pmi

losses/loss_utils.py
HHU-MMBS/TEMI-official-BMVC2023
from abc import ABC, abstractmethod

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

__all__ = [
    "DINOLossMI",
    "WMIBase",
    "PMI",
    "SCAN",
]

import utils
from losses.loss_utils import sim_weight, beta_mi


class StudentTeacherLoss(nn.Module, ABC):

    def __init__(self, out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1):
        super().__init__()
        self.student_temp = student_temp
        self.batchsize = batchsize
        self.out_dim = out_dim
        # we apply a warm up for the teacher temperature because
        # a too high temperature makes the training instable at the beginning
        self.teacher_temp_schedule = np.concatenate((
            np.linspace(warmup_teacher_temp,
                        teacher_temp, warmup_teacher_temp_epochs),
            np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp
        ))

    def student_probs(self, x):
        return F.softmax(x / self.student_temp, dim=-1)

    def teacher_probs(self, x, epoch=0):
        return F.softmax(x / self.teacher_temp_schedule[epoch], dim=-1)

    @abstractmethod
    def forward(self, student_out, teacher_out, epoch):
        pass

    def update_ema(self, output, ema, momentum, use_momentum=True):
        """
        Update exponential moving averages for teacher output.
        """
        if isinstance(output, tuple):
            output = torch.cat(output)
        batch_center = torch.sum(output, dim=0, keepdim=True)
        utils.all_reduce(batch_center)
        batch_center = batch_center / (len(output) * utils.get_world_size())
        if use_momentum:
            return ema * momentum + batch_center * (1 - momentum)
        else:
            return batch_center


class DINOLossMI(StudentTeacherLoss):
    def __init__(self, out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1,
                 center_momentum=0.9, probs_momentum=0.9):
        super().__init__(
            out_dim=out_dim,
            batchsize=batchsize,
            warmup_teacher_temp=warmup_teacher_temp,
            teacher_temp=teacher_temp,
            warmup_teacher_temp_epochs=warmup_teacher_temp_epochs,
            nepochs=nepochs,
            student_temp=student_temp
        )
        self.center_momentum = center_momentum
        self.probs_momentum = probs_momentum
        self.register_buffer("probs_pos", torch.ones(1, out_dim)/out_dim)
        self.register_buffer("probs_neg", torch.ones(1, out_dim)/out_dim)

    def forward(self, student_out, teacher_out, epoch):
        ncrops_student = len(student_out)//self.batchsize
        ncrops_teacher = len(teacher_out)//self.batchsize

        temp = self.teacher_temp_schedule[epoch]
        student_probs = F.softmax(student_out / self.student_temp, dim=-1).chunk(ncrops_student)
        teacher_probs = F.softmax(teacher_out / temp, dim=-1).detach().chunk(ncrops_teacher)

        with torch.no_grad():
            self.probs_pos = self.update_ema(teacher_probs, self.probs_pos,
                                             self.probs_momentum,
                                             use_momentum=True)

        p_k = self.probs_pos

        total_loss = 0.0
        count = 0
        for w in range(ncrops_teacher):
            for v in range(ncrops_student):
                if v == w:
                    # we skip cases where student and teacher operate on the same view for positive examples
                    continue

                loss = - torch.log(torch.sum(teacher_probs[w] * student_probs[v] / p_k, dim=-1))
                total_loss += loss.mean()
                count += 1

        total_loss /= count
        return total_loss


class WMIBase(DINOLossMI):
    def __init__(self, out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1,
                 center_momentum=0.9, reg=False, alpha=0.5, beta=1,
                 weight_mi=True,
                 positive_pmi=False):
        super().__init__(out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                         warmup_teacher_temp_epochs, nepochs, student_temp=student_temp,
                         center_momentum=center_momentum)
        self.out_dim = out_dim
        self.alpha = alpha
        self.reg = reg
        self.weight_mi = weight_mi
        self.min_mi = 0 if positive_pmi else -torch.inf
        self.beta = beta

    def loss(self, p1, p2, pk, beta):
        return beta_mi(p1, p2, pk, beta=beta, clip_min=self.min_mi)

    def weight(self, p1, p2):
        raise not NotImplementedError("weight function not implemented")

    def forward(self, student_out, teacher_out, epoch=0):
        ncrops_student = len(student_out)//self.batchsize
        ncrops_teacher = len(teacher_out)//self.batchsize

        temp = self.teacher_temp_schedule[epoch]
        beta = self.beta
        student_out = (student_out / self.student_temp).chunk(ncrops_student)
        teacher_probs = F.softmax(teacher_out / temp, dim=-1).detach().chunk(ncrops_teacher)

        with torch.no_grad():
            self.probs_pos = self.update_ema(teacher_probs, self.probs_pos,
                                             self.probs_momentum,
                                             use_momentum=True)
            p_k = self.probs_pos

        # weight between teacher pairs
        if self.weight_mi:
            pairs = 0
            weight = 0
            for v1 in range(ncrops_teacher):
                for v2 in list(range(ncrops_teacher))[v1 + 1:]:
                    weight += self.weight(teacher_probs[v1], teacher_probs[v2])
                    pairs += 1
            weight = weight/pairs
        else:
            weight = 1

        # Weighted MI objective
        total_loss = 0.0
        count = 0
        for w in range(ncrops_teacher):
            for v in range(ncrops_student):
                if v == w:
                    # we skip cases where student and teacher operate on the same view for positive examples
                    continue
                loss = weight*self.loss(teacher_probs[w], F.softmax(student_out[v], dim=-1), p_k, beta)
                total_loss += loss.mean()
                count += 1

        # self-entropy regularization
        if self.reg:
            total_aux_loss = 0.0
            for s in range(ncrops_student):
                aux_loss = torch.sum( F.softmax(student_out[s], dim=-1) *  F.log_softmax(student_out[s], dim=-1), dim=-1)
                total_aux_loss += aux_loss.mean() # batch-wise mean
            total_loss = total_loss / count + self.alpha * (total_aux_loss / ncrops_student)
        else:
            total_loss = total_loss / count
        return total_loss


class WMI(WMIBase):

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

    def weight(self, p1, p2):
        return sim_weight(p1, p2)


class PMI(WMIBase):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs, weight_mi=False)


class SCAN(DINOLossMI):
    def __init__(self, out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1,
                 center_momentum=0.9, alpha=5):
        super().__init__(out_dim, batchsize, warmup_teacher_temp, teacher_temp,
                         warmup_teacher_temp_epochs, nepochs, student_temp=student_temp,
                         center_momentum=center_momentum)
        self.out_dim = out_dim
        self.alpha = alpha

    def forward(self, student_out, teacher_out, epoch):
        ncrops_student = len(student_out)// self.batchsize
        ncrops_teacher = len(teacher_out)// self.batchsize

        student_probs = self.student_probs(student_out).chunk(ncrops_student)
        with torch.no_grad():
            teacher_probs = self.teacher_probs(teacher_out, epoch=epoch).chunk(ncrops_teacher)

        # Weighted MI objective
        total_loss = 0.0
        count = 0
        for w in range(ncrops_teacher):
            for v in range(ncrops_student):
                if v == w:
                    # we skip cases where student and teacher operate on the same view for positive examples
                    continue
                loss = -(teacher_probs[w] * student_probs[v]).sum(dim=-1).log()
                total_loss += loss.mean()
                count += 1

        # self-entropy regularization
        total_aux_loss = 0.0
        for pkx in student_probs:
            pk = pkx.mean(dim=0)
            total_aux_loss += (pk * pk.log()).sum()
        return total_loss / count + self.alpha * ( total_aux_loss / ncrops_student)

losses/mi.py
HHU-MMBS/TEMI-official-BMVC2023
from abc import ABC, abstractmethod

import torch

from losses.mi import StudentTeacherLoss

__all__ = [
    'MultiHeadWMI',
    "is_multihead",
    "TEMI",
]

from losses import sim_weight, beta_mi


def multihead_loss(cls):
    """Marks a loss as multiclass compatible."""
    cls.IS_MULTIHEAD = True
    return cls


def is_multihead(cls):
    """Checks if a loss works with multiple heads"""
    return getattr(cls, "IS_MULTIHEAD", False)


@multihead_loss
class MultiHeadWeightProbBase(StudentTeacherLoss, ABC):

    def __init__(self, num_heads, probs_momentum=0.9, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.probs_momentum = probs_momentum
        self.num_heads = num_heads
        self.register_buffer("pk", 1/self.out_dim * torch.ones(self.num_heads, self.out_dim))

    @abstractmethod
    def weight(self, pt1, pt2, epoch, idx):
        pass

    @abstractmethod
    def loss(self, ps, pt, epoch, idx):
        pass

    @property
    def pos_probs(self):
        return self.pk[0]

    def forward(self, student_out, teacher_out, epoch):
        # student_out and teacher_out should be lists of length num_heads
        multi_loss = []
        weight = self.compute_weight(epoch, teacher_out)

        for idx, (s, t) in enumerate(zip(student_out, teacher_out)):
            ncrops_student = len(s) // self.batchsize
            ncrops_teacher = len(t) // self.batchsize

            student_probs = self.student_probs(s).chunk(ncrops_student)
            with torch.no_grad():
                teacher_probs = self.teacher_probs(t, epoch).detach().chunk(ncrops_teacher)
            tmp_loss = 0
            count = 0
            for w in range(ncrops_teacher):
                for v in range(ncrops_student):
                    if v == w:
                        # we skip cases where student and teacher operate on the same view for positive examples
                        continue
                    tmp_loss += (weight * self.loss(student_probs[v], teacher_probs[w], epoch, idx)).mean()
                    count += 1
            multi_loss.append(tmp_loss / count)

        return torch.stack(multi_loss)

    def compute_weight(self, epoch, teacher_out):
        weight_total = 0
        weight_count = 0
        with torch.no_grad():
            for idx, t in enumerate(teacher_out):
                ncrops_teacher = len(t) // self.batchsize

                teacher_probs = self.teacher_probs(t, epoch).detach().chunk(ncrops_teacher)
                self.pk[idx] = self.update_ema(teacher_probs, self.pk[idx], self.probs_momentum, use_momentum=True)

                for v1 in range(ncrops_teacher):
                    for v2 in list(range(ncrops_teacher))[v1 + 1:]:
                        weight_total += self.weight(teacher_probs[v1], teacher_probs[v2], epoch, idx)
                        weight_count += 1
            weight = weight_total / weight_count
        return weight


class MultiHeadWMI(MultiHeadWeightProbBase):

    def __init__(self, beta, *args, positive_pmi=False, **kwargs):
        super().__init__(*args, **kwargs)
        self.beta = beta
        self.min_mi = 0 if positive_pmi else -torch.inf

    def weight(self, pt1, pt2, epoch, idx):
        raise NotImplementedError("Weighting is not implemented for this loss.")

    def loss(self, ps, pt, epoch, idx):
        pk = self.pk[idx]
        return beta_mi(ps, pt, pk, beta=self.beta, clip_min=self.min_mi)


class TEMI(MultiHeadWMI):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def weight(self, pt1, pt2, epoch, idx):
        return sim_weight(pt1, pt2)

losses/multihead_losses.py
HHU-MMBS/TEMI-official-BMVC2023
import argparse
import inspect

import utils

# Load loss, loader and transforms
from model_builders.multi_head import MultiHeadClassifier
from utils import kv_pair


def get_args_parser():
    parser = argparse.ArgumentParser('MI clustering', add_help=False)

    # Model parameters
    parser.add_argument('--arch', default='clip_ViT-B/32', type=str,
                        help="""Name of architecture to train. For quick experiments with ViTs,
                        we recommend using vit_tiny or vit_small.""")
    parser.add_argument('--datapath', default='./data', type=str)
    parser.add_argument('--embed_dim', default=None, type=int)
    parser.add_argument('--patch_size', default=16, type=int, help="""Size in pixels
        of input square patches - default 16 (for 16x16 patches). Using smaller
        values leads to better performance but requires more memory. Applies only
        for ViTs (vit_tiny, vit_small and vit_base). If <16, we recommend disabling
        mixed precision training (--use_fp16 false) to avoid instabilities.""")
    parser.add_argument('--momentum_teacher', default=0.996, type=float, help="""Base EMA
        parameter for teacher update.
        The value is increased to max_momentum_teacher during training with cosine schedule.
        We recommend setting a higher value with small batches: for example use 0.9995 with batch size of 256.""")
    parser.add_argument('--max_momentum_teacher', default=1.0, type=float)
    parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag,
                        help="Whether to use batch normalizations in projection head (Default: False)")
    parser.add_argument('--head_dropout_prob', default=0.0, type=float,
                        help="Dropout probability in projection head (Default: 0.0)")

    # Temperature teacher parameters
    parser.add_argument('--warmup_teacher_temp', default=0.1, type=float,
                        help="""Initial value for the teacher temperature: 0.04 works well in most cases.
                        Try decreasing it if the training loss does not decrease.""")
    parser.add_argument('--teacher_temp', default=0.1, type=float, help="""Final value (after linear warmup)
        of the teacher temperature. For most experiments, anything above 0.07 is unstable. We recommend
        starting with the default value of 0.04 and increase this slightly if needed.""")
    parser.add_argument('--warmup_teacher_temp_epochs', default=30, type=int,
                        help='Number of warmup epochs for the teacher temperature (Default: 30).')

    # Training/Optimization parameters
    parser.add_argument('--use_fp16', type=utils.bool_flag, default=False, help="""Whether or not
        to use half precision for training. Improves training time and memory requirements,
        but can provoke instability and slight decay of performance. We recommend disabling
        mixed precision if the loss is unstable, if reducing the patch size or if training with bigger ViTs.""")
    parser.add_argument('--weight_decay', type=float, default=0.0001, help="""Initial value of the
        weight decay. With ViT, a smaller value at the beginning of training works well.""")
    parser.add_argument('--weight_decay_end', type=float, default=0.0001, help="""Final value of the
        weight decay. We use a cosine schedule for WD and using a larger decay by
        the end of training improves performance for ViTs.""")
    parser.add_argument('--clip_grad', type=float, default=0, help="""Maximal parameter
        gradient norm if using gradient clipping. Clipping with norm .3 ~ 1.0 can
        help optimization for larger ViT architectures. 0 for disabling.""")

    bs_group = parser.add_mutually_exclusive_group()
    bs_group.add_argument('--batch_size_per_gpu', default=64, type=int,
                          help='Per-GPU batch-size : number of distinct images loaded on one GPU.')
    bs_group.add_argument('--batch_size', default=None, type=int,
                          help='Total batch size')

    parser.add_argument('--epochs', default=100, type=int, help='Number of epochs of training.')
    parser.add_argument('--freeze_last_layer', default=1, type=int, help="""Number of epochs
        during which we keep the output layer fixed. Typically doing so during
        the first epoch helps training. Try increasing this value if the loss does not decrease.""")
    parser.add_argument("--lr", default=0.0005, type=float, help="""Learning rate at the end of
        linear warmup (highest LR used during training). The learning rate is linearly scaled
        with the batch size, and specified here for a reference batch size of 256.""")
    parser.add_argument("--warmup_epochs", default=10, type=int,
                        help="Number of epochs for the linear learning-rate warm up.")
    parser.add_argument('--min_lr', type=float, default=1e-6, help="""Target LR at the
        end of optimization. We use a cosine LR schedule with linear warmup.""")
    parser.add_argument('--optimizer', default='adamw', type=str,
                        choices=['adamw', 'sgd', 'lars'],
                        help="""Type of optimizer. We recommend using adamw with ViTs.""")
    parser.add_argument('--drop_path_rate', type=float, default=0.1, help="stochastic depth rate")

    # Multi-crop and aug
    parser.add_argument('--vit_image_size', type=int, default=224, help="""image size that enters vit; 
        must match with patch_size: num_patches = (vit_image_size/patch_size)**2""")
    parser.add_argument('--image_size', type=int, default=32, help="""image size of in-distibution data. 
        negative samples are first resized to image_size and then inflated to vit_image_size. This
        ensures that aux samples have same resolution as in-dist samples""")
    parser.add_argument('--aug_image_size',type=int, default=None,
                       help='Image size for data augmentation. If None, use vit_image_size')
    from augs.augs import IMAGE_AUGMENTATIONS, EMBED_AUGMENTATIONS
    parser.add_argument('--image_aug', choices=IMAGE_AUGMENTATIONS.keys(), default='randaug',
                        help='Augmentation for images')
    parser.add_argument('--embed_aug', choices=EMBED_AUGMENTATIONS.keys(), default='none',
                        help='Augmentations for precomputed embeddings. Only used when --precomputed flag is given')
    parser.add_argument('--aug_args', type=kv_pair, nargs='*', default={})
    parser.add_argument('--num_augs', type=int, default=1)

    parser.add_argument('--pretrained_weights', default='', type=str)
    parser.add_argument('--knn_path', default=None, type=str)
    parser.add_argument('--dataset', default='CIFAR100',
                        choices=['CIFAR100', 'CIFAR10', "STL10", "CIFAR20", "IN1K", "IN50", 'IN100', "IN200"],
                        type=str)
    parser.add_argument('--knn', type=int, default=50, help='Number of nearest neighbors to use')

    parser.add_argument('--output_dir', default=None, type=str, help='Path to save logs and checkpoints.')
    parser.add_argument('--saveckp_freq', default=20, type=int, help='Save checkpoint every x epochs.')
    parser.add_argument('--seed', default=0, type=int, help='Random seed.')
    parser.add_argument('--num_workers', default=6, type=int, help='Number of data loading workers per GPU.')
    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")

    parser.add_argument('--out_dim', default=1000, type=int, help="""Dimensionality of
        the DINO head output. For complex and large datasets large values (like 65k) work well.""")
    parser.add_argument('--norm_last_layer', default=False, type=utils.bool_flag,
                        help="""Whether or not to weight normalize the last layer of the DINO head.
        Not normalizing leads to better performance but can make the training unstable.
        In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""")

    parser.add_argument('--nlayers', default=2, type=int, help='Head layers')
    parser.add_argument('--hidden_dim', default=512, type=int, help="Head's hidden dim")
    parser.add_argument('--bottleneck_dim', default=256, type=int, help="Head's bottleneck dim")
    parser.add_argument('--l2_norm', default=False, action='store_true', help="Whether to apply L2 norm after backbone")
    parser.add_argument('--embed_norm', default=False, action='store_true',
                        help="Whether to normalize embeddings using precomputed mean and std")

    add_from_signature(parser, MultiHeadClassifier, "MultiHeadClassifier")

    parser.add_argument('--disable_ddp', default=False, action='store_true', help="Don't use DDP")

    precomputed_group = parser.add_mutually_exclusive_group()
    precomputed_group.add_argument('--train_backbone', default=False, action='store_true',
                        help="Don't share backbones between teacher and student")
    precomputed_group.add_argument('--precomputed', action='store_true', help="Use precomputed embeddings", default=False)

    parser.add_argument("--loss", default="TEMI", help="The name of one of the classes in losses",
                        choices=['WMI', 'PMI', 'TEMI', 'SCAN'])
    parser.add_argument("--loss-args", type=kv_pair, nargs="*", default={},
                        help="Extra arguments for the loss class")
    parser.add_argument("--loader", default='EmbedNN', help="The name of one of the classes in loaders")
    parser.add_argument("--loader-args", type=kv_pair, nargs="*", default={},
                        help="Extra arguments for the loader class")
    parser.add_argument('--new-run', action='store_true', help="Create a new directory for this run", default=False)

    return parser


def add_from_signature(parser, function, name=None):
    """Add arguments from a function signature to an existing parser."""
    if name is None:
        name = function.__name__
    group = parser.add_argument_group(name)
    signature = inspect.signature(function)
    for name, param in signature.parameters.items():
        default = param.default
        if param.kind == param.VAR_KEYWORD or default is param.empty:
            continue
        try:
            group.add_argument("--" + name, default=param.default, type=type(default))
        except argparse.ArgumentError:
            # Ignore arguments that are already added
            pass


def set_default_args(config):
    if hasattr(config, '__dict__'):
        config = vars(config)
    args = get_args_parser().parse_args([])
    args.__dict__.update(config)
    return args


def process_args(args):
    args.loss_args = dict(args.loss_args)
    args.loader_args = dict(args.loader_args)
    args.aug_args = dict(args.aug_args)
    if not args.knn_path:
        import gen_embeds
        args.knn_path = gen_embeds.get_outpath(args.arch, args.dataset) / 'knn.pt'
    if args.batch_size is not None:
        args.batch_size_per_gpu = None
    return args

main_args.py
HHU-MMBS/TEMI-official-BMVC2023
from .model_builders import *
from .multi_head import MultiHeadClassifier

model_builders/__init__.py
HHU-MMBS/TEMI-official-BMVC2023
import fnmatch
import inspect
import json
import sys
from argparse import Namespace
from pathlib import Path
import os
from typing import Optional, Union

import requests
import clip
import timm
import torch
from torchvision import models as torchvision_models, transforms
from tqdm import tqdm

import utils
from .multi_head import MultiHeadClassifier

from timm.models.helpers import load_state_dict

_AVAILABLE_MODELS = (
    "dino_resnet50",
    "dino_vits16",
    "dino_vitb16",
    "timm_resnet50",
    "timm_vit_small_patch16_224",
    "timm_vit_base_patch16_224",
    "timm_vit_large_patch16_224",
    "convnext_small",
    "convnext_base",
    "convnext_large",
    "msn_vit_small",
    "msn_vit_base",
    "msn_vit_large",
    "mocov3_vit_small",
    "mocov3_vit_base",
    "clip_ViT-B/16",
    "clip_ViT-L/14",
    "clip_RN50",
    "mae_vit_base",
    "mae_vit_large",
    "mae_vit_huge",
)


def available_models(pattern=None):
    if pattern is None:
        return _AVAILABLE_MODELS
    return tuple(fnmatch.filter(_AVAILABLE_MODELS, pattern))


def _load_checkpoint(model, checkpoint_path, use_ema=False, strict=True):
    if os.path.splitext(checkpoint_path)[-1].lower() in ('.npz', '.npy'):
        # numpy checkpoint, try to load via model specific load_pretrained fn
        if hasattr(model, 'load_pretrained'):
            model.load_pretrained(checkpoint_path)
        else:
            raise NotImplementedError('Model cannot load numpy checkpoint')
        return
    state_dict = load_state_dict(checkpoint_path, use_ema)
    msg = model.load_state_dict(state_dict, strict=strict)
    print(msg)


def _download(url: str, filename: Path):
    """from https://stackoverflow.com/a/37573701"""
    response = requests.get(url, stream=True)
    total_size_in_bytes = int(response.headers.get('content-length', 0))
    block_size = 1024  # 1 Kibibyte
    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
    with open(filename, 'wb') as file:
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()
    if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
        raise Exception(f"Could not download from {url}")


_dict_models_urls = {
    "msn": {
        'vit_small_patch16_224': 'https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar',
        'vit_base_patch16_224': 'https://dl.fbaipublicfiles.com/msn/vitb16_600ep.pth.tar',
        'vit_large_patch16_224': 'https://dl.fbaipublicfiles.com/msn/vitl16_600ep.pth.tar',
        "key": 'target_encoder'
    },
    "mae": {
        'vit_base_patch16_224': 'https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth',
        'vit_large_patch16_224': 'https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pth',
        'vit_huge_patch14_224_in21k': 'https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth',
        "key": 'model'
    },

    "mocov3": {
        'vit_small_patch16_224': 'https://dl.fbaipublicfiles.com/moco-v3/vit-s-300ep/vit-s-300ep.pth.tar',
        'vit_base_patch16_224': 'https://dl.fbaipublicfiles.com/moco-v3/vit-b-300ep/vit-b-300ep.pth.tar',
        "key":  "state_dict"
    },
}

_dict_timm_names = {
        "vit_huge": 'vit_huge_patch14_224_in21k',
        "vit_large": 'vit_large_patch16_224',
        "vit_base": 'vit_base_patch16_224',
        "vit_small": 'vit_small_patch16_224',
        "vit_tiny": 'vit_tiny_patch16_224',
        "resnet50": 'resnet50',
        }


def _get_checkpoint_path(model_name: str, timm_base=True):
    if timm_base:
        name = _get_timm_name(model_name)
    prefix = model_name.split("_")[0]
    model_url = _dict_models_urls[prefix][name]
    print(f"Loading {model_url}")
    root = Path('~/.cache/torch/checkpoints').expanduser()
    root.mkdir(parents=True, exist_ok=True)
    path = root / f'{model_name}.pth'
    if not path.is_file():
        print('Downloading checkpoint...')
        _download(model_url, path)
        d = torch.load(path, map_location='cpu')
        ckpt_key_name = _dict_models_urls[prefix]["key"]
        if ckpt_key_name in d.keys():
            state_dict = d[ckpt_key_name]
            state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
            state_dict = {k.replace("momentum_encoder.", ""): v for k, v in state_dict.items()} # for mocov3
            torch.save(state_dict, path)
        else:
            raise KeyError(f"{ckpt_key_name} not found. Only {d.keys()} are available.")
    return path


def _get_timm_name(model_name: str):
    prefix = model_name.split("_")[0]
    # remove prefix
    model_name = model_name.replace("".join([prefix,"_"]),"")
    if model_name in _dict_timm_names.keys():
        return _dict_timm_names[model_name]
    else:
        raise ValueError(f"Model {model_name} not found")


def build_arch(model_name: str):
    timm_name = _get_timm_name(model_name)
    model = timm.create_model(
        timm_name,
        in_chans=3,
        num_classes=0,
        pretrained=False)
    _load_checkpoint(model, _get_checkpoint_path(model_name), strict=False)
    return model


def load_model(config, head=True, split_preprocess=False):
    """
    config/args file
    head=False returns just the backbone for baseline evaluation
    split_preprocess=True returns resizing etc. and normalization/ToTensor as separate transforms
    """
    from main_args import set_default_args
    config = set_default_args(config)

    preprocess = None

    if config.precomputed:
        backbone = config.arch
    elif "timm" in config.arch:  # timm models
        arch = config.arch.replace("timm_", "")
        arch = arch.replace("timm-", "")
        backbone = timm.create_model(arch, pretrained=True, in_chans=3, num_classes=0)
    elif "swag" in config.arch:
        arch = config.arch.replace("swag_", "")
        backbone = torch.hub.load("facebookresearch/swag", model=arch)
        backbone.head = None
    elif "dino" in config.arch:  # dino pretrained models on IN
        # dino_vitb16, dino_vits16
        arch = config.arch.replace("-", "_")
        backbone = torch.hub.load('facebookresearch/dino:main', arch)
    elif "clip" in config.arch: # load clip vit models from openai
        arch = config.arch.replace("clip_", "")
        arch = arch.replace("clip-", "")
        assert arch in clip.available_models()
        clip_model, preprocess = clip.load(arch)
        backbone = clip_model.visual
    elif "mae" in config.arch or "msn" in config.arch or "mocov3" in config.arch:
        backbone = build_arch(config.arch)
    elif "convnext" in config.arch:
        backbone = getattr(torchvision_models, config.arch)(pretrained=True)
        backbone.classifier = torch.nn.Flatten(start_dim=1, end_dim=-1)

    elif config.arch in torchvision_models.__dict__.keys(): # torchvision models
        backbone = torchvision_models.__dict__[config.arch](num_classes=0)
    else:
        print(f"Architecture {config.arch} non supported")
        sys.exit(1)
    if not config.precomputed:
        print(f"Backbone {config.arch} loaded.")
    else:
        print("No backbone loaded, using precomputed embeddings from", config.arch)

    if preprocess is None:
        # imagenet means/stds
        normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        resize = transforms.Resize(config.vit_image_size, transforms.InterpolationMode.BICUBIC)
        if config.dataset=="IN1K":
            preprocess = transforms.Compose([
                transforms.Resize(int(256 * config.vit_image_size / 224)),
                transforms.CenterCrop(config.vit_image_size)
            ])
        else:
            preprocess = resize
        if not split_preprocess:
            preprocess = transforms.Compose([
                preprocess,
                transforms.ToTensor(),
                normalize
            ])
    elif split_preprocess:
        preprocess_aux = preprocess.transforms[:-2]
        normalize_aux = preprocess.transforms[-2:]
        preprocess = transforms.Compose(preprocess_aux)
        normalize = transforms.Compose(normalize_aux)

    if head:
        if getattr(config, "embed_dim", None) is None:
            config.embed_dim = utils.embed_dim(config, backbone)
        # Just get everything via reflection
        mmc_params = inspect.signature(MultiHeadClassifier).parameters
        mmc_args = {k: v for k, v in config.__dict__.items() if k in mmc_params}
        model = MultiHeadClassifier(backbone, **mmc_args)

        if config.embed_norm:
            model.set_mean_std(*load_embed_stats(config, test=False))
        print("Head loaded.")
    else:
        model = backbone.float()

    if split_preprocess:
        return model, preprocess, normalize 
    print(preprocess)   
    return model, preprocess


def _build_from_config(
        precomputed: bool,
        config: Optional[Union[str, Path, Namespace]] = None,
        ckpt_path: Optional[Union[str, Path]] = None):
    if isinstance(config, str) or isinstance(config, Path):
        p = Path(config)
        with open(p, "r") as f:
            config = json.load(f)
        config = Namespace(**config)
    if config is None:
        config = Namespace()
    config.num_heads = 1
    config.precomputed = precomputed
    if ckpt_path is not None:
        # Don't reload norms
        config.embed_norm = False

    d = None
    if ckpt_path is not None:
        d = torch.load(ckpt_path, map_location="cpu")
        if 'teacher' in d:
            d = d['teacher']
        if 'head.best_head_idx' in d:
            best_head_idx = d['head.best_head_idx']
            d2 = {k: v for k, v in d.items() if k in ('embed_mean', 'embed_std')}
            d2['head.best_head_idx'] = torch.tensor(0)
            for k, v in d.items():
                if k.startswith(f'head.heads.{best_head_idx}.'):
                    k = 'head.heads.0.' + k[len(f'head.heads.{best_head_idx}.'):]
                    d2[k] = v
            d = d2
        else:
            d['head.best_head_idx'] = torch.tensor(0)
        config.embed_dim = d['head.heads.0.mlp.0.weight'].size(1)

    model, _ = load_model(config, head=True)
    model.eval()
    if d is not None:
        model.load_state_dict(d, strict=False)

    return model


def build_head_from_config(
        config: Optional[Union[str, Path, Namespace]] = None,
        ckpt_path: Optional[Union[str, Path]] = None):
    """
    config: Either path to hp.json or config namespace
    ckpt_path: Path to checkpoint
    """
    return _build_from_config(True, config, ckpt_path)


def build_model_from_config(
        config: Optional[Union[str, Path, Namespace]] = None,
        ckpt_path: Optional[Union[str, Path]] = None):
    """
    config: Either path to hp.json or config namespace
    ckpt_path: Path to checkpoint
    """
    return _build_from_config(False, config, ckpt_path)


def load_embeds(config=None,
                arch=None,
                dataset=None,
                test=False,
                norm=False,
                datapath='data',
                with_label=False):
    p, test_str = _embedding_path(arch, config, datapath, dataset, test)
    emb = torch.load(p / f'embeddings{test_str}.pt', map_location='cpu')
    if norm:
        emb /= emb.norm(dim=-1, keepdim=True)
    if not with_label:
        return emb
    label = torch.load(p / f'label{test_str}.pt', map_location='cpu')
    return emb, label


def _embedding_path(arch, config, datapath, dataset, test):
    assert bool(config) ^ bool(arch and dataset)
    if config:
        arch = config.arch
        dataset = config.dataset
    import gen_embeds
    test_str = '-test' if test else ''
    p = gen_embeds.get_outpath(arch, dataset, datapath)
    return p, test_str


def load_embed_stats(
        config=None,
        arch=None,
        dset=None,
        test=False,
        datapath='data'):
    p, test_str = _embedding_path(arch, config, datapath, dset, test)
    mean = torch.load(p / f'mean{test_str}.pt', map_location='cpu')
    std = torch.load(p / f'std{test_str}.pt', map_location='cpu')
    return mean, std


model_builders/model_builders.py
HHU-MMBS/TEMI-official-BMVC2023
try:
    # for python version >= 3.8
    from functools import cached_property
except ImportError:
    from functools import lru_cache

    def cached_property(func):
        return property(lru_cache()(func))

import torch
import torch.nn as nn

from utils import backbone_dtype
from utils import trunc_normal_


class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim,
                 use_bn=False, dropout_p=0.0,
                 final_gelu=False, norm_last_layer=True,
                 nlayers=3, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        nlayers = max(nlayers, 1)
        if nlayers == 1:
            self.mlp = nn.Linear(in_dim, bottleneck_dim)
            if final_gelu:
                self.mlp = nn.Sequential(self.mlp, nn.GELU())
        else:
            layers = [nn.Linear(in_dim, hidden_dim)]
            if use_bn:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            if dropout_p > 0:
                layers.append(nn.Dropout(dropout_p))
            for _ in range(nlayers - 2):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                if use_bn:
                    layers.append(nn.BatchNorm1d(hidden_dim))
                layers.append(nn.GELU())
                if dropout_p > 0:
                    layers.append(nn.Dropout(dropout_p))
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
            if final_gelu:
                layers.append(nn.GELU())
            self.mlp = nn.Sequential(*layers)
        self.apply(self._init_weights)
        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
        self.last_layer.weight_g.data.fill_(1)
        if norm_last_layer:
            self.last_layer.weight_g.requires_grad = False

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.mlp(x)
        x = nn.functional.normalize(x, dim=-1, p=2)
        x = self.last_layer(x)
        return x

class MultiHead(nn.Module):

    def __init__(self, dino_head_args, num_heads=1):
        super().__init__()
        self.num_heads = num_heads

        if self.num_heads < 1:
            raise ValueError("Number of heads must be at least 1.")
            # List of all heads
        self.heads = nn.ModuleList([DINOHead(**dino_head_args) for _ in range(self.num_heads)])
        # best head is arbitrary at beginning
        self.register_buffer("best_head_idx", torch.tensor(0))

    @property
    def best_head(self):
        return self.heads[self.best_head_idx]

    def set_losses(self, losses):
        """losses should be (num_heads,) tensor"""
        if self.num_heads == 1:
            return
        if len(losses) != self.num_heads:
            raise ValueError("Number of losses does not match number of heads.")
        self.best_head_idx = torch.argmin(losses)

    def forward(self, x):
        if not self.training or self.num_heads == 1:
            return self.best_head(x)
        return [head(x) for head in self.heads]


class MultiHeadClassifier(nn.Module):
    """Multiple (parallel) heads on top of a backbone."""
    def __init__(self,
                 backbone,
                 embed_dim=512,
                 out_dim=4096,
                 use_bn_in_head=False,
                 head_dropout_prob=0.0,
                 head_final_gelu=False,
                 norm_last_layer=False,
                 req_grad=False,
                 l2_norm=False,
                 nlayers=3,
                 hidden_dim=1024,
                 layer_norm_only=False,
                 bottleneck_dim=256,
                 num_heads=1):
        super().__init__()

        self.register_buffer("embed_mean", torch.zeros(embed_dim))
        self.register_buffer("embed_std", torch.ones(embed_dim))

        precomputed = not isinstance(backbone, nn.Module)
        if layer_norm_only:
            for name, param in backbone.named_parameters():
                if "ln_" in name:
                    #print(name , param.shape)
                    param.requires_grad = True
                else:
                    param.requires_grad = False
                req_grad = True
        elif not precomputed:
            for p in backbone.parameters():
                p.requires_grad = req_grad

        self.backbone = backbone
        if precomputed:
            # self.backbone is just name of backbone that computed embeddings
            assert isinstance(self.backbone, str)
            self.backbone_embed = self.identity_backbone
        else:
            self.dtype = backbone_dtype(self.backbone)
        if not req_grad:
            self.backbone_embed = torch.no_grad()(self.backbone_embed)
        self.l2_norm = l2_norm
        self.embed_dim = embed_dim
        self.req_grad = req_grad

        head_args = dict(
            in_dim=embed_dim,
            out_dim=out_dim,
            use_bn=use_bn_in_head,
            final_gelu=head_final_gelu,
            dropout_p=head_dropout_prob,
            nlayers=nlayers,
            hidden_dim=hidden_dim,
            norm_last_layer=norm_last_layer,
            bottleneck_dim=bottleneck_dim
        )
        self.head = MultiHead(head_args, num_heads)

    def set_mean_std(self, mean, std):
        self.embed_mean.data = mean.clone()
        self.embed_std.data = std.clone()

    def identity_backbone(self, x):
        if isinstance(x, list):
            x = torch.cat(x)
        return (x - self.embed_mean) / self.embed_std

    def backbone_embed(self, x):
        # convert to list
        if not isinstance(x, list):
            x = [x]
        shapes_sorted, sort_idx = torch.sort(torch.Tensor([inp.shape[-1] for inp in x]))
        idx_crops = torch.cumsum(torch.unique_consecutive(shapes_sorted, return_counts=True)[1], 0)
        start_idx = 0
        output = torch.empty((len(x), len(x[0]), self.embed_dim), device=x[0].device)
        for end_idx in idx_crops:
            batch_idx = sort_idx[start_idx:end_idx]  # The indices of tensors of this shape
            _in_batched = torch.cat([x[i] for i in batch_idx])  # Batch them together
            _out = self.backbone(_in_batched.type(self.dtype)).float()

            # accumulate outputs
            _out = torch.stack(_out.chunk(len(batch_idx)))
            output.index_copy_(0, batch_idx.cuda(), _out)
            start_idx = end_idx
        output = torch.cat(torch.unbind(output))
        return (output - self.embed_mean) / self.embed_std

    def apply_head(self, embedded):
        if self.l2_norm:
            embedded /= embedded.norm(dim=-1, keepdim=True)
        # Run the head forward on the concatenated features.
        if not self.req_grad:
            embedded = embedded.detach()
        return self.head(embedded)

    def forward(self, x):
        return self.apply_head(self.backbone_embed(x))

model_builders/multi_head.py
HHU-MMBS/TEMI-official-BMVC2023
import os
import sys
import datetime
import time

import math
import json
from collections import OrderedDict
from pathlib import Path
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils.tensorboard import SummaryWriter

import utils
from augs.augs import IMAGE_AUGMENTATIONS, EMBED_AUGMENTATIONS, AugWrapper
import loaders

from torchvision import models as torchvision_models

import losses
from main_args import get_args_parser, process_args
from model_builders import load_model


class TeacherStudentCombo(nn.Module):

    def __init__(self, student, teacher, args):
        super().__init__()
        # synchronize batch norms (if any)
        if utils.has_batchnorms(student) and not args.disable_ddp:
            student = nn.SyncBatchNorm.convert_sync_batchnorm(student)
            teacher = nn.SyncBatchNorm.convert_sync_batchnorm(teacher)
        # teacher and student start with the same weights
        teacher.load_state_dict(student.state_dict())
        # Hacky
        if not args.train_backbone:
            student.backbone = teacher.backbone
        elif not args.req_grad:
            print('WARNING: args.train_backbone=True, but args.req_grad=False. '
                  'This is probably not what you want.')
        # there is no backpropagation through the teacher, so no need for gradients
        for p in teacher.parameters():
            p.requires_grad = False
        print(f"Student and Teacher are built: they are both {args.arch} network.")

        self.args = args
        self.student = student
        self.teacher = teacher

    def forward(self, images):
        if self.args.train_backbone:
            return self.teacher(images), self.student(images)
        embed = self.teacher.backbone_embed(images)
        return self.teacher.apply_head(embed), self.student.apply_head(embed)

    @property
    def module(self):
        return self

    def student_dict(self):
        if self.args.train_backbone:
            return self.student.state_dict()
        return OrderedDict([(k, v) for k, v in self.student.state_dict().items() if "backbone" not in k])

    @property
    def trainable_student(self):
        if self.args.train_backbone:
            return self.student
        return self.student.head

    def teacher_dict(self):
        if self.args.train_backbone:
            return self.teacher.state_dict()
        return OrderedDict([(k, v) for k, v in self.teacher.state_dict().items() if "backbone" not in k])

    @property
    def trainable_teacher(self):
        if self.args.train_backbone:
            return self.teacher
        return self.teacher.head


def train_dino(args, writer):
    if not args.disable_ddp:
        utils.init_distributed_mode(args)
    if args.batch_size is not None:
        args.batch_size_per_gpu = args.batch_size // utils.get_world_size()
    utils.fix_random_seeds(args.seed)
    print("git:\n  {}\n".format(utils.get_sha()))
    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
    cudnn.benchmark = True

    student, _, normalize = load_model(args, split_preprocess=True)
    teacher, _ = load_model(args)

    if not args.precomputed:
        aug = IMAGE_AUGMENTATIONS[args.image_aug](num_augs=args.num_augs, **args.aug_args)
        transform = AugWrapper(
            vit_image_size=args.vit_image_size,
            aug_image_size=args.aug_image_size,
            global_augs=aug,
            normalize=normalize,
            image_size=args.image_size
        )
    else:
        aug = EMBED_AUGMENTATIONS[args.embed_aug](num_augs=args.num_augs, **args.aug_args)
        transform = AugWrapper(
            global_augs=aug
        )

    dataset = getattr(loaders, args.loader)(
        knn_path=args.knn_path,
        datapath=args.datapath,
        k=args.knn,
        transform=transform, dataset=args.dataset,
        precompute_arch=args.arch if args.precomputed else None,
        **args.loader_args)

    if not args.disable_ddp:
        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
    else:
        sampler = None
    data_loader = torch.utils.data.DataLoader(
        dataset,
        shuffle=(sampler is None),
        sampler=sampler,
        batch_size=args.batch_size_per_gpu,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=True,
    )
    print(f"In-distribution Data loaded: there are {len(dataset)} images.")
    print("len dataloader", len(data_loader))

    student_teacher_model = TeacherStudentCombo(teacher=teacher, student=student, args=args)
    # move networks to gpu
    student_teacher_model = student_teacher_model.cuda()
    if not args.disable_ddp:
        student_teacher_model = nn.parallel.DistributedDataParallel(student_teacher_model, device_ids=[args.gpu])


    # ============ preparing loss ... ============
    loss_class = getattr(losses, args.loss)
    dino_loss_args = dict(
            out_dim=args.out_dim,
            batchsize=args.batch_size_per_gpu,
            warmup_teacher_temp=args.warmup_teacher_temp,
            teacher_temp=args.teacher_temp,
            warmup_teacher_temp_epochs=args.warmup_teacher_temp_epochs,
            nepochs=args.epochs,
            **args.loss_args)
    if losses.is_multihead(loss_class):
        dino_loss_args.update(num_heads=args.num_heads)
        dino_loss = loss_class(**dino_loss_args).cuda()
    elif args.num_heads == 1:
        dino_loss = loss_class(**dino_loss_args).cuda()
    else:
        dino_loss = nn.ModuleList([loss_class(**dino_loss_args) for _ in range(args.num_heads)]).cuda()

    # ============ preparing optimizer ... ============
    params_groups = utils.get_params_groups(student_teacher_model.module.trainable_student)
    if args.optimizer == "adamw":
        optimizer = torch.optim.AdamW(params_groups)  # to use with ViTs
    elif args.optimizer == "sgd":
        optimizer = torch.optim.SGD(params_groups, lr=0, momentum=0.9)  # lr is set by scheduler
    elif args.optimizer == "lars":
        optimizer = utils.LARS(params_groups)  # to use with convnet and large batches
    else:
        raise ValueError("Unknown optimizer: {}".format(args.optimizer))
    # for mixed precision training
    fp16_scaler = None
    if args.use_fp16:
        fp16_scaler = torch.cuda.amp.GradScaler()

    # ============ init schedulers ... ============
    bs_factor = (args.batch_size_per_gpu * utils.get_world_size()) / 256.
    lr_schedule = utils.cosine_scheduler(
        args.lr * bs_factor,  # linear scaling rule
        args.min_lr * bs_factor,
        args.epochs, len(data_loader),
        warmup_epochs=args.warmup_epochs,
    )
    wd_schedule = utils.cosine_scheduler(
        args.weight_decay,
        args.weight_decay_end,
        args.epochs, len(data_loader),
    )
    # momentum parameter is increased to 1. during training with a cosine schedule
    momentum_schedule = utils.cosine_scheduler(args.momentum_teacher, args.max_momentum_teacher,
                                               args.epochs, len(data_loader))
    print(f"Loss, optimizer and schedulers ready.")

    # ============ optionally resume training ... ============
    to_restore = {"epoch": 0}
    utils.restart_from_checkpoint(
        os.path.join(args.output_dir, "checkpoint.pth"),
        run_variables=to_restore,
        student=student_teacher_model.module.student,
        teacher=student_teacher_model.module.teacher,
        optimizer=optimizer,
        fp16_scaler=fp16_scaler,
        dino_loss=dino_loss,
    )
    start_epoch = to_restore["epoch"]

    start_time = time.time()
    print("Starting DINO training !")
    for epoch in range(start_epoch, args.epochs):
        if not args.disable_ddp:
            data_loader.sampler.set_epoch(epoch)
        # ============ training one epoch of DINO ... ============
        train_stats = train_one_epoch(student_teacher_model, dino_loss,
            data_loader, optimizer, lr_schedule, wd_schedule, momentum_schedule,
            epoch, fp16_scaler, args, writer)

        # ============ writing logs ... ============
        save_dict = {
            'student': student_teacher_model.module.student_dict(),
            'teacher': student_teacher_model.module.teacher_dict(),
            'optimizer': optimizer.state_dict(),
            'epoch': epoch + 1,
            'args': args,
            'dino_loss': dino_loss.state_dict(),
        }
        if fp16_scaler is not None:
            save_dict['fp16_scaler'] = fp16_scaler.state_dict()
        utils.save_on_master(save_dict, os.path.join(args.output_dir, 'checkpoint.pth'))
        if args.saveckp_freq and epoch % args.saveckp_freq == 0:
            utils.save_on_master(save_dict, os.path.join(args.output_dir, f'checkpoint{epoch:04}.pth'))
        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                     'epoch': epoch}
        if utils.is_main_process():
            with (Path(args.output_dir) / "log.txt").open("a") as f:
                f.write(json.dumps(log_stats) + "\n")

        try:       
            torch.set_printoptions(profile="full")
            if epoch % 10 == 0:
                d_loss = dino_loss[0] if hasattr(dino_loss, "__getitem__") else dino_loss
                print("highest probs:", torch.topk(d_loss.probs_pos * 100, 50)[0])
                print("lowest probs:", torch.topk(d_loss.probs_pos * 100, 50, largest=False)[0])
            torch.set_printoptions(profile="default")
        except:
            print(" ")

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))


def train_one_epoch(student_teacher_model, dino_loss, data_loader,
                    optimizer, lr_schedule, wd_schedule, momentum_schedule,epoch,
                    fp16_scaler, args, writer):
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Epoch: [{}/{}]'.format(epoch, args.epochs)
    for it, data in enumerate(metric_logger.log_every(data_loader, 10, header)):
        images, _ = data
        
        # update weight decay and learning rate according to their schedule
        it = len(data_loader) * epoch + it  # global training iteration        
        for i, param_group in enumerate(optimizer.param_groups):
            param_group["lr"] = lr_schedule[it]
            if i == 0:  # only the first group is regularized
                param_group["weight_decay"] = wd_schedule[it]
        # move images to gpu
        images = [im.cuda(non_blocking=True) for im in images]
        
        # teacher and student forward passes + compute dino loss
        with torch.cuda.amp.autocast(fp16_scaler is not None):
            teacher_out, student_out = student_teacher_model(images)
            if losses.is_multihead(dino_loss) or args.num_heads == 1:
                head_losses = dino_loss(student_out, teacher_out, epoch=epoch)
            else:
                head_losses = torch.stack([d(s, t, epoch=epoch) for d, s, t in zip(dino_loss, student_out, teacher_out)])
            loss = head_losses.mean()

        if not math.isfinite(loss.item()):
            print("Loss is {}, stopping training".format(loss.item()), flush=True)
            sys.exit(1)

        # student update
        optimizer.zero_grad()
        param_norms = None
        if fp16_scaler is None:
            loss.backward()
            if args.clip_grad:
                param_norms = utils.clip_gradients(student_teacher_model, args.clip_grad)
            utils.cancel_gradients_last_layer(epoch, student_teacher_model,
                                              args.freeze_last_layer)
            optimizer.step()
        else:
            fp16_scaler.scale(loss).backward()
            if args.clip_grad:
                fp16_scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
                param_norms = utils.clip_gradients(student_teacher_model, args.clip_grad)
            utils.cancel_gradients_last_layer(epoch, student_teacher_model,
                                              args.freeze_last_layer)
            fp16_scaler.step(optimizer)
            fp16_scaler.update()

        # EMA update for the teacher
        with torch.no_grad():
            m = momentum_schedule[it]  # momentum parameter
            s_head_params = student_teacher_model.module.trainable_student.parameters()
            t_head_params = student_teacher_model.module.trainable_teacher.parameters()
            for param_q, param_k in zip(s_head_params, t_head_params):
                param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)

        # logging
        torch.cuda.synchronize()
        metric_logger.update(loss=loss.item())
        metric_logger.update_raw(head_losses=head_losses)
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
        metric_logger.update(wd=optimizer.param_groups[0]["weight_decay"])
        if utils.is_main_process():
            writer.add_scalar("Train loss step", loss, it)

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()

    if utils.is_main_process() and args.num_heads > 1:
        avg_loss = metric_logger.meters['head_losses'].global_avg
        student_teacher_model.module.teacher.head.set_losses(avg_loss)
        student_teacher_model.module.student.head.set_losses(avg_loss)

    if utils.is_main_process():
        if args.num_heads == 1:
            writer.add_scalar("Train loss epoch", torch.Tensor([metric_logger.meters['loss'].global_avg]), epoch)
        else:
            avg_loss = metric_logger.meters['head_losses'].global_avg
            writer.add_scalars("Train loss epoch",
                               {f"head{i}": loss for i, loss in enumerate(avg_loss)},
                               epoch)

        d_loss = dino_loss[0] if hasattr(dino_loss, "__getitem__") else dino_loss
        if hasattr(d_loss, 'probs_pos'):
            writer.add_histogram("p(k) over Epochs", d_loss.probs_pos, epoch)
    print("Averaged stats:", metric_logger)
    return {k: meter.global_avg for k, meter in metric_logger.scalar_meters.items()}


def default_out_dir(loss, dset):
    return f"./experiments/{loss}-{dset}"


def make_out_dir(args):
    if args.output_dir is None:
        args.output_dir = default_out_dir(args.loss, args.loader)
    if args.new_run:
        n = 1
        dir_name = args.output_dir
        while Path(args.output_dir).is_dir():
            n += 1
            args.output_dir = f"{dir_name}{n}"
    Path(args.output_dir).mkdir(parents=True, exist_ok=not args.new_run)


def main():
    parser = get_args_parser()
    args = parser.parse_args()
    args = process_args(args)

    make_out_dir(args)
    writer = None
    if utils.is_main_process():
        writer = SummaryWriter(args.output_dir)
    with open(os.path.join(args.output_dir, "hp.json"), 'wt') as f:
        json.dump(vars(args), f, indent=4, default=str)
    train_dino(args, writer)


if __name__ == '__main__':
    main()

train_main.py
HHU-MMBS/TEMI-official-BMVC2023
# Copyright (c) XXXXXXXXXX
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Misc functions.

Mostly copy-paste from torchvision references or other public repos like DETR:
https://github.com/facebookresearch/detr/blob/master/util/misc.py
"""
import argparse
import os
import sys
import time
import math
import random
import datetime
import subprocess
import warnings
from ast import literal_eval
from collections import defaultdict, deque, OrderedDict

import numpy as np
import torch
from torch import nn
import torch.distributed as dist
from PIL import ImageFilter, ImageOps

from sklearn.metrics import normalized_mutual_info_score as nmi
from sklearn.metrics import adjusted_mutual_info_score as adjusted_nmi
from sklearn.metrics import adjusted_rand_score as adjusted_rand_index
from scipy.optimize import linear_sum_assignment


class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """
    def __init__(self, p=0.5, image_size=224, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min * (image_size/224.)
        self.radius_max = radius_max * (image_size/224.)

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )


class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """
    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img


def load_pretrained_weights(model, pretrained_weights, checkpoint_key, head=False, head_only=False):
    if not head and head_only:
        raise ValueError("head_only is True but head is False")
    if os.path.isfile(pretrained_weights):
        state_dict = torch.load(pretrained_weights, map_location="cpu")
        if checkpoint_key is not None and checkpoint_key in state_dict:
            print(f"Take key {checkpoint_key} in provided checkpoint dict")
            state_dict = state_dict[checkpoint_key]
        
        # remove `module.` prefix
        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
        if not head:
            # remove `backbone.` prefix induced by multicrop wrapper    
            state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
        if head_only:
            state_dict = OrderedDict([(k, v) for k, v in state_dict.items() if not 'backbone' in k])

        msg = model.load_state_dict(state_dict, strict=False)
        #print(f"Pretrained path {pretrained_weights}")
        print('Pretrained path {} and loaded with msg: {}'.format(pretrained_weights, msg))
    else:
        print("There is no reference weights available for this model => We use current weights (possibly random).")


def clip_gradients(model, clip):
    norms = []
    for name, p in model.named_parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            norms.append(param_norm.item())
            clip_coef = clip / (param_norm + 1e-6)
            if clip_coef < 1:
                p.grad.data.mul_(clip_coef)
    return norms


def cancel_gradients_last_layer(epoch, model, freeze_last_layer):
    if epoch >= freeze_last_layer:
        return
    for n, p in model.named_parameters():
        if "last_layer" in n:
            p.grad = None


def restart_from_checkpoint(ckp_path, run_variables=None, **kwargs):
    """
    Re-start from checkpoint
    """
    if not os.path.isfile(ckp_path):
        return
    print("Found checkpoint at {}".format(ckp_path))

    # open checkpoint file
    checkpoint = torch.load(ckp_path, map_location="cpu")

    # key is what to look for in the checkpoint file
    # value is the object to load
    # example: {'state_dict': model}
    for key, value in kwargs.items():
        if key in checkpoint and value is not None:
            d = {k.replace("module.", ""): v for k, v in checkpoint[key].items()}
            try:
                msg = value.load_state_dict(d)
                print("=> loaded '{}' from checkpoint '{}' with msg {}".format(key, ckp_path, msg))
            except RuntimeError:
                try:
                    msg = value.load_state_dict(d, strict=False)
                    print("=> loaded '{}' from checkpoint '{}' with msg {}".format(key, ckp_path, msg))
                except ValueError:
                    print("=> failed to load '{}' from checkpoint: '{}'".format(key, ckp_path))
        else:
            print("=> key '{}' not found in checkpoint: '{}'".format(key, ckp_path))

    # re load variable important for the run
    if run_variables is not None:
        for var_name in run_variables:
            if var_name in checkpoint:
                run_variables[var_name] = checkpoint[var_name]


def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0):
    warmup_schedule = np.array([])
    warmup_iters = warmup_epochs * niter_per_ep
    if warmup_epochs > 0:
        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)

    iters = np.arange(epochs * niter_per_ep - warmup_iters)
    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))

    schedule = np.concatenate((warmup_schedule, schedule))
    assert len(schedule) == epochs * niter_per_ep
    return schedule


def bool_flag(s):
    """
    Parse boolean arguments from the command line.
    """
    FALSY_STRINGS = {"off", "false", "0"}
    TRUTHY_STRINGS = {"on", "true", "1"}
    if s.lower() in FALSY_STRINGS:
        return False
    elif s.lower() in TRUTHY_STRINGS:
        return True
    else:
        raise argparse.ArgumentTypeError("invalid value for a boolean flag")


def fix_random_seeds(seed=31):
    """
    Fix random seeds.
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)


class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20, fmt=None):
        if fmt is None:
            fmt = "{median:.6f} ({global_avg:.6f})"
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def is_scalar(self):
        return isinstance(self.total, (float, int))

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        scalar = not isinstance(self.total, torch.Tensor)
        if scalar:
            total = torch.tensor(self.total, dtype=torch.float64, device="cuda")
        else:
            total = self.total
        count = torch.tensor(self.count, dtype=torch.int64, device="cuda")
        dist.barrier()
        dist.all_reduce(total)
        dist.all_reduce(count)
        self.count = int(count)
        self.total = total.item() if scalar else total

    def deq_to_tensor(self):
        if not isinstance(self.deque[-1], torch.Tensor):
            tensor_list = [torch.tensor(v) for v in self.deque]
        else:
            tensor_list = list(self.deque)
        return torch.stack(tensor_list).T

    @property
    def median(self):
        return self.deq_to_tensor().median().item()

    @property
    def avg(self):
        return self.deq_to_tensor().mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deq_to_tensor())

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        # TODO maybe improve this
        if isinstance(self.value, torch.Tensor):
            return '...'
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value)


def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict


class MetricLogger(object):
    def __init__(self, delimiter="\t"):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    @property
    def scalar_meters(self):
        return {k: v for k, v in self.meters.items() if v.is_scalar()}

    def update(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, torch.Tensor) and v.numel() == 1:
                v = v.item()
                assert isinstance(v, (float, int))
            self.meters[k].update(v)

    def update_raw(self, **kwargs):
        for k, v in kwargs.items():
            self.meters[k].update(v)

    def __getattr__(self, attr):
        if attr in self.meters:
            return self.meters[attr]
        if attr in self.__dict__:
            return self.__dict__[attr]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, attr))

    def __str__(self):
        loss_str = []
        for name, meter in self.scalar_meters.items():
            loss_str.append(
                "{}: {}".format(name, str(meter))
            )
        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):
        for meter in self.meters.values():
            meter.synchronize_between_processes()

    def add_meter(self, name, meter):
        self.meters[name] = meter

    def log_every(self, iterable, print_freq, header=None):
        i = 0
        if not header:
            header = ''
        start_time = time.time()
        end = time.time()
        iter_time = SmoothedValue(fmt='{avg:.6f}')
        data_time = SmoothedValue(fmt='{avg:.6f}')
        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
        if torch.cuda.is_available():
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}',
                'max mem: {memory:.0f}'
            ])
        else:
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}'
            ])
        MB = 1024.0 * 1024.0
        for obj in iterable:
            data_time.update(time.time() - end)
            yield obj
            iter_time.update(time.time() - end)
            if i % print_freq == 0 or i == len(iterable) - 1:
                eta_seconds = iter_time.global_avg * (len(iterable) - i)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if torch.cuda.is_available():
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time),
                        memory=torch.cuda.max_memory_allocated() / MB))
                else:
                    print(log_msg.format(
                        i, len(iterable), eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time)))
            i += 1
            end = time.time()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        print('{} Total time: {} ({:.6f} s / it)'.format(
            header, total_time_str, total_time / len(iterable)))


def get_sha():
    cwd = os.path.dirname(os.path.abspath(__file__))

    def _run(command):
        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
    sha = 'N/A'
    diff = "clean"
    branch = 'N/A'
    try:
        sha = _run(['git', 'rev-parse', 'HEAD'])
        subprocess.check_output(['git', 'diff'], cwd=cwd)
        diff = _run(['git', 'diff-index', 'HEAD'])
        diff = "has uncommited changes" if diff else "clean"
        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
    except Exception:
        pass
    message = f"sha: {sha}, status: {diff}, branch: {branch}"
    return message


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()


def all_reduce(value, *args, **kwargs):
    if is_dist_avail_and_initialized():
        dist.all_reduce(value, *args, **kwargs)


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()


def is_main_process():
    return get_rank() == 0


def save_on_master(*args, **kwargs):
    if is_main_process():
        torch.save(*args, **kwargs)


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


def init_distributed_mode(args):
    # launched with torch.distributed.launch
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    # launched with submitit on a slurm cluster
    elif 'SLURM_PROCID' in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    # launched naively with `python main_dino.py`
    # we manually add MASTER_ADDR and MASTER_PORT to env variables
    elif torch.cuda.is_available():
        print('Will run the code on one GPU.')
        args.rank, args.gpu, args.world_size = 0, 0, 1
        os.environ['MASTER_ADDR'] = '127.0.2.2'
        os.environ['MASTER_PORT'] = '29504'
    else:
        print('Does not support training without GPU.')
        sys.exit(1)

    dist.init_process_group(
        backend="nccl",
        init_method=args.dist_url,
        world_size=args.world_size,
        rank=args.rank,
    )

    torch.cuda.set_device(args.gpu)
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    dist.barrier()
    setup_for_distributed(args.rank == 0)


def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
    return [correct[:k].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]


def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


class LARS(torch.optim.Optimizer):
    """
    Almost copy-paste from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
    """
    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, eta=0.001,
                 weight_decay_filter=None, lars_adaptation_filter=None):
        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum,
                        eta=eta, weight_decay_filter=weight_decay_filter,
                        lars_adaptation_filter=lars_adaptation_filter)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            for p in g['params']:
                dp = p.grad

                if dp is None:
                    continue

                if p.ndim != 1:
                    dp = dp.add(p, alpha=g['weight_decay'])

                if p.ndim != 1:
                    param_norm = torch.norm(p)
                    update_norm = torch.norm(dp)
                    one = torch.ones_like(param_norm)
                    q = torch.where(param_norm > 0.,
                                    torch.where(update_norm > 0,
                                                (g['eta'] * param_norm / update_norm), one), one)
                    dp = dp.mul(q)

                param_state = self.state[p]
                if 'mu' not in param_state:
                    param_state['mu'] = torch.zeros_like(p)
                mu = param_state['mu']
                mu.mul_(g['momentum']).add_(dp)

                p.add_(mu, alpha=-g['lr'])


def get_params_groups(model):
    regularized = []
    not_regularized = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # we do not regularize biases nor Norm parameters
        if name.endswith(".bias") or len(param.shape) == 1:
            not_regularized.append(param)
        else:
            regularized.append(param)
    return [{'params': regularized}, {'params': not_regularized, 'weight_decay': 0.}]


def has_batchnorms(model):
    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
    for name, module in model.named_modules():
        if isinstance(module, bn_types):
            return True
    return False


def multi_scale(samples, model):
    v = None
    for s in [1, 1/2**(1/2), 1/2]:  # we use 3 different scales
        if s == 1:
            inp = samples.clone()
        else:
            inp = nn.functional.interpolate(samples, scale_factor=s, mode='bilinear', align_corners=False)
        feats = model(inp).clone()
        if v is None:
            v = feats
        else:
            v += feats
    v /= 3
    v /= v.norm()
    return v

# copied from https://github.com/elad-amrani/self-classifier/blob/39351384277d4541a0b7525a66770cacdd8f12c3/src/cls_eval.py#L406
def compute_metrics(targets, preds, min_samples_per_class=5, superclass_mapping=None, print_results=False):

    val_nmi = nmi(targets, preds)
    val_adjusted_nmi = adjusted_nmi(targets, preds)
    val_adjusted_rand_index = adjusted_rand_index(targets, preds)

    # compute accuracy
    num_classes = max(targets.max(), preds.max()) + 1
    count_matrix = np.zeros((num_classes, num_classes), dtype=np.int32)
    for ii in range(preds.shape[0]):
        count_matrix[preds[ii], targets[ii]] += 1
    reassignment = np.dstack(linear_sum_assignment(count_matrix.max() - count_matrix))[0]

    if len(np.unique(preds)) > len(np.unique(targets)):  # if using over-clustering, append remaining clusters to best option
        for cls_idx in np.unique(preds):
            if reassignment[cls_idx, 1] not in targets:
                reassignment[cls_idx, 1] = count_matrix[cls_idx].argmax()

    if superclass_mapping is not None:
        count_matrix = np.zeros((num_classes, num_classes), dtype=np.int32)
        for ii in range(preds.shape[0]):
            count_matrix[preds[ii], superclass_mapping[targets[ii]]] += 1
        for ii in range(len(reassignment[:, 1])):
            reassignment[ii, 1] = superclass_mapping[reassignment[ii, 1]]
    acc = count_matrix[reassignment[:, 0], reassignment[:, 1]].sum().astype(np.float32) / preds.shape[0]
    
    if print_results:
        print('=> number of samples: {}'.format(len(targets)))
        print('=> number of unique assignments: {}'.format(len(set(preds))))
        print('=> NMI: {:.3f}%'.format(val_nmi * 100.0))
        print('=> Adjusted NMI: {:.3f}%'.format(val_adjusted_nmi * 100.0))
        print('=> Adjusted Rand-Index: {:.3f}%'.format(val_adjusted_rand_index * 100.0))
        print('=> Accuracy: {:.3f}%'.format(acc * 100.0))

    # extract max accuracy classes
    num_samples_per_class = count_matrix[reassignment[:, 0], :].sum(axis=1)
    acc_per_class = np.where(num_samples_per_class >= min_samples_per_class,
                             count_matrix[reassignment[:, 0], reassignment[:, 1]] / num_samples_per_class, 0)
    max_acc_classes = np.argsort(acc_per_class)[::-1]
    acc_per_class = acc_per_class[max_acc_classes]
    num_samples_per_class = num_samples_per_class[max_acc_classes]

    return acc * 100, val_nmi * 100.0, val_adjusted_nmi * 100.0, val_adjusted_rand_index * 100.0 


def _backbone_param(model):
    try:
        return model.conv1.weight
    except AttributeError:
        return next(model.parameters())


def backbone_dtype(model):
    if not isinstance(model, nn.Module):
        return torch.float
    return _backbone_param(model).dtype


@torch.no_grad()
def embed_dim(args, model):
    from model_builders import load_embeds
    try:
        return load_embeds(args).shape[-1]
    except Exception:
        pass
    if isinstance(model, nn.Module):
        p = _backbone_param(model)
        dummy_in = torch.empty(1, 3, args.vit_image_size, args.vit_image_size,
                               device=p.device, dtype=p.dtype)
        dummy_out = model(dummy_in)
        return dummy_out.size(-1)
    raise ValueError('Could not infer embed_dim')


def kv_pair(s):
    # For extra arparse arguments
    k, v = s.split("=")
    try:
        # v is float/int etc, parse it
        v = literal_eval(v)
    except (ValueError, SyntaxError):
        pass
    return k, v

utils.py
HHU-MMBS/TEMI-official-BMVC2023