Official implementation of "Towards Distribution-Agnostic Generalized Category Discovery" (NIPS 2023)
cifar_10_root = 'cifar10/'
cifar_100_root = 'cifar100/'
imagenet_root = '/ImageNet/'
exp_root = 'dev_outputs/'
config.py
JianhongBai/BaCon
import torch
from torch.utils.data import Dataset
import os
from PIL import Image
import cv2
import random
import numpy as np
import pickle
from torchvision import datasets

class LT_Dataset(Dataset):

  def __init__(self, root, txt, transform=None, returnPath=False):
    self.img_path = []
    self.targets = []
    self.root = root
    self.transform = transform
    self.returnPath = returnPath
    self.txt = txt

    with open(txt) as f:
      for line in f:
        self.img_path.append(os.path.join(root, line.split()[0]))
        self.targets.append(int(line.split()[1]))

    self.uq_idxs = np.array(range(len(self)))

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

  def __getitem__(self, index):

    path = self.img_path[index]
    label = self.targets[index]

    with open(path, 'rb') as f:
      sample = Image.open(f).convert('RGB')

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

    uq_idx = self.uq_idxs[index]
    return sample, label, uq_idx


def unpickle(file):
  with open(file, 'rb') as fo:
    dict = pickle.load(fo)
  return dict


class Unsupervised_LT_Dataset(LT_Dataset):
  def __init__(self, returnIdx=False, returnLabel=False, return_ood_flag=False, **kwds):
    super().__init__(**kwds)
    self.returnIdx = returnIdx
    self.returnLabel = returnLabel
    self.return_ood_flag = return_ood_flag

  def __getitem__(self, index):
    path = self.img_path[index]
    label = self.labels[index]

    if not os.path.isfile(path):
      path = path + ".gz"

    with open(path, 'rb') as f:
      sample = Image.open(f).convert('RGB')

    samples = [self.transform(sample), self.transform(sample)]
    if self.returnIdx and (not self.returnLabel):
      return torch.stack(samples), index
    elif (not self.returnIdx) and self.returnLabel:
      return torch.stack(samples), label
    elif self.returnIdx and self.returnLabel:
      return torch.stack(samples), label, index
    elif self.return_ood_flag:
      return torch.stack(samples), -1
    else:
      return torch.stack(samples)


class LT_Dataset_Meta(Dataset):

  def __init__(self, root, txt, classRange, transform=None, returnPath=False):
    self.img_path = []
    self.labels = []
    self.root = root
    self.transform = transform
    self.returnPath = returnPath

    self.txt = txt
    with open(txt) as f:
      for line in f:
        self.img_path.append(os.path.join(root, line.split()[0]))
        self.labels.append(int(line.split()[1]))

    self.setClassRange(classRange)

  def setClassRange(self, classRange):
    self.epochList = []
    self.classRange = classRange
    for idx, label in enumerate(self.labels):
      if label in self.classRange:
        self.epochList.append(idx)

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

  def __getitem__(self, index):
    sampleIdx = self.epochList[index]

    path = self.img_path[sampleIdx]
    label = self.labels[sampleIdx]
    label = self.classRange.index(label)

    with open(path, 'rb') as f:
      sample = Image.open(f).convert('RGB')

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

    if not self.returnPath:
      return sample, label, index
    else:
      return sample, label, index, path.replace(self.root, '')


class Supervised_Folder_Dataset(datasets.ImageFolder):

  def __init__(self, root, transform=None, target_transform=None):
    super().__init__(root, transform, target_transform)

  def __getitem__(self, index):
    path, target = self.samples[index]
    sample = self.loader(path)
    if self.transform is not None:
      sample = self.transform(sample)
    if self.target_transform is not None:
      target = self.target_transform(target)

    return sample, target, index

class Unsupervised_Folder_Dataset(datasets.ImageFolder):

  def __init__(self, root, transform=None, target_transform=None, return_ood_flag=False):
    super().__init__(root, transform, target_transform)

  def __getitem__(self, index):
    path, target = self.samples[index]
    sample = self.loader(path)
    if self.transform is not None:
      samples = [self.transform(sample), self.transform(sample)]
    if self.target_transform is not None:
      target = self.target_transform(target)
    if self.return_ood_flag:
      return torch.stack(samples), -1
    return torch.stack(samples)


data/LT_dataset.py
JianhongBai/BaCon
from torchvision import transforms

import torch

def get_transform(transform_type='imagenet', image_size=32, args=None):
    
    if transform_type == 'imagenet':

        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
        interpolation = args.interpolation
        crop_pct = args.crop_pct

        train_transform = transforms.Compose([
            transforms.Resize(int(image_size / crop_pct), interpolation),
            transforms.RandomCrop(image_size),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ColorJitter(),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=torch.tensor(mean),
                std=torch.tensor(std))
        ])

        test_transform = transforms.Compose([
            transforms.Resize(int(image_size / crop_pct), interpolation),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=torch.tensor(mean),
                std=torch.tensor(std))
        ])
    
    else:

        raise NotImplementedError

    return (train_transform, test_transform)

data/augmentations/__init__.py
JianhongBai/BaCon
from torchvision.datasets import CIFAR10, CIFAR100
from copy import deepcopy
import numpy as np
import torch
from data.data_utils import subsample_instances
from config import cifar_100_root, cifar_10_root

class CustomCIFAR10(CIFAR10):

    def __init__(self, sublist=[], *args, **kwargs):

        super(CustomCIFAR10, self).__init__(*args, **kwargs)

        self.txt = sublist
        self.uq_idxs = np.array(range(len(self)))
        if len(sublist) > 0:
            self.data = self.data[sublist]
            self.targets = np.array(self.targets)[sublist].tolist()
            self.uq_idxs = self.uq_idxs[sublist]

    def __getitem__(self, item):

        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]
        return img, label, uq_idx

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


class CustomCIFAR100(CIFAR100):

    def __init__(self, sublist=[], return_false_mask=False, return_true_mask=False, *args, **kwargs):
        super(CustomCIFAR100, self).__init__(*args, **kwargs)
        self.txt = sublist

        self.uq_idxs = np.array(range(len(self)))
        self.return_false_mask = return_false_mask
        self.return_true_mask = return_true_mask

        if len(sublist) > 0:
            self.data = self.data[sublist]
            self.targets = np.array(self.targets)[sublist].tolist()
            self.uq_idxs = self.uq_idxs[sublist]

    def __getitem__(self, item):
        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]
        return img, label, uq_idx

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


def subsample_dataset(dataset, idxs):

    # Allow for setting in which all empty set of indices is passed

    if len(idxs) > 0:

        dataset.data = dataset.data[idxs]
        dataset.targets = np.array(dataset.targets)[idxs].tolist()
        dataset.uq_idxs = dataset.uq_idxs[idxs]

        return dataset

    else:

        return None


def subsample_classes(dataset, include_classes=(0, 1, 8, 9)):

    cls_idxs = [x for x, t in enumerate(dataset.targets) if t in include_classes]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    # dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.targets)

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.targets == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs

def get_lt_dist(cls_num, img_max, imb_factor):
    img_num_per_cls = []
    for cls_idx in range(cls_num):
        num = img_max * (imb_factor ** (cls_idx / (cls_num - 1.0)))
        img_num_per_cls.append(int(num))
    return torch.tensor(img_num_per_cls)


