inference_60.py

# The testing module requires faiss
from functools import partial
import sys
# So if you don't have that, then this import will break
from pml import trainers
from vit_pytorch.swin import build_model
from pml import losses, miners, samplers, testers, utils
# import losss
import torch.nn as nn
import record_keeper
from utils import common_functions as c_f
import pml.utils.logging_presets as logging_presets
import pml
import pml as pytorch_metric_learning
from torchvision import datasets, models, transforms
import torchvision
import logging
import sklearn
logging.getLogger().setLevel(logging.INFO)
import os
# from pytorch_pretrained_vit import ViT
from pml.losses.base_metric_loss_function import BaseMetricLossFunction
from pml.testers.base_tester import BaseTester
from vit_pytorch.pvt import PyramidVisionTransformer
from vit_pytorch.CausalLevit import LeViT_384
import sys
logging.info("pytorch-metric-learning VERSION %s" % pytorch_metric_learning.__version__)
logging.info("record_keeper VERSION %s" % record_keeper.__version__)

from sklearn.metrics import accuracy_score
# from efficientnet_pytorch import EfficientNet
import torch
import numpy as np
import pickle

import hydra
from omegaconf import DictConfig

# reprodcibile
np.random.seed(42)
torch.manual_seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


class Loss_n(BaseMetricLossFunction):
    def __init__(self):
        super().__init__()
        self.n_pair_loss = losses.NPairsLoss()
        self.angular_loss = losses.AngularLoss(alpha=45)

    def compute_loss(self, embeddings, labels, indices_tuple):
        dict_angular = self.angular_loss.compute_loss(embeddings, labels, indices_tuple)
        dict_npair = self.n_pair_loss.compute_loss(embeddings, labels, indices_tuple)
        losses = 0.01 * dict_angular['loss']['losses'] + 0.1 * dict_npair['loss']['losses']
        dict_angular['loss']['losses'] = losses
        return dict_angular


def calibration_augmentation(base_means, base_cov, embedding_and_labels):
    n_shot = 1
    n_ways = 20
    support_data = np.nan_to_num(embedding_and_labels['samples'][0])
    support_label = embedding_and_labels['samples'][1]
    sampled_data = []
    sampled_label = []
    num_sampled = int(10 / n_shot)
    for i in range(20):
        mean, cov = distribution_calibration(support_data[np.squeeze(support_label == i), :], base_means, base_cov, k=4)
        sampled_data.append(np.random.multivariate_normal(mean=mean, cov=cov, size=num_sampled))
        sampled_label.extend([support_label[i]] * num_sampled)
    sampled_data = np.concatenate([sampled_data[:]]).reshape(n_ways * n_shot * num_sampled, -1)
    X_aug = np.concatenate([support_data, sampled_data])
    Y_aug = np.concatenate([support_label, sampled_label])
    return X_aug, Y_aug


def distribution_calibration(query, base_means, base_cov, k, alpha=0.21):
    dist = []
    for i in range(len(base_means)):
        dist.append(np.linalg.norm(query - base_means[i]))
    index = np.argpartition(dist, k)[:k]

    mean = np.concatenate([np.array(base_means)[index], query])
    calibrated_mean = np.mean(mean, axis=0)
    # print(base_cov)
    calibrated_cov = np.mean(np.array(base_cov)[index], axis=0) + alpha

    return calibrated_mean, calibrated_cov