def get_cifar_100_datasets(train_transform, test_transform, args=None):

    whole_training_set = CustomCIFAR100(root=cifar_100_root, train=True, transform=train_transform, download=True)
    whole_training_set.target_transform = None

    l_k_uq_idxs = torch.load(f'data_uq_idxs/cifar100_k80_imb{args.imb_ratio}/l_k_uq_idxs.pt')
    unl_unk_uq_idxs = torch.load(f'data_uq_idxs/cifar100_k80_imb{args.imb_ratio}/unl_unk_uq_idxs.pt')
    unl_k_uq_idxs = torch.load(f'data_uq_idxs/cifar100_k80_imb{args.imb_ratio}/unl_k_uq_idxs.pt')

    train_dataset_labelled = subsample_dataset(deepcopy(whole_training_set), l_k_uq_idxs)
    lt_unlabeled_known_dataset = subsample_dataset(deepcopy(whole_training_set), unl_k_uq_idxs)
    lt_unlabeled_unknown_dataset = subsample_dataset(deepcopy(whole_training_set), unl_unk_uq_idxs)
    train_dataset_unlabelled = torch.utils.data.ConcatDataset([lt_unlabeled_known_dataset, lt_unlabeled_unknown_dataset])

    test_dataset = CustomCIFAR100(root=cifar_100_root, transform=test_transform, train=False)

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'test': test_dataset,
    }

    return all_datasets

def get_cifar_10_datasets(train_transform, test_transform, args=None):

    whole_training_set = CustomCIFAR10(root=cifar_10_root, train=True, transform=train_transform, download=True)
    whole_training_set.target_transform = None

    l_k_uq_idxs = torch.load(f'data_uq_idxs/cifar10_k5_imb{args.imb_ratio}/l_k_uq_idxs.pt')
    unl_unk_uq_idxs = torch.load(f'data_uq_idxs/cifar10_k5_imb{args.imb_ratio}/unl_unk_uq_idxs.pt')
    unl_k_uq_idxs = torch.load(f'data_uq_idxs/cifar10_k5_imb{args.imb_ratio}/unl_k_uq_idxs.pt')

    train_dataset_labelled = subsample_dataset(deepcopy(whole_training_set), l_k_uq_idxs)
    lt_unlabeled_known_dataset = subsample_dataset(deepcopy(whole_training_set), unl_k_uq_idxs)
    lt_unlabeled_unknown_dataset = subsample_dataset(deepcopy(whole_training_set), unl_unk_uq_idxs)
    train_dataset_unlabelled = torch.utils.data.ConcatDataset([lt_unlabeled_known_dataset, lt_unlabeled_unknown_dataset])
    
    test_dataset = CustomCIFAR10(root=cifar_10_root, transform=test_transform, train=False)

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'test': test_dataset,
    }

    return all_datasets

data/cifar.py
JianhongBai/BaCon
import numpy as np
from torch.utils.data import Dataset

def subsample_instances(dataset, prop_indices_to_subsample=0.8):

    np.random.seed(0)
    subsample_indices = np.random.choice(range(len(dataset)), replace=False,
                                         size=(int(prop_indices_to_subsample * len(dataset)),))

    return subsample_indices

class MergedDataset(Dataset):

    """
    Takes two datasets (labelled_dataset, unlabelled_dataset) and merges them
    Allows you to iterate over them in parallel
    """

    def __init__(self, labelled_dataset, unlabelled_dataset):

        self.labelled_dataset = labelled_dataset
        self.unlabelled_dataset = unlabelled_dataset
        self.target_transform = None

    def __getitem__(self, item):

        if item < len(self.labelled_dataset):
            img, label, uq_idx = self.labelled_dataset[item]
            labeled_or_not = 1

        else:

            img, label, uq_idx = self.unlabelled_dataset[item - len(self.labelled_dataset)]
            labeled_or_not = 0


        return img, label, uq_idx, np.array([labeled_or_not])

    def __len__(self):
        return len(self.unlabelled_dataset) + len(self.labelled_dataset)

data/data_utils.py
JianhongBai/BaCon
from data.data_utils import MergedDataset

from data.cifar import get_cifar_10_datasets, get_cifar_100_datasets
from data.imagenet import get_imagenet_100_datasets
from copy import deepcopy
import pickle
import os


get_dataset_funcs = {
    'cifar10': get_cifar_10_datasets,
    'cifar100': get_cifar_100_datasets,
'imagenet100': get_imagenet_100_datasets,
}


def get_datasets(dataset_name, train_transform, test_transform, args):

    """
    :return: train_dataset: MergedDataset which concatenates labelled and unlabelled
             test_dataset,
             unlabelled_train_examples_test,
             datasets
    """

    if dataset_name not in get_dataset_funcs.keys():
        raise ValueError

    # Get datasets
    get_dataset_f = get_dataset_funcs[dataset_name]
    datasets = get_dataset_f(train_transform=train_transform, test_transform=test_transform, args=args)

    # Train split (labelled and unlabelled classes) for training
    train_dataset = MergedDataset(labelled_dataset=deepcopy(datasets['train_labelled']),
                                  unlabelled_dataset=deepcopy(datasets['train_unlabelled']))

    test_dataset = datasets['test']
    unlabelled_train_examples_test = None

    return train_dataset, test_dataset, unlabelled_train_examples_test, datasets


def get_class_splits(args):

    if args.dataset_name == 'cifar10':
        args.train_classes = range(5)
        args.unlabeled_classes = range(5, 10)
    elif args.dataset_name == 'cifar100':
        args.train_classes = range(80)
        args.unlabeled_classes = range(80, 100)
    elif args.dataset_name == 'imagenet100':
        args.train_classes = range(50)
        args.unlabeled_classes = range(50, 100)


    return args

data/get_datasets.py
JianhongBai/BaCon
import torchvision
import numpy as np

import os

from copy import deepcopy
from data.data_utils import subsample_instances
from config import imagenet_root
import torch

class ImageNetBase(torchvision.datasets.ImageFolder):

    def __init__(self, root, transform):

        super(ImageNetBase, self).__init__(root, transform)

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, item):

        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]

        return img, label, uq_idx


def subsample_dataset(dataset, idxs, args):


    img_path_ = []
    for i in idxs:
        img_path_.append(dataset.img_path[i])
    dataset.img_path = img_path_

    dataset.targets = np.array(dataset.targets)[idxs].tolist()
    dataset.uq_idxs = dataset.uq_idxs[idxs]
    return dataset


def subsample_classes(dataset, include_classes=list(range(1000)), args=None):

    cls_idxs = [x for x, t in enumerate(dataset.targets) if t in include_classes]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs, args)

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = list(set(train_dataset.targets))

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(np.array(train_dataset.targets) == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs

def get_lt_dist(cls_num, img_max, imb_factor):
    img_num_per_cls = []
    for cls_idx in range(cls_num):
        num = img_max * (imb_factor ** (cls_idx / (cls_num - 1.0)))
        img_num_per_cls.append(int(num))
    return torch.tensor(img_num_per_cls)

def get_imagenet_100_datasets(train_transform, test_transform, train_classes=range(50),
                           prop_train_labels=0.5, split_train_val=False, seed=0, args=None):

    from data.LT_dataset import LT_Dataset
    txt_path = "splits/imagenet100/imageNet_100_LT_train.txt"
    test_txt_path = "splits/imagenet100/ImageNet_100_test.txt"
    total_class = 100
    args.anno_ratio = 50
    seed = 416
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)

    whole_training_set = LT_Dataset(root=imagenet_root, txt=txt_path, transform=train_transform)

    whole_training_set.target_transform = None

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    known_classes = train_classes
    unknown_classes = range(len(train_classes), total_class)
    whole_known_classes_set = subsample_classes(deepcopy(whole_training_set), include_classes=known_classes,
                                                args=args)

    labeled_known_idxs = []
    unlabeled_known_idxs = []
    known_classes = np.unique(whole_known_classes_set.targets)
    for cls in known_classes:
        cls_idxs = np.where(whole_known_classes_set.targets == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False,
                              size=((int((1 - args.anno_ratio / 100) * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        labeled_known_idxs.extend(t_)
        unlabeled_known_idxs.extend(v_)

    lt_labeled_known_dataset = subsample_dataset(deepcopy(whole_known_classes_set), labeled_known_idxs, args)
    lt_unlabeled_known_dataset = subsample_dataset(deepcopy(whole_known_classes_set), unlabeled_known_idxs,
                                                   args)

    lt_unlabeled_unknown_dataset = subsample_classes(deepcopy(whole_training_set),
                                                     include_classes=range(len(train_classes), total_class),
                                                     args=args)

    # Get test set for all classes
    test_dataset = LT_Dataset(root=imagenet_root, txt=test_txt_path, transform=test_transform)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = lt_labeled_known_dataset
    train_dataset_unlabelled = torch.utils.data.ConcatDataset(
        [lt_unlabeled_known_dataset, lt_unlabeled_unknown_dataset])
    val_dataset_labelled = None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
        'bl_train_unlabelled': None
    }
    return all_datasets
data/imagenet.py
JianhongBai/BaCon
import numpy as np
import torch

device = torch.device('cuda:0')

def compute_est_adjustment(tro=1.0, args=None):
    """compute the base probabilities"""

    dist = args.est_dist
    label_freq_array = np.array(dist.cpu())
    label_freq_array = label_freq_array / label_freq_array.sum()
    adjustments = np.log(label_freq_array ** tro + 1e-12)
    adjustments = torch.from_numpy(adjustments)
    adjustments = adjustments.to(device)
    return adjustments
logits_adj_utils.py
JianhongBai/BaCon
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD, lr_scheduler
from torch.utils.data import DataLoader
from tqdm import tqdm
import sys
sys.path.append('./')
from data.augmentations import get_transform
from data.get_datasets import get_datasets, get_class_splits

from util.general_utils import AverageMeter, init_experiment
from util.cluster_and_log_utils import log_accs_from_preds
from config import exp_root
from model.loss import info_nce_logits, SupConLoss, DistillLoss, ContrastiveLearningViewGenerator, get_params_groups
from copy import deepcopy
from sklearn.cluster import KMeans
# from cuml.cluster import KMeans as cuKMeans
from data.cifar import CustomCIFAR100, cifar_100_root
import random
from model.dist_est import dist_est
from model.reg_loss import compute_reg_loss
from model.softconloss import compute_softconloss
import os
import sys

class CE_Head(nn.Module):
    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True,
                 nlayers=3, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        self.apply(self._init_weights)
        self.last_layer = nn.utils.weight_norm(nn.Linear(in_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):
            torch.nn.init.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 = nn.functional.normalize(x, dim=-1, p=2)
        logits = self.last_layer(x)
        return logits


def get_mean_lr(optimizer):
    return torch.mean(torch.Tensor([param_group['lr'] for param_group in optimizer.param_groups])).item()


def train_dual(ce_backbone, ce_head, cl_backbone, cl_head, train_loader, test_loader, args):
    def set_model(train=False, eval=False):
        assert (train ^ eval)
        if train:
            student_ce.train()
            student_cl.train()
        elif eval:
            student_ce.eval()
            student_cl.eval()

    from util.cluster_and_log_utils import set_args_mmf
    set_args_mmf(args, train_loader)

    save_path = 'checkpoints/' + os.path.join(*args.model_path.split('/')[3:-1])
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    student_ce = nn.Sequential(ce_backbone, ce_head).to(device)
    student_cl = nn.Sequential(cl_backbone, cl_head).to(device)

    params_groups_cl = list(cl_head.parameters()) + list(cl_backbone.parameters())
    params_groups_ce = get_params_groups(student_ce)
    optimizer_ce = SGD(params_groups_ce, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    exp_lr_scheduler_ce = lr_scheduler.CosineAnnealingLR(
        optimizer_ce,
        T_max=args.epochs,
        eta_min=args.lr * 1e-3,
    )

    optimizer_cl = SGD(params_groups_cl, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    exp_lr_scheduler_cl = lr_scheduler.CosineAnnealingLR(
        optimizer_cl,
        T_max=args.epochs,
        eta_min=args.lr * 1e-3,
    )
    args.current_epoch = 0
    best_test_acc_all_cl = -1
    cluster_criterion = DistillLoss(
        args.warmup_teacher_temp_epochs,
        args.epochs,
        args.n_views,
        args.warmup_teacher_temp,
        args.teacher_temp,
        args=args,
    )

    set_model(eval=True)
    est_count = dist_est(ce_backbone, cl_backbone, ce_head, train_all_test_trans_loader, args)
    set_model(train=True)


    for epoch in range(args.epochs):
        args.current_epoch = epoch
        loss_record_ce = AverageMeter()
        loss_record_cl = AverageMeter()
        for batch_idx, batch in enumerate(train_loader):
            images_, class_labels, uq_idxs, mask_lab = batch
            mask_lab = mask_lab[:, 0]
            class_labels, mask_lab = class_labels.cuda(non_blocking=True), mask_lab.cuda(non_blocking=True).bool()
            images = torch.cat(images_, dim=0).cuda(non_blocking=True)

            x = ce_backbone.prepare_tokens(images)

            for i, blk in enumerate(ce_backbone.blocks):
                if i < args.grad_from_block:
                    x = blk(x)   # get fixed feature

            for i, blk in enumerate(ce_backbone.blocks):
                if i >= args.grad_from_block:
                    ce_backbone_feature = blk(x)
            ce_backbone_feature = ce_backbone.norm(ce_backbone_feature)
            ce_backbone_feature = ce_backbone_feature[:, 0]
            student_out = ce_head(ce_backbone_feature)

            for i, blk in enumerate(cl_backbone.blocks):
                if i >= args.grad_from_block:
                    cl_backbone_feature = blk(x)
            cl_backbone_feature = cl_backbone.norm(cl_backbone_feature)
            cl_backbone_feature = cl_backbone_feature[:, 0]
            cl_proj_feature = cl_head(cl_backbone_feature)

            ####################### COMPUTE LOSS #######################
            pstr = ''
            teacher_out = student_out.detach()

            # clustering, unsup
            cluster_loss = cluster_criterion(student_out, teacher_out, epoch)

            # clustering, sup
            sup_logits = torch.cat([f[mask_lab] for f in (student_out / 0.1).chunk(2)], dim=0)
            sup_labels = torch.cat([class_labels[mask_lab] for _ in range(2)], dim=0)

            cls_loss = nn.CrossEntropyLoss()(sup_logits, sup_labels)

            me_max_loss = compute_reg_loss(student_out, est_count, args)

            pstr += f'cls_loss: {cls_loss.item():.2f} '
            pstr += f'cps_loss: {cluster_loss.item():.2f} '
            pstr += f'reg_loss: {me_max_loss.item():.2f} '
            cluster_loss += args.memax_weight * me_max_loss
            loss_ce = (1 - args.sup_weight) * cluster_loss + args.sup_weight * cls_loss

            loss_record_ce.update(loss_ce.item(), class_labels.size(0))
            optimizer_ce.zero_grad()
            loss_ce.backward()
            optimizer_ce.step()

            loss_cl = 0
            # for CL part
            # represent learning, unsup
            cl_proj_feature = torch.nn.functional.normalize(cl_proj_feature, dim=-1)

            contrastive_logits, contrastive_labels = info_nce_logits(features=cl_proj_feature)
            contrastive_loss = torch.nn.CrossEntropyLoss()(contrastive_logits, contrastive_labels)

            # representation learning, sup
            sup_cl_proj_feature = torch.cat([f[mask_lab].unsqueeze(1) for f in cl_proj_feature.chunk(2)], dim=1)
            sup_con_labels = class_labels[mask_lab]

            if epoch >= args.ce_warmup:
                sup_con_loss = SupConLoss()(sup_cl_proj_feature, labels=sup_con_labels)
                soft_con_loss = compute_softconloss(student_out, cl_proj_feature, sup_con_labels, sup_cl_proj_feature, mask_lab, args)
                loss_cl += ((1 - args.sup_weight) * contrastive_loss + (args.sup_weight / 2) * sup_con_loss)
                loss_cl += (args.sup_weight / 2) * soft_con_loss
            else:
                sup_con_loss = SupConLoss()(sup_cl_proj_feature, labels=sup_con_labels)
                loss_cl += ((1 - args.sup_weight) * contrastive_loss + args.sup_weight * sup_con_loss)

            pstr += f'sup_con_loss: {sup_con_loss.item():.2f} '
            pstr += f'contrastive_loss: {contrastive_loss.item():.2f} '

            loss_record_cl.update(loss_cl.item(), class_labels.size(0))
            optimizer_cl.zero_grad()
            loss_cl.backward()
            optimizer_cl.step()

            if batch_idx % args.print_freq == 0:
                args.logger.info('Epoch: [{}][{}/{}]\t loss_ce {:.5f} loss_cl {:.5f}\t {}'
                                 .format(epoch, batch_idx, len(train_loader), loss_ce.item(), loss_cl.item(), pstr))

        args.logger.info('Train Epoch: {} Avg Loss_ce: {:.2f} Avg Loss_cl: {:.2f}'.format(epoch, loss_record_ce.avg, loss_record_cl.avg))

        if (epoch+1) % args.est_freq == 0:
            set_model(eval=True)
            est_count = dist_est(ce_backbone, cl_backbone, ce_head, train_all_test_trans_loader, args)

        if epoch % args.test_freq == 0:
            args.logger.info('Testing on disjoint test set...')
            with torch.no_grad():
                all_acc_test_cl, old_acc_test_cl, new_acc_test_cl, acc_list_cl, cl_ind_map = test(
                    student_cl,
                    test_loader,
                    epoch=epoch,
                    save_name='Test ACC',
                    args=args,
                    train_loader=train_loader)
            
            args.logger.info(
                    'Test Accuracies CL: All {:.1f} | Old {:.1f} | New {:.1f}'.format(all_acc_test_cl,
                                                                                      old_acc_test_cl,
                                                                                      new_acc_test_cl))

        # Step schedule
        exp_lr_scheduler_ce.step()
        exp_lr_scheduler_cl.step()

        if epoch % args.test_freq == 0 and all_acc_test_cl > best_test_acc_all_cl:

            best_test_acc_new_cl = new_acc_test_cl
            best_test_acc_old_cl = old_acc_test_cl
            best_test_acc_all_cl = all_acc_test_cl

            save_dict_cl = {
                'ce_backbone': ce_backbone.state_dict(),
                'ce_head': ce_head.state_dict(),
                'cl_backbone': cl_backbone.state_dict(),
                'cl_head': cl_head.state_dict(),
            }

            torch.save(save_dict_cl, save_path + f'/model_epoch{epoch}.pt')
            args.logger.info("model saved to {}.".format(save_path))

        if epoch >= args.stop_epoch:
            break

    args.logger.info(
        f'Metrics with best model on test set: All: {best_test_acc_all_cl:.1f} Old: {best_test_acc_old_cl:.1f} New: {best_test_acc_new_cl:.1f} ')


def test(model, test_loader, epoch, save_name, args, train_loader):
    model.eval()

    all_feats = []
    targets = np.array([])
    mask = np.array([])
    print('Collating features...')
    # First extract all features
    for batch_idx, batch in enumerate(test_loader):
        batch = batch[:3]
        (images, label, _) = batch
        images = images.cuda()

        # Pass features through base model and then additional learnable transform (linear layer)
        feats = model[0](images)  # follow GCD: clustering on normalized backbone feature

        feats = torch.nn.functional.normalize(feats, dim=-1)

        all_feats.append(feats.detach().cpu().numpy())
        targets = np.append(targets, label.cpu().numpy())
        mask = np.append(mask, np.array([True if x.item() in range(len(args.train_classes))
                                         else False for x in label]))

    # -----------------------
    # K-MEANS
    # -----------------------
    print('Fitting K-Means...')
    all_feats = np.concatenate(all_feats)
    kmeans = KMeans(n_clusters=args.num_labeled_classes + args.num_unlabeled_classes, random_state=0).fit(all_feats)
    preds = kmeans.labels_
    print('Done!')

    all_acc, old_acc, new_acc, acc_list, ind_map = log_accs_from_preds(y_true=targets, y_pred=preds, mask=mask,
                                                    T=epoch, eval_funcs=args.eval_funcs, save_name=save_name,
                                                    args=args, train_loader=train_loader)

    return all_acc, old_acc, new_acc, acc_list, ind_map


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', default=256, type=int)
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--eval_funcs', type=list, default=['v2'])
    parser.add_argument('--dataset_name', type=str, default='cifar100')
    parser.add_argument('--prop_train_labels', type=float, default=0.5)
    parser.add_argument('--grad_from_block', type=int, default=11)
    parser.add_argument('--lr', type=float, default=0.1)
    parser.add_argument('--gamma', type=float, default=0.1)
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--weight_decay', type=float, default=5e-5)
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--exp_root', type=str, default=exp_root)
    parser.add_argument('--transform', type=str, default='imagenet')
    parser.add_argument('--sup_weight', type=float, default=0.35)
    parser.add_argument('--n_views', default=2, type=int)
    parser.add_argument('--memax_weight', type=float, default=1)
    parser.add_argument('--warmup_teacher_temp', default=0.07, type=float)
    parser.add_argument('--teacher_temp', default=0.04, type=float)
    parser.add_argument('--warmup_teacher_temp_epochs', default=30, type=int)
    parser.add_argument('--fp16', action='store_true', default=False)
    parser.add_argument('--print_freq', default=10, type=int)
    parser.add_argument('--exp_name', default='cifar100', type=str)
    parser.add_argument('--local_rank', default=-1, type=int)
    parser.add_argument('--p', default=1.1, type=float)
    parser.add_argument('--test_freq', default=1, type=int)
    parser.add_argument('--ce_warmup', default=1, type=int)
    parser.add_argument('--est_freq', default=10, type=int)
    parser.add_argument('--labeled_classes', default=80, type=int)
    parser.add_argument('--config_file', default='', type=str)
    parser.add_argument('--alpha', default=0.8, type=float)
    parser.add_argument('--beta', default=0.5, type=float)
    parser.add_argument('--tro', default=0.5, type=float)
    parser.add_argument('--stop_epoch', default=200, type=int)
    parser.add_argument('--imb_ratio', default=100, type=int)

    # ----------------------
    # INIT
    # ----------------------

    args = parser.parse_args()
    pid = os.getpid()
    print('MY PID：', pid)

    if args.config_file != '':
        with open('configs/' + args.config_file, 'r') as f:
            args = parser.parse_args(f.read().split())


    device = torch.device('cuda:0')
    args = get_class_splits(args)
    total_class = 100 if args.dataset_name != 'cifar10' else 10

    args.train_classes = range(args.labeled_classes)
    args.unlabeled_classes = range(args.labeled_classes, total_class)

    args.num_labeled_classes = len(args.train_classes)
    args.num_unlabeled_classes = len(args.unlabeled_classes)
    args.num_classes = args.num_labeled_classes + args.num_unlabeled_classes
    args.mlp_out_dim = args.num_labeled_classes + args.num_unlabeled_classes
    init_experiment(args, runner_name=['BaCon'])
    args.logger.info(f'Using evaluation function {args.eval_funcs[0]} to print results')
    
    torch.backends.cudnn.benchmark = True

    args.interpolation = 3
    args.crop_pct = 0.875

    backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')
    
    args.image_size = 224
    args.feat_dim = 768
    args.num_mlp_layers = 3

    # ----------------------
    # HOW MUCH OF BASE MODEL TO FINETUNE
    # ----------------------
    for m in backbone.parameters():
        m.requires_grad = False

    # Only finetune layers from block 'args.grad_from_block' onwards
    for name, m in backbone.named_parameters():
        if 'block' in name:
            block_num = int(name.split('.')[1])
            if block_num >= args.grad_from_block:
                m.requires_grad = True

    args.logger.info('model build')

    # --------------------
    # CONTRASTIVE TRANSFORM
    # --------------------
    train_transform, test_transform = get_transform(args.transform, image_size=args.image_size, args=args)
    train_transform = ContrastiveLearningViewGenerator(base_transform=train_transform, n_views=args.n_views)
    # --------------------
    # DATASETS
    # --------------------
    train_dataset, test_dataset, unlabelled_train_examples_test, datasets = get_datasets(args.dataset_name,
                                                                                         train_transform,
                                                                                         test_transform,
                                                                                         args)

    seed = torch.randint(0, 100000, (1,)).item()
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

    # --------------------
    # SAMPLER
    # Sampler which balances labelled and unlabelled examples in each batch
    # --------------------
    label_len = len(train_dataset.labelled_dataset)
    unlabelled_len = len(train_dataset.unlabelled_dataset)

    sample_weights = [1 if i < label_len else label_len / unlabelled_len for i in range(len(train_dataset))]
    sample_weights = torch.DoubleTensor(sample_weights)
    sampler = torch.utils.data.WeightedRandomSampler(sample_weights, num_samples=len(train_dataset))

    train_all_test_trans = deepcopy(train_dataset)
    train_labelled_test_trans = deepcopy(train_dataset.labelled_dataset)
    train_unlabelled_test_trans = deepcopy(train_dataset.unlabelled_dataset)

    train_all_test_trans.labelled_dataset.transform = test_transform

    train_all_test_trans.unlabelled_dataset.datasets[0].transform = test_transform
    train_all_test_trans.unlabelled_dataset.datasets[1].transform = test_transform
    train_unlabelled_test_trans.datasets[0].transform = test_transform
    train_unlabelled_test_trans.datasets[1].transform = test_transform
    train_labelled_test_trans.transform = test_transform


    # --------------------
    # DATALOADERS
    # --------------------
    train_loader = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False,
                              sampler=sampler, drop_last=True, pin_memory=True)
    test_loader_unlabelled = DataLoader(unlabelled_train_examples_test, num_workers=args.num_workers,
                                        batch_size=256, shuffle=False, pin_memory=False)
    test_loader_labelled = DataLoader(test_dataset, num_workers=args.num_workers,
                                      batch_size=256, shuffle=False, pin_memory=False)

    train_all_test_trans_loader = DataLoader(train_all_test_trans, num_workers=args.num_workers,
                                      batch_size=256, shuffle=False, pin_memory=False)
    train_labelled_test_trans_loader = DataLoader(train_labelled_test_trans, num_workers=args.num_workers,
                                      batch_size=256, shuffle=False, pin_memory=False)
    train_unlabelled_test_trans_loader = DataLoader(train_unlabelled_test_trans, num_workers=args.num_workers,
                                      batch_size=256, shuffle=False, pin_memory=False)

    # ----------------------
    # PROJECTION HEAD
    # ----------------------

    from model import vision_transformer as vits


    cl_head = vits.__dict__['DINOHead'](in_dim=args.feat_dim, out_dim=65536, nlayers=args.num_mlp_layers)
    ce_head = CE_Head(in_dim=args.feat_dim, out_dim=args.mlp_out_dim, nlayers=args.num_mlp_layers)

    cl_backbone = deepcopy(backbone)
    ce_backbone = deepcopy(backbone)

    # ----------------------
    # TRAIN
    # ----------------------

    ce_backbone = ce_backbone.to(device)
    ce_head = ce_head.to(device)
    cl_backbone = cl_backbone.to(device)
    cl_head = cl_head.to(device)
    
    train_dual(ce_backbone, ce_head, cl_backbone, cl_head, train_loader, test_loader_labelled, args)



model/bacon.py
JianhongBai/BaCon
import torch
import numpy as np
from sklearn.cluster import KMeans
from logits_adj_utils import compute_est_adjustment
# from cuml.cluster import KMeans as cuKMeans
import argparse

device = torch.device('cuda:0')


@torch.no_grad()
def dist_est(ce_backbone, cl_backbone, ce_head, train_loader_test_trans, args, debug=False):

    est_count = torch.ones(args.num_classes).cuda()
    with torch.no_grad():
        all_feats_cl = []
        targets = np.array([])
        mask_lab = np.array([])
        uq_idxs = np.array([])
        # First extract all features
        for batch_idx, batch in enumerate(train_loader_test_trans):
            images, class_labels, uq_idxs_, mask_lab_ = batch
            mask_lab_ = mask_lab_[:, 0]
            images = images.cuda(non_blocking=True)
            class_labels, mask_lab_ = class_labels.cuda(non_blocking=True), mask_lab_.cuda(
                non_blocking=True).bool()
            cl_backbone_feature = cl_backbone(images)

            uq_idxs = np.append(uq_idxs, uq_idxs_.cpu().numpy())
            targets = np.append(targets, class_labels.cpu().numpy())
            mask_lab = np.append(mask_lab, mask_lab_.cpu().bool().numpy())

            cl_feats = torch.nn.functional.normalize(cl_backbone_feature, dim=-1)
            all_feats_cl.append(cl_feats.detach().cpu().numpy())

        mask_lab = mask_lab.astype(bool)
        all_feats = np.concatenate(all_feats_cl)

        kmeans = KMeans(n_clusters=args.num_labeled_classes + args.num_unlabeled_classes,
                        random_state=0).fit(all_feats)
        all_preds = kmeans.labels_

        all_preds = torch.from_numpy(all_preds)
        _, ins_num = torch.unique(all_preds, return_counts=True)


    def align(y_pred, y_true):
        assert y_pred.size == y_true.size
        # D = max(y_pred.max(), y_true.max()) + 1
        D = args.num_classes
        w = np.zeros((D, D), dtype=float)
        for i in range(y_pred.size):
            w[y_pred[i], y_true[i]] += 1
        w /= ins_num[:D].view(-1, 1).repeat(1, D).numpy()

        from scipy.optimize import linear_sum_assignment as linear_assignment
        ind = linear_assignment(w.max() - w)
        y_pred_id, y_true_id = ind
        return torch.from_numpy(y_pred_id), torch.from_numpy(y_true_id)

    labeled_preds = all_preds[mask_lab]
    labeled_targets = targets[mask_lab]

    y_pred_id, y_true_id = align(y_pred=labeled_preds.cpu().numpy().astype(int),
                                                    y_true=labeled_targets.astype(int))
    _, idx = y_true_id.sort()
    assert torch.all(idx == y_pred_id[idx])
    # _, idx = labeled_targets.sort()[1]
    labeled_targets = torch.from_numpy(labeled_targets)
    _, known_labeled_ins_num = torch.unique(labeled_targets, return_counts=True)

    idx1 = torch.argsort(known_labeled_ins_num)
    known_cluster_ins_num = ins_num[idx[:args.num_labeled_classes]]
    known_cluster_ins_num = known_cluster_ins_num.to(est_count.dtype).cuda()
    idx2 = torch.argsort(known_cluster_ins_num)

    est_count[:args.num_labeled_classes][idx1] = known_cluster_ins_num[idx2]
    est_count[args.num_labeled_classes:] = ins_num[idx[args.num_labeled_classes:]].sort(descending=True)[0]

    args.est_dist = est_count.detach().cuda()
    args.est_adjustment = compute_est_adjustment(tro=args.tro, args=args)
    return est_count.detach().cuda()

model/dist_est.py
JianhongBai/BaCon
import random

import torch
import torch.nn as nn
import numpy as np
from torch.nn import functional as F
from torch.nn.modules.loss import _Loss
import json


class ContrastiveLearningViewGenerator(object):
    """Take two random crops of one image as the query and key."""

    def __init__(self, base_transform, n_views=2):
        self.base_transform = base_transform
        self.n_views = n_views

    def __call__(self, x):
        if not isinstance(self.base_transform, list):
            return [self.base_transform(x) for i in range(self.n_views)]
        else:
            return [self.base_transform[i](x) for i in range(self.n_views)]


class SupConLoss(torch.nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR
    From: https://github.com/HobbitLong/SupContrast"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)

        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss


class SemiConLoss(torch.nn.Module):
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07, args=None):
        super(SemiConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature
        self.args = args

    def forward(self, features, soft_mask=None, ambiguous_mask=None):
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)

        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        # mask = mask.repeat(anchor_count, contrast_count)
        if not len(logits) == len(soft_mask):
            soft_mask = soft_mask.repeat(anchor_count, contrast_count)

        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(soft_mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )

        soft_mask = soft_mask * logits_mask
        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        mean_log_prob_pos = (soft_mask * log_prob).sum(1) / soft_mask.sum(1)
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()


        return loss

def info_nce_logits(features, n_views=2, temperature=1.0, device='cuda'):

    b_ = 0.5 * int(features.size(0))

    labels = torch.cat([torch.arange(b_) for i in range(n_views)], dim=0)
    labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()
    labels = labels.to(device)

    features = F.normalize(features, dim=1)

    similarity_matrix = torch.matmul(features, features.T)

    # discard the main diagonal from both: labels and similarities matrix
    mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)
    labels = labels[~mask].view(labels.shape[0], -1)
    similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)

    # select and combine multiple positives
    positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)

    # select only the negatives the negatives
    negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)

    logits = torch.cat([positives, negatives], dim=1)
    labels = torch.zeros(logits.shape[0], dtype=torch.long).to(device)

    logits = logits / temperature
    return logits, labels


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.}]