class OneShotTester(BaseTester):

    def __init__(self, end_of_testing_hook=None):
        super().__init__()
        self.max_accuracy = 0.0
        self.embedding_filename = ""
        self.end_of_testing_hook = end_of_testing_hook

    def __get_correct(self, output, target, topk=(1,)):
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))
            #             print(correct)
        return correct

    def __accuracy(self, output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            correct = self.__get_correct(output, target, topk)
            batch_size = target.size(0)
            res = []
            for k in topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100.0 / batch_size))
            return res

    def do_knn_and_accuracies(self, accuracies, embeddings_and_labels, split_name, tag_suffix=''):
        # print(embeddings_and_labels)
        # train_embeddings = embeddings_and_labels['train'][0]
        # train_labels = embeddings_and_labels['train'][1]
        # print(train_embeddings.shape)
        # print(train_labels.shape)
        query_embeddings = embeddings_and_labels["val"][0]
        query_labels = embeddings_and_labels["val"][1]
        reference_embeddings_1 = embeddings_and_labels["samples"][0]
        reference_labels_1 = embeddings_and_labels["samples"][1]
        #print(reference_labels_1)
        reference_embeddings = np.zeros((12, 128))
        #print(reference_embeddings_1.shape)
        #sys.exit()
        for i in range(12):
            #mask = reference_labels_1 == i+100
            #mask = np.squeeze(mask)
            reference_embeddings[i,:] = reference_embeddings_1[5*i:5*i+5,:].mean(axis=0)
        reference_labels = np.arange(0,20)

        #reference_labels = embeddings_and_labels["samples"][1]
        knn_indices, knn_distances = utils.stat_utils.get_knn(reference_embeddings.astype('float32'), query_embeddings.astype('float32'), 1, False)
        knn_labels = reference_labels[knn_indices][:, 0]

        accuracy = accuracy_score(knn_labels, query_labels)
        f_1_score = sklearn.metrics.f1_score(query_labels, knn_labels, average='macro')
        precision = sklearn.metrics.precision_score(query_labels, knn_labels, average='macro')
        recall = sklearn.metrics.recall_score(query_labels, knn_labels, average='macro')
        logging.info('accuracy:{}'.format(accuracy))
        logging.info('f_1_score:{}'.format(f_1_score))
        logging.info('precision:{}'.format(precision))
        logging.info('recall:{}'.format(recall))
        
        query_embeddings = embeddings_and_labels["val"][0]
        query_labels = embeddings_and_labels["val"][1]
        reference_embeddings_1 = embeddings_and_labels["samples"][0]
        reference_labels_1 = embeddings_and_labels["samples"][1]
        # print(reference_labels_1)
        reference_embeddings = np.zeros((12, 128))
        # print(reference_embeddings_1.shape)
        # sys.exit()
        for i in range(12):
            # mask = reference_labels_1 == i+100
            # mask = np.squeeze(mask)
            reference_embeddings[i, :] = reference_embeddings_1[3 * i:3 * i + 1, :].mean(axis=0)
            query_embeddings = np.concatenate([query_embeddings, reference_embeddings_1[3 * i+1:3 * i + 3, :]], axis=0)
            query_labels = np.concatenate([query_labels,reference_labels_1[3 * i+1:3 * i + 3, :]], axis=0)
        reference_labels = np.arange(0, 12)


        # reference_labels = embeddings_and_labels["samples"][1]
        knn_indices, knn_distances = utils.stat_utils.get_knn(reference_embeddings.astype('float32'),
                                                              query_embeddings.astype('float32'), 1, False)
        knn_labels = reference_labels[knn_indices][:, 0]

        accuracy = accuracy_score(knn_labels, query_labels)
        f_1_score = sklearn.metrics.f1_score(query_labels, knn_labels, average='macro')
        precision = sklearn.metrics.precision_score(query_labels, knn_labels, average='macro')
        recall = sklearn.metrics.recall_score(query_labels, knn_labels, average='macro')
        logging.info('accuracy:{}'.format(accuracy))
        logging.info('f_1_score:{}'.format(f_1_score))
        logging.info('precision:{}'.format(precision))
        logging.info('recall:{}'.format(recall))
        #print(accuracy
        
    def do_knn_and_accuracies_aug(self, accuracies, embeddings_and_labels, split_name, tag_suffix=''):
        # print(embeddings_and_labels)
        print("test")
        train_embeddings = embeddings_and_labels['train'][0]
        train_labels = embeddings_and_labels['train'][1]
        # print(train_labels.shape)
        # print(train_embeddings.shape)
        base_means = []
        base_cov = []
        for key in range(100):
            feature = train_embeddings[np.squeeze(train_labels == key, axis=1), :]
            mean = np.mean(feature, axis=0)
            cov = np.cov(feature.T)
            base_means.append(mean)
            base_cov.append(cov)
            # print(cov.shape)
        x_aug, y_aug = calibration_augmentation(base_means, base_cov, embeddings_and_labels)
        # print(train_embeddings.shape)
        # print(train_labels.shape)
        reference_embedding = []
        reference_labels = []
        for key in range(20):
            reference_embedding.append(x_aug[np.squeeze(y_aug == key), :].mean(0))
            reference_labels.append(key)
        reference_embeddings = np.stack(reference_embedding, axis=0)
        reference_labels = np.stack(reference_labels, axis=0)
        query_embeddings = embeddings_and_labels["val"][0]
        query_labels = embeddings_and_labels["val"][1]
        # reference_embeddings = #embeddings_and_labels["samples"][0]
        # reference_labels = #embeddings_and_labels["samples"][1]
        # print(reference_embeddings.shape)
        # print(query_embeddings.shape)
        knn_indices, knn_distances = utils.stat_utils.get_knn(reference_embeddings.astype(np.float32), query_embeddings,
                                                              1, False)
        knn_labels = reference_labels[knn_indices][:, 0]
        accuracy = accuracy_score(knn_labels, query_labels)
        print(accuracy)
        accuracies["accuracy"] = accuracy
        keyname = self.accuracies_keyname("mean_average_precision_at_r")  # accuracy as keyname not working
        accuracies[keyname] = accuracy