class DistillLoss(nn.Module):
    def __init__(self, warmup_teacher_temp_epochs, nepochs,
                 ncrops=2, warmup_teacher_temp=0.07, teacher_temp=0.04,
                 student_temp=0.1, args=None):
        super().__init__()
        self.args = args
        self.student_temp = student_temp
        self.ncrops = ncrops
        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 forward(self, student_output, teacher_output, epoch):
        """
        Cross-entropy between softmax outputs of the teacher and student networks.
        """
        student_out = student_output / self.student_temp
        student_out = student_out.chunk(self.ncrops)
        # teacher centering and sharpening
        temp = self.teacher_temp_schedule[epoch]

        teacher_out = F.softmax(teacher_output / temp, dim=-1)
        teacher_out = teacher_out.detach().chunk(2)

        total_loss = 0
        n_loss_terms = 0

        for iq, q in enumerate(teacher_out):
            for v in range(len(student_out)):
                if v == iq:
                    # we skip cases where student and teacher operate on the same view
                    continue
                loss = torch.sum(-q * F.log_softmax(student_out[v], dim=-1), dim=-1)

                total_loss += loss.mean()
                n_loss_terms += 1
        total_loss /= n_loss_terms
        return total_loss

model/loss.py
JianhongBai/BaCon
import torch
import numpy as np
import math
import torch.nn.functional as F


def compute_reg_loss(student_out, est_count, args):

    avg_probs = (student_out / 0.1).softmax(dim=1).mean(dim=0)
    avg_probs = avg_probs * 1 / (est_count) ** args.p
    avg_probs /= avg_probs.sum()
    me_max_loss = - torch.sum(torch.log(avg_probs ** (-avg_probs))) + math.log(float(len(avg_probs)))
    return me_max_loss

model/reg_loss.py
JianhongBai/BaCon
import torch
from model.loss import SemiConLoss
sample_count = torch.zeros(100).cuda()


def compute_softconloss(student_out, cl_proj_feature, sup_con_labels, sup_cl_proj_feature, mask_lab, args):

    logits = (student_out / 0.1) - args.est_adjustment
    view1_logits, view2_logits = logits.softmax(dim=1).chunk(2)
    soft_labels = (view1_logits + view2_logits) / 2
    
    known_class_sampling_rate = ((1 / args.est_dist) * args.est_dist.min()) ** (args.alpha)
    existing_class_idx, _ = torch.unique(sup_con_labels, return_counts=True)
    sampling_rate = ((1 / args.est_dist) * args.est_dist.min()) ** (args.beta)
    sampling_rate[existing_class_idx] = known_class_sampling_rate[existing_class_idx]
    batch_confidence, batch_preds = soft_labels[~mask_lab].max(dim=1)

    sampling_mask = torch.zeros(len(soft_labels)).bool().cuda()
    mask = torch.arange(len(soft_labels)).cuda()

    for cls_idx in torch.unique(batch_preds):
        cls_ins_num = (batch_preds == cls_idx).sum()
        cls_sample_num = torch.bernoulli(torch.tensor([sampling_rate[cls_idx]] * cls_ins_num)).sum().int()
        cls_conf = batch_confidence[batch_preds == cls_idx]
        val, idx = torch.topk(cls_conf, k=cls_sample_num)
        sampling_mask[mask[~mask_lab][batch_preds == cls_idx][idx]] = True

    semicon_soft_labels = soft_labels[(~mask_lab) & sampling_mask]
    semicon_feats = torch.cat([f[(~mask_lab) & sampling_mask].unsqueeze(1) for f in cl_proj_feature.chunk(2)], dim=1)

    sup_labels_one_hot = soft_labels[mask_lab]
    sup_cl_proj_feature = torch.cat([sup_cl_proj_feature, semicon_feats], dim=0)
    sup_con_labels = torch.cat([sup_labels_one_hot, semicon_soft_labels], dim=0)
    soft_mask = (sup_con_labels.unsqueeze(1) * sup_con_labels.unsqueeze(0)).sum(dim=2)
    soft_mask[torch.eye(len(soft_mask), device=soft_mask.device).bool()] = 1

    sup_con_loss = SemiConLoss(args=args)(sup_cl_proj_feature, soft_mask=soft_mask.detach())
    return sup_con_loss