class MLP(nn.Module):
    # layer_sizes[0] is the dimension of the input
    # layer_sizes[-1] is the dimension of the output
    def __init__(self, layer_sizes, final_relu=False):
        super().__init__()
        layer_list = []
        layer_sizes = [int(x) for x in layer_sizes]
        num_layers = len(layer_sizes) - 1
        final_relu_layer = num_layers if final_relu else num_layers - 1
        for i in range(len(layer_sizes) - 1):
            input_size = layer_sizes[i]
            curr_size = layer_sizes[i + 1]
            if i < final_relu_layer:
                layer_list.append(nn.ReLU(inplace=True))
            layer_list.append(nn.Linear(input_size, curr_size))
        self.net = nn.Sequential(*layer_list)
        self.last_linear = self.net[-1]

    def forward(self, x):
        out = self.net(x)
        # print(out.size())
        return out


class SignalModule(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # self.trunk_signal = torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
        # ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16, mlp_dim=21,dropout=0.1,emb_dropout=0.1)

        self.trunk_signal = LeViT_384(num_classes=512, distillation=True,
                                      pretrained=False, fuse=False)
        self.trunk_fft = LeViT_384(num_classes=512, distillation=True,
                                   pretrained=False, fuse=False, in_chans=6)

        """PyramidVisionTransformer(
            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
            norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
        )
        self.trunk_fft = PyramidVisionTransformer(
            patch_size=4, in_chans=6, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
            norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
        )"""

        """
        self.trunk_fft = LeViT(
            image_size = 256,
            num_classes = 20,
            stages = 3,             # number of stages
            dim = (256, 384, 512),  # dimensions at each stage
            depth = 4,              # transformer of depth 4 at each stage
            heads = (4, 6, 8),      # heads at each stage
            mlp_mult = 2,
            dropout = 0.1
         )
        """

        # self.trunk_signal.fc = Identity()
        # self.trunk_fft.fc =Identity()
        # ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16,channels=6, mlp_dim=21,dropout=0.1,emb_dropout=0.1)
        # self.conv_fusion = nn.Sequential(
        #    nn.Conv2d(1024, 512, kernel_size=3, padding=1, bias=False),
        #    nn.BatchNorm2d(1024, eps=1e-3, momentum=0.01),
        #    nn.ReLU()
        # )
        self.MLP_2 = MLP([1024, 512])
        # self.MLP = MLP([6,3])

    def forward(self, x):
        # print(x.size())
        batch_size = x.size()[0]
        c_1 = self.trunk_signal(x)

        data_fft = torch.rfft(x.permute(0, 1, 3, 2), signal_ndim=1)
        # print(data_fft[:,:,:,:,1])
        # sys.exit()
        real = data_fft[:, :, :, :, 0].permute(0, 1, 3, 2)
        imag = data_fft[:, :, :, :, 1].permute(0, 1, 3, 2)
        norm = torch.nn.functional.normalize(torch.sqrt(torch.pow(real, 2) + torch.pow(imag, 2)), dim=-2)
        angle = torch.nn.functional.normalize(torch.atan2(real, imag), dim=-2)
        data_fusion = torch.cat([norm, angle], dim=1)
        # print(data_fusion.size())
        # sys.exit()
        container = torch.zeros([batch_size, 6, 256, 256]).cuda()
        container[:, :, :129, :] = data_fusion[:, :, :, :]
        container[:, :, 129:, :] = torch.flip(container, dims=[2])[:, :, :127, :]
        c_2 = self.trunk_fft(container)
        c_2 = self.trunk_signal(x)
        # print(c_1.size(), c2.size())
        fusion = self.MLP_2(torch.cat([c_1, c_2], dim=1))

        return fusion


# This is for replacing the last layer of a pretrained network.
# This code is from https://github.com/KevinMusgrave/powerful_benchmarker
class Identity(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


def get_datasets(data_dir, cfg, mode="train"):
    common_transforms = []
    train_transforms = []
    test_transforms = []
    # if cfg.transform.transform_resize_match:
    common_transforms.append(transforms.Resize((256, 256)))

    if cfg.transform.transform_random_resized_crop:
        train_transforms.append(transforms.RandomResizedCrop(cfg.transform.transform_resize))
    if cfg.transform.transform_random_horizontal_flip:
        train_transforms.append(torchvision.transforms.RandomHorizontalFlip(p=0.5))
    if cfg.transform.transform_random_rotation:
        train_transforms.append(
            transforms.RandomRotation(cfg.transform.transform_random_rotation_degrees))  # , fill=255))
    if cfg.transform.transform_random_shear:
        train_transforms.append(torchvision.transforms.RandomAffine(0,
                                                                    shear=(
                                                                        cfg.transform.transform_random_shear_x1,
                                                                        cfg.transform.transform_random_shear_x2,
                                                                        cfg.transform.transform_random_shear_y1,
                                                                        cfg.transform.transform_random_shear_y2
                                                                    ),
                                                                    fillcolor=255))
    if cfg.transform.transform_random_perspective:
        train_transforms.append(transforms.RandomPerspective(distortion_scale=cfg.transform.transform_perspective_scale,
                                                             p=0.5,
                                                             interpolation=3)
                                )
    if cfg.transform.transform_random_affine:
        train_transforms.append(transforms.RandomAffine(degrees=(cfg.transform.transform_degrees_min,
                                                                 cfg.transform.transform_degrees_max),
                                                        translate=(cfg.transform.transform_translate_a,
                                                                   cfg.transform.transform_translate_b),
                                                        fillcolor=255))
    data_transforms = {
        'train': transforms.Compose(common_transforms + train_transforms + [transforms.ToTensor()]),
        'test': transforms.Compose(common_transforms + [transforms.ToTensor()]),
    }

    #train_dataset = datasets.ImageFolder(os.path.join(data_dir, "train"),
    #                                    data_transforms["train"])

    # for the final model we can join train, validation, validation samples datasets
    print(mode)
    if mode == "final_train":
        # train_dataset = torch.utils.data.ConcatDataset([train_dataset,
        #        val_dataset,
        #        val_samples_dataset])

        test_dataset = datasets.ImageFolder(os.path.join(data_dir, "test"),
                                            data_transforms["test"])

        samples_dataset = datasets.ImageFolder(os.path.join(data_dir, "samples"),
                                               data_transforms["test"])
        return test_dataset, samples_dataset
    else:
        if mode == "train":
            val_dataset = datasets.ImageFolder(os.path.join(data_dir, "val"),
                                               data_transforms["test"])

            val_samples_dataset = datasets.ImageFolder(os.path.join(data_dir, "val_samples"),
                                                       data_transforms["test"])
            return val_dataset, val_samples_dataset

        if mode == "test":
            return test_dataset, samples_dataset


@hydra.main(config_path="config/config.yaml")
def train_app(cfg):
    print(cfg.pretty())

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # trunk = PyramidVisionTransformer(
    #    patch_size=4, embed_dims=[128, 256, 512, 768], num_heads=[2, 4, 8, 12], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
    #    norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 10, 60, 3], sr_ratios=[8, 4, 2, 1],
    #     )
    # trunk = PyramidVisionTransformer(
    #    patch_size=32, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
    #    norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
    #    )
    trunk = LeViT_384(num_classes=512, distillation=True,
                      pretrained=False, fuse=False)

    # trunk = LeViT_384_cascaded(num_classes=512, distillation=True,
    #          pretrained=False, fuse=False)
    # device=torch.device("cpu")
    # Set trunk model and replace the softmax layer with an identity function
    #trunk = torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
    # trunk = SignalModule(cfg)
    # trunk=ViT(image_size=256,patch_size=64, num_classes=21, dim=512,depth=6, heads=16, mlp_dim=21,dropout=0.1,e

    # trunk= torchvision.models.__dict__[cfg.model.model_name](pretrained=cfg.model.pretrained)
    # trunk = models.resnet18(pretrained=False)
    # trunk = models.alexnet(pretrained=True)
    # trunk = models.resnet50(pretrained=True)
    # trunk = models.resnet152(pretrained=True)
    # trunk = models.wide_resnet50_2(pretrained=True)
    # trunk = EfficientNet.from_pretrained('efficientnet-b2')
    # trunk = ViT('B_16_imagenet1k', pretrained=True)
    #trunk = build_model()
    trunk_output_size = 768 #trunk.fc.in_features
    trunk.head = Identity()
    trunk.head_2 = Identity()
    path ='/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-25/14-41-33/example_saved_models_c_2_cat/levit60_rep_twostage__proPMBFA_re_nturgbd_dataset_100_20_noise/' #levit60twostage
    #path = "/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-01/20-05-33/example_saved_models_c_2_cat/Swin_retrain_PMBFA/"
    #path_resnet = "/home/kpeng/oneshot_metriclearning/sl-dml/outputs/2021-09-27/11-38-48/example_saved_models/resnet18_baseline_one_shot_nturgbd120/"
    #path = '/home/kpeng/oneshot_metriclearning/models/output_transformer_dml/2021-09-21/13-34-59/example_saved_models_c_2_cat/model_c2_cat_levit_NTU_ONE_SHOT_SWAP_AXIS_model_resnet18_cl_cross_entropy_ml_triplet_margin_miner_multi_similarity_mix_ml_0.50_mix_cl_0.50_resize_256_emb_size_128_class_size_21_opt_rmsprop_lr_0.00_m_0.00_wd_0.00/'
    #path_levit = '/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-01/20-15-35/example_saved_models_c_2_cat/Levit_retrain_PMBFA/'
    #path_levit_pmbfa = '/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-09/07-52-58/example_saved_models_c_2_cat/levit_re_PMBFA_nturgbd60_48_12/'
    #path_ = '/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-08/17-09-26/example_saved_models_c_2_cat/Swin_re_PMBFA_nturgbd60_48_12/'

    #path = '/home/kpeng/oneshot_metriclearning/transformer-sl-dml/outputs/2021-10-09/20-45-17/example_saved_models_c_2_cat/swin_PMBFA_nturgbd60_48_12/'
    embedder = MLP([trunk_output_size, cfg.embedder.size])
    embedder.load_state_dict(torch.load(path+'embedder_best23.pth'))
    embedder = torch.nn.DataParallel(embedder.to(device)).eval()
    classifier = MLP([cfg.embedder.size, 49])
    classifier.load_state_dict(torch.load(path+'classifier_best23.pth'))
    classifier = torch.nn.DataParallel(classifier.to(device)).eval()
    trunk.load_state_dict(torch.load(path + 'trunk_best23.pth'))
    trunk = torch.nn.DataParallel(trunk.to(device)).eval()
    # Set the datasets
    data_dir = os.environ["DATASET_FOLDER"] + "/" + cfg.dataset.data_dir
    print("Data dir: " + data_dir)

    val_dataset, val_samples_dataset = get_datasets(data_dir, cfg, mode=cfg.mode.type)
    print("Testset: ", len(val_dataset), "Samplesset: ", len(val_samples_dataset))


    batch_size = cfg.trainer.batch_size
    num_epochs = cfg.trainer.num_epochs


    # Package the above stuff into dictionaries.
    models = {"trunk": trunk, "embedder": embedder, "classifier": classifier}
    dataset_dict = {"samples": val_samples_dataset, "val": val_dataset}

    # We can specify loss weights if we want to. This is optional

    experiment_name = 'test'
    hooks = logging_presets.get_hook_container(record_keeper)
    dataset_dict = {"samples": val_samples_dataset, "val": val_dataset}
    model_folder = "example_saved_models_c__cat/%s/" % (experiment_name)
    # Create the tester
    tester = OneShotTester(
        end_of_testing_hook=hooks.end_of_testing_hook,

    )
    tester.test(dataset_dict=dataset_dict,epoch=1,trunk_model=trunk, embedder_model=embedder, classifier_model=classifier, splits_to_eval=['samples','val'])


if __name__ == "__main__":
    train_app()