model/softconloss.py
JianhongBai/BaCon
# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.
"""
Mostly copy-paste from timm library.
https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
"""
import math
from functools import partial

import torch
import torch.nn as nn

from torch.nn.init import trunc_normal_

def drop_path(x, drop_prob: float = 0., training: bool = False):
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x, attn


class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, return_attention=False):
        y, attn = self.attn(self.norm1(x))
        x = x + self.drop_path(y)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        if return_attention:
            return x, attn
        else:
            return x


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        num_patches = (img_size // patch_size) * (img_size // patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x


class VisionTransformer(nn.Module):
    """ Vision Transformer """
    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
                 drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
        super().__init__()
        self.num_features = self.embed_dim = embed_dim

        self.patch_embed = PatchEmbed(
            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)

        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        self.apply(self._init_weights)

    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)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)

    def prepare_tokens(self, x):
        B, nc, w, h = x.shape
        x = self.patch_embed(x)  # patch linear embedding

        # add the [CLS] token to the embed patch tokens
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # add positional encoding to each token
        x = x + self.interpolate_pos_encoding(x, w, h)

        return self.pos_drop(x)

    def forward(self, x, return_all_patches=False):
        x = self.prepare_tokens(x)
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)

        if return_all_patches:
            return x
        else:
            return x[:, 0]

    def get_last_selfattention(self, x):
        x = self.prepare_tokens(x)
        for i, blk in enumerate(self.blocks):
            if i < len(self.blocks) - 1:
                x = blk(x)
            else:
                # return attention of the last block
                x, attn = blk(x, return_attention=True)
                x = self.norm(x)
                return x, attn

    def get_intermediate_layers(self, x, n=1):
        x = self.prepare_tokens(x)
        # we return the output tokens from the `n` last blocks
        output = []
        for i, blk in enumerate(self.blocks):
            x = blk(x)
            if len(self.blocks) - i <= n:
                output.append(self.norm(x))
        return output


def vit_tiny(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vit_small(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


def vit_base(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
    return model


class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim, use_bn=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)
        else:
            layers = [nn.Linear(in_dim, hidden_dim)]
            if use_bn:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            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())
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
            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 VisionTransformerWithLinear(nn.Module):

    def __init__(self, base_vit, num_classes=200):

        super().__init__()

        self.base_vit = base_vit
        self.fc = nn.Linear(768, num_classes)

    def forward(self, x, return_features=False):

        features = self.base_vit(x)
        features = torch.nn.functional.normalize(features, dim=-1)
        logits = self.fc(features)

        if return_features:
            return logits, features
        else:
            return logits

    @torch.no_grad()
    def normalize_prototypes(self):
        w = self.fc.weight.data.clone()
        w = torch.nn.functional.normalize(w, dim=1, p=2)
        self.fc.weight.copy_(w)

model/vision_transformer.py
JianhongBai/BaCon
import numpy as np
from scipy.optimize import linear_sum_assignment as linear_assignment
import torch

def set_args_mmf(args, train_loader):
    train_loader_targets = train_loader.dataset.labelled_dataset.targets + train_loader.dataset.unlabelled_dataset.datasets[0].targets + train_loader.dataset.unlabelled_dataset.datasets[1].targets

    labeled_num = len(torch.unique(torch.tensor(train_loader.dataset.labelled_dataset.targets)))
    print(f'Labeled Classes Number: {labeled_num}')
    print(f'Labeled Set Length: {len(train_loader.dataset.labelled_dataset.targets)}')

    print(f'Unlabeled Set Length: {len(train_loader_targets)}')

    train_loader_targets = torch.tensor(train_loader_targets)
    known_train_loader_targets = train_loader_targets[train_loader_targets < args.num_labeled_classes]
    unknown_train_loader_targets = train_loader_targets[train_loader_targets >= args.num_labeled_classes]

    k_cls_idx, ins_num = torch.unique(known_train_loader_targets, return_counts=True)

    val, idx = torch.sort(ins_num, descending=True)

    if len(k_cls_idx) == 5:  # CIFAR-10
        many_thre = val[1]
        few_thre = val[2]
    else:
        many_thre = val[int(1 / 3 * len(k_cls_idx))]
        few_thre = val[int(2 / 3 * len(k_cls_idx))]

    known_many_cls = k_cls_idx[ins_num > many_thre]
    known_few_cls = k_cls_idx[ins_num < few_thre]
    known_median_cls = k_cls_idx[(ins_num <= many_thre) & (ins_num >= few_thre)]
    assert len(k_cls_idx) == (len(known_many_cls) + len(known_few_cls) + len(known_median_cls))

    uk_cls_idx, ins_num = torch.unique(unknown_train_loader_targets, return_counts=True)
    val, idx = torch.sort(ins_num, descending=True)

    if len(k_cls_idx) == 5:  # CIFAR-10
        many_thre = val[1]
        few_thre = val[2]
    else:
        many_thre = val[int(1 / 3 * len(uk_cls_idx))]
        few_thre = val[int(2 / 3 * len(uk_cls_idx))]

    unknown_many_cls = uk_cls_idx[ins_num > many_thre]
    unknown_few_cls = uk_cls_idx[ins_num < few_thre]
    unknown_median_cls = uk_cls_idx[(ins_num <= many_thre) & (ins_num >= few_thre)]

    args.known_many_cls = known_many_cls
    args.known_median_cls = known_median_cls
    args.known_few_cls = known_few_cls
    args.unknown_many_cls = unknown_many_cls
    args.unknown_median_cls = unknown_median_cls
    args.unknown_few_cls = unknown_few_cls
    args.k_cls_idx = k_cls_idx
    print(
        f'Known Many: {len(known_many_cls)}  Known Med: {len(known_median_cls)}  Known Few: {len(known_few_cls)} \n Novel Many: {len(unknown_many_cls)}  Novel Med: {len(unknown_median_cls)}  Novel Few: {len(unknown_few_cls)}')


def calculate_acc(ind_map, w, list_targets):
    acc_list = []
    for targets in list_targets:
        if targets is None:
            acc_list.append(-1.)

            continue
        targets = targets.numpy().tolist()
        acc = 0
        total_instances = 0

        for i in targets:
            acc += w[ind_map[i], i]
            total_instances += sum(w[:, i])

        if total_instances != 0:
            acc /= total_instances
            acc_list.append(acc * 100)
        else:
            acc_list.append(-1.)

    return acc_list


def split_cluster_acc_v2(y_true, y_pred, mask, train_loader, args):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`

    # Return
        accuracy, in [0,1]
    """

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    total_acc = sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size
    total_acc *= 100

    if y_true.max() == args.k_cls_idx.max().item():
        args.unknown_many_cls = args.unknown_median_cls = args.unknown_few_cls = None
    acc_list = calculate_acc(ind_map=ind_map, w=w, list_targets=[args.known_many_cls, args.known_median_cls, args.known_few_cls,
                                                      args.unknown_many_cls, args.unknown_median_cls, args.unknown_few_cls])

    old_acc = 0
    total_old_instances = 0
    total_pred_old_instances = 0
    for i in old_classes_gt:
        old_acc += w[ind_map[i], i]
        total_old_instances += sum(w[:, i])
        total_pred_old_instances += sum(w[ind_map[i], :])

    if total_old_instances != 0:
        old_acc /= total_old_instances
        old_acc *= 100
    else:
        old_acc = -1.

    new_acc = 0
    total_new_instances = 0
    total_pred_new_instances = 0
    for i in new_classes_gt:
        new_acc += w[ind_map[i], i]
        total_new_instances += sum(w[:, i])
        total_pred_new_instances += sum(w[ind_map[i], :])

    if total_new_instances != 0:
        new_acc /= total_new_instances
        new_acc *= 100
    else:
        new_acc = -1.

    return total_acc, old_acc, new_acc, acc_list, ind_map


def split_cluster_acc_v2_balanced(y_true, y_pred, mask):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}

    old_acc = np.zeros(len(old_classes_gt))
    total_old_instances = np.zeros(len(old_classes_gt))
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(len(new_classes_gt))
    total_new_instances = np.zeros(len(new_classes_gt))
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_acc /= total_old_instances
    new_acc /= total_new_instances
    total_acc, old_acc, new_acc = total_acc.mean(), old_acc.mean(), new_acc.mean()
    return total_acc, old_acc, new_acc


EVAL_FUNCS = {
    'v2': split_cluster_acc_v2,
    'v2b': split_cluster_acc_v2_balanced
}

def log_accs_from_preds(y_true, y_pred, mask, eval_funcs, save_name, T=None,
                        print_output=True, args=None, train_loader=None):

    """
    Given a list of evaluation functions to use (e.g ['v1', 'v2']) evaluate and log ACC results

    :param y_true: GT labels
    :param y_pred: Predicted indices
    :param mask: Which instances belong to Old and New classes
    :param T: Epoch
    :param eval_funcs: Which evaluation functions to use
    :param save_name: What are we evaluating ACC on
    :param writer: Tensorboard logger
    :return:
    """

    mask = mask.astype(bool)
    y_true = y_true.astype(int)
    y_pred = y_pred.astype(int)

    for i, f_name in enumerate(eval_funcs):

        acc_f = EVAL_FUNCS[f_name]
        all_acc, old_acc, new_acc, acc_list, ind_map = acc_f(y_true, y_pred, mask, train_loader, args)
        log_name = f'{save_name}_{f_name}'

        if i == 0:
            to_return = (all_acc, old_acc, new_acc, acc_list, ind_map)

        if print_output:
            old_std = np.array([acc_list[0], acc_list[1], acc_list[2]]).std()
            new_std = np.array([acc_list[3], acc_list[4], acc_list[5]]).std()
            print_str = f'Epoch {T}, {log_name}: All {all_acc:.1f} | Old {old_acc:.1f} | New {new_acc:.1f} | '
            print_str2 = f'KMany {acc_list[0]:.1f} | KMed {acc_list[1]:.1f} | KFew {acc_list[2]:.1f} | Std {old_std:.1f}'
            print_str3 = f'UMany {acc_list[3]:.1f} | UMed {acc_list[4]:.1f} | UFew {acc_list[5]:.1f} | Std {new_std:.1f}'

            try:
                args.logger.info(print_str)
                args.logger.info(print_str2)
                args.logger.info(print_str3)
            except:
                print(print_str)
                print(print_str2)
                print(print_str3)
    return to_return

util/cluster_and_log_utils.py
JianhongBai/BaCon
import os
import torch
import inspect

from datetime import datetime
from loguru import logger
# from torch.utils.tensorboard import SummaryWriter

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):

        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):

        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def init_experiment(args, runner_name=None, exp_id=None):

    args.cuda = torch.cuda.is_available()

    # Get filepath of calling script
    if runner_name is None:
        runner_name = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))).split(".")[-2:]

    root_dir = os.path.join(args.exp_root, *runner_name)

    if not os.path.exists(root_dir):
        os.makedirs(root_dir)

    # Either generate a unique experiment ID, or use one which is passed
    if exp_id is None:

        if args.exp_name is None:
            raise ValueError("Need to specify the experiment name")
        # Unique identifier for experiment
        now = '{}_({:02d}.{:02d}.{}_|_'.format(args.exp_name, datetime.now().day, datetime.now().month, datetime.now().year) + \
              datetime.now().strftime("%S.%f")[:-3] + ')'

        log_dir = os.path.join(root_dir, 'log', now)
        while os.path.exists(log_dir):
            now = '({:02d}.{:02d}.{}_|_'.format(datetime.now().day, datetime.now().month, datetime.now().year) + \
                  datetime.now().strftime("%S.%f")[:-3] + ')'

            log_dir = os.path.join(root_dir, 'log', now)

    else:

        log_dir = os.path.join(root_dir, 'log', f'{exp_id}')

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
        
        
    logger.add(os.path.join(log_dir, 'log.txt'))
    args.logger = logger
    args.log_dir = log_dir

    # Instantiate directory to save models to
    model_root_dir = os.path.join(args.log_dir, 'checkpoints')
    if not os.path.exists(model_root_dir):
        os.mkdir(model_root_dir)

    args.model_dir = model_root_dir
    args.model_path = os.path.join(args.model_dir, 'model.pt')

    print(f'Experiment saved to: {args.log_dir}')

    args.writer = None # SummaryWriter(log_dir=args.log_dir)

    hparam_dict = {}

    for k, v in vars(args).items():
        if isinstance(v, (int, float, str, bool, torch.Tensor)):
            hparam_dict[k] = v

    # args.writer.add_hparams(hparam_dict=hparam_dict, metric_dict={})
    # hparam_dict = {}
    #
    # for k, v in vars(args).items():
    #     if isinstance(v, (int, float, str, bool, torch.Tensor)):
    #         hparam_dict[k] = v

    print(runner_name)
    print(args)

    return args
util/general_utils.py
JianhongBai/BaCon