robotics.py

# Modified from https://github.com/bearpaw/pytorch-classification
from __future__ import print_function

import argparse
import os
import shutil
import time
import random
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data as data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import models.imagenet as customized_models

import numpy as np
import h5py

from utils import Bar, Logger, AverageMeter, accuracy, mkdir_p, savefig

# Models
default_model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

customized_models_names = sorted(name for name in customized_models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(customized_models.__dict__[name]))

for name in customized_models.__dict__:
    if name.islower() and not name.startswith("__") and callable(customized_models.__dict__[name]):
        models.__dict__[name] = customized_models.__dict__[name]

model_names = default_model_names + customized_models_names

# Parse arguments
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')

# Datasets
parser.add_argument('--train', type=str, help='path to train h5 dataset')
parser.add_argument('--val', type=str, help='path to val h5 dataset')
parser.add_argument('--val-on-train', type=str, default=None, help='path to val_on_train h5 dataset')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
# Optimization options
parser.add_argument('--epochs', default=90, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--repeat-epochs', default=1, type=int, metavar='N',
                    help='repeat training batch in the same epoch')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--train-batch', default=256, type=int, metavar='N',
                    help='train batchsize (default: 256)')
parser.add_argument('--test-batch', default=200, type=int, metavar='N',
                    help='test batchsize (default: 200)')
parser.add_argument('--skip-val', action='store_true', help='skip validation during training')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--onecyclelr', action='store_true', help='use onecyclelr')
parser.add_argument('--drop', '--dropout', default=0, type=float,
                    metavar='Dropout', help='Dropout ratio')
parser.add_argument('--schedule', type=int, nargs='+', default=[30, 60],
                        help='Decrease learning rate at these epochs.')
parser.add_argument('--gamma', type=float, default=0.1, help='LR is multiplied by gamma on schedule.')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
# Checkpoints
parser.add_argument('-c', '--checkpoint', default='checkpoint', type=str, metavar='PATH',
                    help='path to save checkpoint (default: checkpoint)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
# Architecture
parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
                    choices=model_names,
                    help='model architecture: ' +
                        ' | '.join(model_names) +
                        ' (default: resnet18)')
# Openset-specific
parser.add_argument('--use-clip', action='store_true', help='whether to use CLIP model')
parser.add_argument('--clip-repo', type=str, default='clip', choices=['clip', 'open_clip'])
parser.add_argument('--clip-model', type=str, default='ViT-L/14')
parser.add_argument('--clip-dataset', type=str, default='openai', choices=['openai', 'laion400m_e31', 'laion400m_e32', 'laion2b_s32b_b82k'])
parser.add_argument('--clip-align-image-classification', type=int, default=1, help="whether to use L-cls (vision-language alignment loss for classification) (0/1)")
parser.add_argument('--chatgpt-raw-text-file', type=str, default=None, help="file containing label descriptions generated by chatgpt")
parser.add_argument('--chatgpt-neg-text-file', type=str, default=None, help="file containing negative label descriptions generated by chatgpt")
parser.add_argument('--clip-align-image-mse', action='store_true', help="Whether to use L-mse")
parser.add_argument('--clip-align-image-mse-unnorm', action='store_true', help="Whether to use L-mse (unnormalized version)")
parser.add_argument('--clip-align-image-contrastive', action='store_true', help="Whether to use L-im-cst")
parser.add_argument('--clip-align-image-contrastive-mode', type=str, default='bidirectional', choices=['single', 'bidirectional'])
parser.add_argument('--temperature', type=float, default=0.07, help="Temperature")
parser.add_argument('--few-shot-num', type=int, default=0, help='number of few-shot examples')
parser.add_argument('--few-shot-method', type=str, default='None', help='few-shot mode, None or finetune')
parser.add_argument('--local-aggregation', action='store_true', help="Perform local aggregation following DualCoOp")
parser.add_argument('--prompt-learner', action='store_true', help='whether to use prompt learner (dual prompts, CoOp)')
parser.add_argument('--prompt-learner-nctx', type=int, default=8, help='number of CoOp context tokens')
# Miscs
parser.add_argument('--manualSeed', type=int, help='Manual seed')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='#valuate model on validation set')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='Use pre-trained model')
parser.add_argument('--advprop', default=False, action='store_true',
                    help='Use advprop or not')
parser.add_argument('--use-adam', action='store_true', help='Whether to use adam optimizer')
#Device options
parser.add_argument('--gpu-id', default='0', type=str,
                    help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--clip-gpu-id', default='0', type=str,
                    help='id(s) for CUDA_VISIBLE_DEVICES for the CLIP model')

args = parser.parse_args()
args.clip_align_image_classification = bool(args.clip_align_image_classification)
print("clip_align_image_classification = ", args.clip_align_image_classification)
state = {k: v for k, v in args._get_kwargs()}

# Use CUDA
use_cuda = torch.cuda.is_available()
assert use_cuda

# Random seed
if args.manualSeed is None:
    args.manualSeed = random.randint(1, 10000)
random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
if use_cuda:
    cuda_device = f"cuda:{args.gpu_id}"
    torch.cuda.manual_seed_all(args.manualSeed)

best_acc = 0  # best test accuracy

def main():
    global best_acc
    start_epoch = args.start_epoch  # start from epoch 0 or last checkpoint epoch

    if not os.path.isdir(args.checkpoint):
        mkdir_p(args.checkpoint)

    # Build CLIP teacher model
    if args.use_clip:
        clip_device = f'cuda:{args.clip_gpu_id}'
        print("clip_device", clip_device)
        if args.clip_repo == 'clip':
            import clip
            clip_model, clip_preprocess_orig = clip.load(args.clip_model, device=clip_device)
            if 'ViT' in args.clip_model or args.clip_model in ['RN50']:
                clip_preprocess = transforms.Compose([
                    transforms.Resize(size=224, interpolation=transforms.InterpolationMode.BICUBIC, max_size=None, antialias=None),
                    transforms.CenterCrop(size=(224, 224)),
                    transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
                ])           
            elif args.clip_model == 'RN50x16':
                clip_preprocess = transforms.Compose([
                    transforms.Resize(size=384, interpolation=transforms.InterpolationMode.BICUBIC, max_size=None, antialias=None),
                    transforms.CenterCrop(size=(384, 384)),
                    transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
                ])
            else:
                raise NotImplementedError()
        elif args.clip_repo == 'open_clip':
            import open_clip
            clip_model, _, clip_preprocess_orig = open_clip.create_model_and_transforms(args.clip_model, pretrained=args.clip_dataset)
            clip_preprocess = transforms.Compose([
                transforms.Resize(size=224, interpolation=transforms.InterpolationMode.BICUBIC, max_size=None, antialias=None),
                transforms.CenterCrop(size=(224, 224)),
                transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
            ])       
        else:
            raise NotImplementedError(f"Unknown CLIP repo: {args.clip_repo}")
        print("clip_preprocess", clip_preprocess)
        clip_model.to(clip_device).eval()
        clip_model = clip_model.to(torch.float32)
        # Prevent CLIP model from updating during training
        for m in clip_model.parameters():
            m.requires_grad = False
    else:
        clip_model = clip_preprocess = clip_device = None
        
    # Data loading code
    train_h5 = args.train
    val_h5 = args.val
    val_on_train_h5 = args.val_on_train
    
    if args.advprop:
        normalize = transforms.Lambda(lambda img: img * 2.0 - 1.0)
    else:
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                        std=[0.229, 0.224, 0.225])
    # x = [base rgb, base_depth, hand_rgb, hand_depth]
    normalize_all = transforms.Lambda(lambda x: torch.cat([normalize(x[:, :3, :, :]), x[:, 3:4, :, :], normalize(x[:, 4:7, :, :]), x[:, 7:, :, :]], dim=1))

    # Build transformations
    if 'vit' not in args.arch:
        train_transform = transforms.Compose([
                transforms.RandomHorizontalFlip(),
                normalize_all,
            ])
        # to keep all instances in input image
        pre_train_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize([224, 224], interpolation=transforms.InterpolationMode.BICUBIC, max_size=None, antialias=None),
        ])
    else:
        raise NotImplementedError()
        
    # Dataset model ids, hardcoded for now
    train_model_ids = ['002_master_chef_can', '006_mustard_bottle', '011_banana',  
                 '013_apple', '015_peach', 
                '024_bowl', '030_fork', '032_knife',
                '035_power_drill', '037_scissors', 
                '048_hammer', '056_tennis_ball', '058_golf_ball', '062_dice']
    val_model_ids = ['003_cracker_box', '012_strawberry', 
                 '014_lemon', '016_pear', '017_orange',
                '021_bleach_cleanser', '026_sponge', '031_spoon',
                '036_wood_block', '043_phillips_screwdriver',
                '053_mini_soccer_ball', '077_rubiks_cube']
    all_model_ids = train_model_ids + val_model_ids
    
    # Build datasets
    from custom_data_loader import CLIPH5Dataset
    if not args.evaluate and args.few_shot_num > 0 and args.few_shot_method == 'finetune':
        from custom_data_loader import FSCLIPH5Dataset
        train_dataset = FSCLIPH5Dataset(train_h5, 
                                        val_h5,
                                        all_model_ids,
                                        transform=pre_train_transform,
                                        clip_model=clip_model,
                                        clip_preprocess=clip_preprocess,
                                        clip_device=clip_device,
                                        fs_num=args.few_shot_num,)
    else:
        train_dataset = CLIPH5Dataset(train_h5, 
                                      train_model_ids,
                                      transform=pre_train_transform,
                                      clip_model=clip_model,
                                      clip_preprocess=clip_preprocess,
                                      clip_device=clip_device,)
    
    if args.few_shot_num > 0 and args.few_shot_method == 'finetune':
        from custom_data_loader import FewShotSampler
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.train_batch,
            num_workers=args.workers, pin_memory=True,
            sampler=FewShotSampler(train_dataset, args.train_batch))
    else:
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.train_batch, shuffle=True,
            num_workers=args.workers, pin_memory=True)

    val_transform = transforms.Compose([
            normalize_all,
        ])
    val_dataset = CLIPH5Dataset(val_h5, 
                                all_model_ids,
                                transform=pre_train_transform,
                                clip_model=None)
    val_loader = torch.utils.data.DataLoader(
        val_dataset, 
        batch_size=args.test_batch, shuffle=False,
        num_workers=args.workers, pin_memory=True)
    
    val_on_train_dataset = None
    val_on_train_loader = None
    if val_on_train_h5 is not None:
        val_on_train_dataset = CLIPH5Dataset(
            val_on_train_h5, 
            train_model_ids,
            transform=pre_train_transform,
            clip_model=None)
        val_on_train_loader = torch.utils.data.DataLoader(
            val_on_train_dataset, 
            batch_size=args.test_batch, shuffle=False,
            num_workers=args.workers, pin_memory=True)

    # Create student model
    extra_args = dict()
    if args.clip_model == 'ViT-B/32':
        clip_feats_dim = 512
    elif args.clip_model in ['RN50']:
        clip_feats_dim = 1024
    else:
        clip_feats_dim = 768
    if args.use_clip:
        extra_args['fc_out_dim'] = clip_feats_dim # match clip feature dimension
    else:
        extra_args['fc_out_dim'] = 2 * len(all_model_ids) 
    if ('efficientnet' in args.arch or 'vit' in args.arch) and 'fc_out_dim' in extra_args:
        extra_args['num_classes'] = extra_args['fc_out_dim'] # "num_classes" here refers to feature dim; we are not doing regular cross entropy here
        extra_args.pop('fc_out_dim')
    print("Model creation extra args:", extra_args)
    
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        if 'efficientnet' in args.arch:
            extra_args['advprop'] = args.advprop
            # model = EfficientNet.from_pretrained(args.arch, advprop=args.advprop, **extra_args)
        model = models.__dict__[args.arch](pretrained=True, **extra_args)
    else:
        print("=> creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch](**extra_args)
        
    if 'efficientnet' in args.arch:
        model._modules['features'][0][0] = nn.Conv2d(8, 32, kernel_size=(3,3), stride=(2,2), padding=(1,1), bias=False)
        
        if args.local_aggregation:
            if not args.use_clip:
                local_proj_out_d = 2 * len(all_model_ids)
            else:
                local_proj_out_d = clip_feats_dim
            model.local_projection = nn.Conv2d(1280, local_proj_out_d, kernel_size=(1,1), stride=(1,1), bias=False)
            def local_aggregation_forward(model, x):
                out_local = model.features(x)
                out_global = model.classifier(model.avgpool(out_local).flatten(1))
                return model.local_projection(out_local), out_global
            model.forward = lambda x: local_aggregation_forward(model, x)
    else:
        raise NotImplementedError()
    
    # for multi-label classification
    model.global_logit_bias = 0.0 # nn.Parameter(torch.tensor(0.0))
    model.global_score_bias = nn.Parameter(torch.tensor(0.0), requires_grad=True)
    
    print("Final model:", model)
    model = model.to(cuda_device)

    cudnn.benchmark = True
    print('    Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
    
    # If using CLIP as teacher, process language descriptions and build prompt learner (if specified)
    chatgpt_lines = []
    if args.chatgpt_raw_text_file is not None:
        with open(args.chatgpt_raw_text_file, 'r') as f:
            for line in f:
                line = line.strip()
                if len(line) > 0:
                    chatgpt_lines.append(line)
    chatgpt_neg_lines = []
    if args.chatgpt_neg_text_file is not None:
        with open(args.chatgpt_neg_text_file, 'r') as f:
            for line in f:
                line = line.strip()
                if len(line) > 0:
                    chatgpt_neg_lines.append(line)
                        
    prompt_learner = text_encoder = None
    if args.use_clip:
        print("train class_to_idx", train_dataset.class_to_idx)
        print("test class_to_idx", val_dataset.class_to_idx)
        if args.chatgpt_raw_text_file is not None:
            extra_txt_fxn = lambda dset, x: chatgpt_lines[dset.class_to_idx[x]]
        else:
            extra_txt_fxn = lambda dset, x: ""
        train_text_labels = ["You can find a " + x + " in this image, except a robot arm. " + extra_txt_fxn(train_dataset, x) 
                             for x in train_dataset.class_to_idx.keys()]
        val_text_labels = ["You can find a " + x + " in this image, except a robot arm. " + extra_txt_fxn(val_dataset, x)
                            for x in val_dataset.class_to_idx.keys()]
        if val_on_train_dataset is not None:
            val_on_train_text_labels = ["You can find a " + x + " in this image, except a robot arm. " + extra_txt_fxn(val_on_train_dataset, x)
                                        for x in val_on_train_dataset.class_to_idx.keys()]
        else:
            val_on_train_text_labels = None
            
        if args.prompt_learner:
            from models.misc.prompt_learner import PromptLearner, TextEncoder
            if args.chatgpt_raw_text_file is not None:
                txt_fxn = lambda dset, x: x + " . " + chatgpt_lines[dset.class_to_idx[x]]
            else:
                txt_fxn = lambda dset, x: x
            if args.chatgpt_neg_text_file is not None:
                neg_txt_fxn = lambda dset, x: x + " . " + chatgpt_neg_lines[dset.class_to_idx[x]] + " . " + x + " . "
            else:
                neg_txt_fxn = lambda dset, x: x
            gen_txt_label_fxn = lambda dset: [txt_fxn(dset, x) for x in dset.class_to_idx.keys()] + [neg_txt_fxn(dset, x) for x in dset.class_to_idx.keys()]
            prompt_learner = PromptLearner(clip_model, gen_txt_label_fxn(train_dataset), \
                                            gen_txt_label_fxn(val_dataset), 
                                            gen_txt_label_fxn(val_on_train_dataset) if val_on_train_dataset is not None else None,
                                            device=cuda_device,
                                            mode="dual", cocoop=False,
                                            n_ctx=args.prompt_learner_nctx,
                                            use_neg_classname=True)
            text_encoder = TextEncoder(clip_model)
        else:
            prompt_learner = text_encoder = None
            
        print("train_text_labels", train_text_labels)
        print("val_text_labels", val_text_labels)
        
            
        if args.clip_repo == 'clip': 
            train_text_features = clip_model.encode_text(clip.tokenize(train_text_labels, truncate=True).to(clip_device)).float().detach()
            val_text_features = clip_model.encode_text(clip.tokenize(val_text_labels, truncate=True).to(clip_device)).float().detach()
            if val_on_train_text_labels is not None:
                val_on_train_text_features = clip_model.encode_text(clip.tokenize(val_on_train_text_labels, truncate=True).to(clip_device)).float().detach()
            else:
                val_on_train_text_features = None
        elif args.clip_repo == 'open_clip':
            tokenize = open_clip.tokenizer.tokenize
            train_text_features = clip_model.encode_text(tokenize(train_text_labels).to(clip_device)).float().detach()
            val_text_features = clip_model.encode_text(tokenize(val_text_labels).to(clip_device)).float().detach()
            if val_on_train_text_labels is not None:
                val_on_train_text_features = clip_model.encode_text(tokenize(val_on_train_text_labels).to(clip_device)).float().detach()
            else:
                val_on_train_text_features = None
    else:
        train_text_features = val_text_features = val_on_train_text_features = None
        
    if args.few_shot_method not in [None, 'None', 'finetune']:
        raise NotImplementedError()

    # L-cls for multi-label classification
    from misc_losses import AsymmetricLoss
    criterion_loss = AsymmetricLoss(gamma_neg=2, gamma_pos=1, clip=0.05, 
                                    disable_torch_grad_focal_loss=True, use_sigmoid=False)
    criterion = lambda inputs, targets: criterion_loss(inputs.softmax(dim=0)[0], targets, global_score_bias=torch.tanh(model.global_score_bias))
    
    # Define optimizer and scheduler
    params_to_optimize = list(model.parameters())
    if args.use_clip and prompt_learner is not None:
        for name, param in prompt_learner.named_parameters():
            if param.requires_grad:
                params_to_optimize.append(param)
                # print("Prompt learner param to optimize:", name)
    if not args.use_adam:
        optimizer = optim.SGD(params_to_optimize, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    else:
        optimizer = optim.AdamW(params_to_optimize, lr=args.lr, weight_decay=0.01) # hardcoded
    scheduler = None
    if args.onecyclelr:
        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=args.lr, steps_per_epoch=1, epochs=args.epochs)

    # Resume
    title = 'Robotics-' + args.arch
    if args.resume:
        # Load checkpoint.
        print(f'==> Resuming from checkpoint.. {args.resume}')
        assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'
        # args.checkpoint = os.path.dirname(args.resume)
        checkpoint = torch.load(args.resume, map_location=cuda_device)
        best_acc = checkpoint['best_acc']
        if args.few_shot_method != 'finetune':
            start_epoch = checkpoint['epoch']
        else:
            start_epoch = 0
        
        pth1 = Path(os.path.join(args.resume, 'checkpoint.pth.tar')).parent.absolute()
        pth2 = Path(os.path.join(args.checkpoint, 'log.txt')).parent.absolute()
        if pth1 == pth2:
            logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True)
        else:
            logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=False)
            logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss', 
                              'Train Acc.', 'Train Prec.', 'Train Recl.',
                              'Valid Acc.', 'Val Prec.', 'Val Recl.',
                              'Val On Train Acc.', 'Val On Train Prec.', 'Val On Train Recl.'])
        
        missing_keys, unexpected_keys = model.load_state_dict(checkpoint['state_dict'], strict=False)
        if len(missing_keys) > 0:
            print("Missing model keys:", missing_keys)
        if len(unexpected_keys) > 0:
            print("Unexpected model keys:", unexpected_keys)
        if args.few_shot_method != 'finetune':
            optimizer.load_state_dict(checkpoint['optimizer'])
        if args.onecyclelr and args.few_shot_method != 'finetune' and checkpoint['scheduler'] is not None:
            scheduler.load_state_dict(checkpoint['scheduler'])
            
        if prompt_learner is not None:
            if 'prompt_learner' in checkpoint.keys():
                prompt_learner_state_dict = checkpoint['prompt_learner']
                keys = list(prompt_learner_state_dict.keys())
                for k in keys:
                    if 'token_prefix' in k or 'token_suffix' in k:
                        prompt_learner_state_dict.pop(k)
                prompt_learner.load_state_dict(checkpoint['prompt_learner'], strict=False)
            else:
                print("No prompt learner in checkpoint, initializing prompt learner from scratch...")
    else:
        logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
        logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss', 
                        'Train Acc.', 'Train Prec.', 'Train Recl.',
                        'Valid Acc.', 'Val Prec.', 'Val Recl.',
                        'Val On Train Acc.', 'Val On Train Prec.', 'Val On Train Recl.'])
    
    avg_logger = Logger(os.path.join(args.checkpoint, 'log_avg.txt'), title=title)
    avg_logger.set_names(['Avg Train Loss.', 'Avg Val Loss.', 
                          'Avg Train Acc.', 'Avg Train Prec.', 'Avg Train Recl.',
                          'Avg Val Acc.', 'Avg Val Prec.', 'Avg Val Recl.',
                          'Avg Val On Train Acc.', 'Avg Val On Train Prec.', 'Avg Val On Train Recl.'])
    avg_train_loss = AverageMeter()
    avg_val_loss = AverageMeter()
    avg_train_acc = AverageMeter() 
    avg_train_prec = AverageMeter() 
    avg_train_recl = AverageMeter() 
    avg_val_acc = AverageMeter()
    avg_val_prec = AverageMeter()
    avg_val_recl = AverageMeter()
    avg_val_on_train_acc = AverageMeter()
    avg_val_on_train_prec = AverageMeter()
    avg_val_on_train_recl = AverageMeter()

    if args.evaluate:
        print('\nEvaluation only')
        test_loss, test_acc, test_precision, test_recall = test(val_loader, model, criterion, start_epoch, use_cuda, val_transform, 
                                   val_text_features, clip_model, 
                                   local_aggregation=args.local_aggregation,
                                   prompt_learner=prompt_learner, text_encoder=text_encoder, prompt_mode='test',
                                   n_train_cls=len(train_model_ids), n_val_cls=len(val_model_ids))
        print(' Test Loss:  %.8f, Test Acc:  %.2f, Test Prec: %.2f, Test Recl: %.2f' % (test_loss, test_acc, test_precision, test_recall))
        return

    # Train and val
    for epoch in range(start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch, scheduler)

        print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, optimizer.param_groups[0]['lr']))
        
        for _ in range(args.repeat_epochs):
            train_loss, train_acc, train_precision, train_recall = train(train_loader, model, criterion, optimizer, epoch, args.epochs, use_cuda, train_transform, 
                                        train_text_features, clip_model, 
                                        prompt_learner=prompt_learner,
                                        text_encoder=text_encoder,
                                        prompt_mode='train',
                                        local_aggregation=args.local_aggregation,
                                        clip_align_image_classification=args.clip_align_image_classification,
                                        clip_align_image_mse=args.clip_align_image_mse,
                                        clip_align_image_mse_unnorm=args.clip_align_image_mse_unnorm,
                                        clip_align_image_contrastive=args.clip_align_image_contrastive,
                                        clip_align_image_contrastive_mode=args.clip_align_image_contrastive_mode,)
        if not args.skip_val:
            test_loss, test_acc, test_precision, test_recall = test(val_loader, model, criterion, epoch, use_cuda, val_transform, 
                                        val_text_features, clip_model,
                                        target_remap=None, # we don't have target offset here as val set uses all model ids 
                                        local_aggregation=args.local_aggregation,
                                        prompt_learner=prompt_learner, text_encoder=text_encoder, prompt_mode='test',
                                        n_train_cls=len(train_model_ids), n_val_cls=len(val_model_ids))
            val_on_train_acc = 0.0
            val_on_train_precision = 0.0
            val_on_train_recall = 0.0
            if (args.epochs - epoch <= 5) and (val_on_train_loader is not None):
                # in distribution accuracy
                _, val_on_train_acc, val_on_train_precision, val_on_train_recall = test(
                    val_on_train_loader, model, criterion, epoch, use_cuda, val_transform,
                    val_on_train_text_features, clip_model, 
                    local_aggregation=args.local_aggregation,
                    prompt_learner=prompt_learner, text_encoder=text_encoder, prompt_mode='val_on_train')
            # append logger file
            logger.append([optimizer.param_groups[0]['lr'], train_loss, test_loss, 
                           train_acc, train_precision, train_recall, test_acc, test_precision, test_recall,
                           val_on_train_acc, val_on_train_precision, val_on_train_recall])
            
            if args.epochs - epoch <= 5:
                avg_train_loss.update(train_loss)
                avg_val_loss.update(test_loss)
                avg_train_acc.update(train_acc)
                avg_train_prec.update(train_precision)
                avg_train_recl.update(train_recall)
                avg_val_acc.update(test_acc)
                avg_val_prec.update(test_precision)
                avg_val_recl.update(test_recall)
                avg_val_on_train_acc.update(val_on_train_acc)
                avg_val_on_train_prec.update(val_on_train_precision)
                avg_val_on_train_recl.update(val_on_train_recall)
                if args.epochs - epoch == 1:
                    avg_logger.append([avg_train_loss.avg, avg_val_loss.avg, 
                                       avg_train_acc.avg, avg_train_prec.avg, avg_train_recl.avg,
                                       avg_val_acc.avg, avg_val_prec.avg, avg_val_recl.avg,
                                       avg_val_on_train_acc.avg, avg_val_on_train_prec.avg, avg_val_on_train_recl.avg])
                    avg_logger.close()       

            # save model
            is_best = test_acc > best_acc
            best_acc = max(test_acc, best_acc)
            save_dict = {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'acc': test_acc,
                    'best_acc': best_acc,
                    'optimizer' : optimizer.state_dict(),
                    'scheduler' : scheduler.state_dict() if scheduler is not None else None,
                }
            if prompt_learner is not None:
                save_dict['prompt_learner'] = prompt_learner.state_dict()
            save_checkpoint(save_dict, is_best, checkpoint=args.checkpoint)

    logger.close()
    logger.plot()
    savefig(os.path.join(args.checkpoint, 'log.eps'))

    print('Best acc:')
    print(best_acc)

def train(train_loader, model, criterion, optimizer, epoch, total_epochs, use_cuda, 
          train_transform, 
          train_text_features=None, clip_model=None, 
          prompt_learner=None,
          text_encoder=None,
          prompt_mode='train',
          local_aggregation=False,
          clip_align_image_classification=False, 
          clip_align_image_mse=False, clip_align_image_mse_unnorm=False,
          clip_align_image_contrastive=False,
          clip_align_image_contrastive_mode='single',):
    # switch to train mode
    model.train()

    batch_time = AverageMeter()
    data_time = AverageMeter()
    total_losses = AverageMeter()
    losses = AverageMeter()
    aux_mse_losses = AverageMeter()
    aux_contrastive_losses = AverageMeter()
    overall_correct = AverageMeter()
    precision = AverageMeter()
    recall = AverageMeter()
    end = time.time()

    bar = Bar('Processing', max=len(train_loader))
    
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        
        # measure data loading time
        data_time.update(time.time() - end)
        if clip_model is not None:
            (base_rgb, base_depth, hand_rgb, hand_depth), (outputs_clip_base, outputs_clip_hand) = inputs
        else:
            outputs_clip_base = outputs_clip_hand = None
            (base_rgb, base_depth, hand_rgb, hand_depth) = inputs
            
        inputs = torch.cat([base_rgb, base_depth, hand_rgb, hand_depth], dim=1)
        inputs = train_transform(inputs)
        assert inputs.shape[1] == 8 and inputs.shape[2] == 224 and inputs.shape[3] == 224
        targets = targets.float()
        
        if use_cuda:
            inputs, targets = inputs.to(cuda_device), targets.to(cuda_device)
        if clip_model is not None:
            outputs_clip_base = outputs_clip_base.to(cuda_device)
            outputs_clip_hand = outputs_clip_hand.to(cuda_device)
        inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)

        # compute output
        outputs = model(inputs)
        if local_aggregation:
            outputs_local, outputs = outputs # local [B, C, H, W]; global [B, C]
        
        loss = None
        aux_mse_loss = None
        aux_contrastive_loss = None
        if clip_model is not None:
            outputs_norm = nn.functional.normalize(outputs, dim=-1)
            if local_aggregation:
                outputs_local_norm = nn.functional.normalize(outputs_local, dim=1)
            ## input: normalized image features from distilled network, output: unnormalized clip text features
            if prompt_learner is not None and text_encoder is not None:
                assert prompt_mode == 'train'
                prompts_pos, prompts_neg = prompt_learner(outputs_norm) # each [1, n_cls, n_ctx, ctx_dim]
                prompts_pos = prompts_pos[0]
                prompts_neg = prompts_neg[0]
                train_text_features = None # override the input train_text_features since the prompt is adaptive
                train_pos_text_features = text_encoder(
                    prompts_pos, prompt_learner.train_tokenized_prompts[:prompts_pos.shape[0]]
                )[None, ...].tile(outputs_norm.shape[0], 1, 1) # [batch_size, n_cls, d]
                train_neg_text_features = text_encoder(
                    prompts_neg, prompt_learner.train_tokenized_prompts[-prompts_neg.shape[0]:]
                )[None, ...].tile(outputs_norm.shape[0], 1, 1) 
                cur_train_text_features = torch.stack([train_pos_text_features, train_neg_text_features], dim=0) # [2, batch_size, n_cls, d]
                # cur_train_text_features = cur_train_text_features - cur_train_text_features.mean(dim=(0,1,2), keepdim=True) # [2, batch_size, n_cls, d]
            else:
                cur_train_text_features = train_text_features
                            
            cur_train_text_features_norm = nn.functional.normalize(cur_train_text_features, dim=-1)
            if not local_aggregation:
                if cur_train_text_features_norm.dim() == 4:
                    assert prompt_learner is not None and text_encoder is not None
                    classify_outputs = torch.einsum('ni,pnci->pnc', outputs_norm, cur_train_text_features_norm) # [2, batch_size, n_cls]
                else:
                    raise NotImplementedError
            else:
                if cur_train_text_features_norm.dim() == 4:
                    assert prompt_learner is not None and text_encoder is not None
                    classify_outputs = torch.einsum('nihw,pnci->pnchw', outputs_local_norm, cur_train_text_features_norm) # p=2
                else:
                    raise NotImplementedError
                classify_outputs = ((classify_outputs / args.temperature).flatten(-2).softmax(dim=-1).view(classify_outputs.shape) * classify_outputs)
                classify_outputs = classify_outputs.sum(dim=(-1,-2)) # [2, batch_size, n_cls]
                
            if clip_align_image_classification:
                classify_outputs = (classify_outputs + model.global_logit_bias) / args.temperature # temperature
                loss = criterion(classify_outputs, targets)
                
            outputs_clip_base_norm = nn.functional.normalize(outputs_clip_base, dim=1) # channel is always the second dimension of clip features
            outputs_clip_hand_norm = nn.functional.normalize(outputs_clip_hand, dim=1)
                            
            if (clip_align_image_mse or clip_align_image_mse_unnorm or clip_align_image_contrastive):
                
                cur_temp = args.temperature
                
                if clip_align_image_mse:
                    # L-mse
                    out_norm = outputs_norm
                    aux_mse_loss = ((out_norm - outputs_clip_hand_norm) ** 2).sum(dim=1).mean()
                if clip_align_image_mse_unnorm:
                    # L-mse, alternative version
                    out = outputs
                    aux_mse_loss = torch.log(1 + ((out - outputs_clip_hand) ** 2).sum(dim=1).mean() / 10.0)
                if clip_align_image_contrastive:
                    # L-im-cst
                    contrastive_mat_hand = torch.einsum('ni,ci->nc', outputs_norm, outputs_clip_hand_norm) / cur_temp
                    contrastive_labels = torch.arange(contrastive_mat_hand.size(0), device=contrastive_mat_hand.device)
                        
                    if clip_align_image_contrastive_mode == 'single':
                        aux_contrastive_loss = F.cross_entropy(contrastive_mat_hand, contrastive_labels)
                    elif clip_align_image_contrastive_mode == 'bidirectional':
                        aux_contrastive_loss = 0.5 * (
                            F.cross_entropy(contrastive_mat_hand, contrastive_labels)
                            + F.cross_entropy(contrastive_mat_hand.T, contrastive_labels)
                        )
                    else:
                        raise NotImplementedError()
                        
        else:
            B = outputs.size(0)
            if not local_aggregation:
                classify_outputs = outputs.reshape(B, 2, -1).permute(1,0,2)
            else:
                H, W = outputs_local.size(-2), outputs_local.size(-1)
                classify_outputs = outputs_local.reshape(B, 2, -1, H, W).permute(1,0,2,3,4)
                classify_outputs = (classify_outputs.flatten(-2).softmax(dim=-1).view(classify_outputs.shape) * classify_outputs)
                classify_outputs = classify_outputs.sum(dim=(-1,-2))
            classify_outputs = classify_outputs[..., :targets.shape[-1]]
            loss = criterion(classify_outputs, targets)
            
        # measure accuracy and record loss
        if loss is not None:
            losses.update(loss.item(), inputs.size(0))
        if aux_mse_loss is not None:
            aux_mse_losses.update(aux_mse_loss.item(), inputs.size(0))
        if aux_contrastive_loss is not None:
            aux_contrastive_losses.update(aux_contrastive_loss.item(), inputs.size(0))
        
        with torch.no_grad():
            assert classify_outputs.shape[0] == 2, "Not Implemented"
            pred_data = classify_outputs.softmax(dim=0)[0]
            pred_data = torch.clamp(pred_data + model.global_score_bias, 0.0, 1.0)
            correct = ((pred_data > 0.5) * targets + (pred_data <= 0.5) * (1 - targets))
            overall_correct.update(100.0 * (correct.mean()).item(), inputs.size(0))
            precision.update(100.0 * (((pred_data > 0.5) * targets).sum() / (pred_data > 0.5).sum()).item(), inputs.size(0))
            recall.update(100.0 * (((pred_data > 0.5) * targets).sum() / (targets > 0.5).sum()).item(), inputs.size(0))

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss_for_bkward = 0.0
        if loss is not None:
            loss_for_bkward += loss
        if aux_mse_loss is not None:
            loss_for_bkward += aux_mse_loss
        if aux_contrastive_loss is not None:
            loss_for_bkward += aux_contrastive_loss
        total_losses.update(loss_for_bkward.item(), inputs.size(0))
        loss_for_bkward.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        # plot progress
        bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Tot Loss: {total_loss:.4f} | Classification loss: {loss:.4f} | L-MSE: {aux_mse_loss:.4f} | L-im-cst: {aux_contrastive_loss:.4f} | Overall Correct: {overall_correct: .4f} | Precision: {precision: .4f} | Recall: {recall: .4f}'.format(
                    batch=batch_idx + 1,
                    size=len(train_loader),
                    data=data_time.val,
                    bt=batch_time.val,
                    total=bar.elapsed_td,
                    eta=bar.eta_td,
                    total_loss=total_losses.avg,
                    loss=losses.avg if loss is not None else 0.0,
                    aux_mse_loss=aux_mse_losses.avg if aux_mse_loss is not None else 0.0,
                    aux_contrastive_loss=aux_contrastive_losses.avg if aux_contrastive_loss is not None else 0.0,
                    overall_correct=overall_correct.avg,
                    precision=precision.avg,
                    recall=recall.avg,
                    )
        bar.next()
    bar.finish()
    return (total_losses.avg, overall_correct.avg, precision.avg, recall.avg)

def test(val_loader, model, criterion, epoch, use_cuda, 
         val_transform, val_text_features=None, clip_model=None, 
         target_remap=None, local_aggregation=False,
         prompt_learner=None, text_encoder=None, prompt_mode='test',
         n_train_cls=None, n_val_cls=None):
    global best_acc

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    overall_correct = AverageMeter()
    precision = AverageMeter()
    recall = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    bar = Bar('Processing', max=len(val_loader))
    avg_accuracy_per_class = None
    avg_precision_per_class = None
    avg_recall_per_class = None
            
    tot_idx = 0
    for batch_idx, (inputs, targets) in enumerate(val_loader):
        # measure data loading time
        data_time.update(time.time() - end)
        
        tot_idx += targets.shape[0]
        
        (base_rgb, base_depth, hand_rgb, hand_depth) = inputs # note that the validation datasets should not return clip image features
        inputs = torch.cat([base_rgb, base_depth, hand_rgb, hand_depth], dim=1)
        inputs = val_transform(inputs)
        targets = targets.float()

        if use_cuda:
            inputs, targets = inputs.to(cuda_device), targets.to(cuda_device)
        inputs, targets = torch.autograd.Variable(inputs, volatile=True), torch.autograd.Variable(targets)

        # compute output
        with torch.no_grad():
            outputs = model(inputs)
            if local_aggregation:
                outputs_local, outputs = outputs # local [B, C, H, W]; global [B, C]
            cur_val_text_features = val_text_features
            
        if clip_model is not None:
            outputs_norm = nn.functional.normalize(outputs, dim=-1)
            if local_aggregation:
                outputs_local_norm = nn.functional.normalize(outputs_local, dim=1)
            ## input: normalized clip image features, output: normalized clip text features
            if prompt_learner is not None and text_encoder is not None:
                with torch.no_grad():            
                    if prompt_mode == 'train':
                        tokenized_prompts = prompt_learner.train_tokenized_prompts
                    elif prompt_mode == 'test':
                        tokenized_prompts = prompt_learner.val_tokenized_prompts
                    elif prompt_mode == 'val_on_train':
                        tokenized_prompts = prompt_learner.val_on_train_tokenized_prompts
                    else: 
                        raise NotImplementedError()
                    prompts_pos, prompts_neg = prompt_learner(outputs_norm.to(cuda_device), mode=prompt_mode)
                    prompts_pos = prompts_pos[0]
                    prompts_neg = prompts_neg[0]
                    cur_val_text_features = None # override the input val_text_features since the prompt is adaptive
                    val_pos_text_features = text_encoder(
                        prompts_pos, tokenized_prompts[:prompts_pos.shape[0]]
                    )[None, ...].tile(outputs_norm.shape[0], 1, 1) # [batch_size, n_cls, d]
                    val_neg_text_features = text_encoder(
                        prompts_neg, tokenized_prompts[-prompts_neg.shape[0]:]
                    )[None, ...].tile(outputs_norm.shape[0], 1, 1) 
                    cur_val_text_features = torch.stack([val_pos_text_features, val_neg_text_features], dim=0) # [2, batch_size, n_cls, d]
                    # cur_val_text_features = cur_val_text_features - cur_val_text_features.mean(dim=(0,1,2), keepdim=True) # [2, batch_size, n_cls, d]
                    cur_val_text_features_norm = nn.functional.normalize(cur_val_text_features, dim=-1)
                    if not local_aggregation:
                        classify_outputs = torch.einsum('ni,pnmi->pnm', outputs_norm, cur_val_text_features_norm) # [2, batch_size, n_cls]
                    else:
                        classify_outputs = torch.einsum('nihw,pnmi->pnmhw', outputs_local_norm, cur_val_text_features_norm) # p=2
            else:
                raise NotImplementedError()
                
            if local_aggregation:
                classify_outputs = ((classify_outputs / args.temperature).flatten(-2).softmax(dim=-1).view(classify_outputs.shape) * classify_outputs)
                classify_outputs = classify_outputs.sum(dim=(-1,-2)) # [batch_size, n_cls]
                
            classify_outputs = (classify_outputs + model.global_logit_bias)
            classify_outputs = classify_outputs / args.temperature # temperature
        else:
            B = outputs.size(0)
            if not local_aggregation:
                classify_outputs = outputs.reshape(B, 2, -1).permute(1,0,2)
            else:
                H, W = outputs_local.size(-2), outputs_local.size(-1)
                classify_outputs = outputs_local.reshape(B, 2, -1, H, W).permute(1,0,2,3,4)
                classify_outputs = (classify_outputs.flatten(-2).softmax(dim=-1).view(classify_outputs.shape) * classify_outputs)
                classify_outputs = classify_outputs.sum(dim=(-1,-2))
            if target_remap is not None:
                assert len(target_remap) == 2
                classify_outputs = classify_outputs[..., target_remap[0]:target_remap[1]]
            
        loss = criterion(classify_outputs, targets)
        
        with torch.no_grad():
            assert classify_outputs.shape[0] == 2, "Not implemented"
            pred_data = classify_outputs.softmax(dim=0)[0]
            pred_data = torch.clamp(pred_data + model.global_score_bias, 0.0, 1.0)
            correct = ((pred_data > 0.5) * targets + (pred_data <= 0.5) * (1 - targets))
            overall_correct.update(100.0 * (correct.mean()).item(), inputs.size(0))
            precision.update(100.0 * (((pred_data > 0.5) * targets).sum() / (pred_data > 0.5).sum()).item(), inputs.size(0))
            recall.update(100.0 * (((pred_data > 0.5) * targets).sum() / (targets > 0.5).sum()).item(), inputs.size(0))
        
        if avg_accuracy_per_class is None:
            avg_accuracy_per_class = [[0.0, 0.0] for _ in range(classify_outputs.shape[-1])]
            avg_precision_per_class = [[0.0, 0.0] for _ in range(classify_outputs.shape[-1])]
            avg_recall_per_class = [[0.0, 0.0] for _ in range(classify_outputs.shape[-1])]
        for i in range(classify_outputs.shape[-1]):
            avg_accuracy_per_class[i][0] += correct[:, i].sum().item()
            avg_accuracy_per_class[i][1] += correct[:, i].shape[0]
            avg_precision_per_class[i][0] += ((pred_data[:,i] > 0.5) * targets[:,i]).sum().item()
            avg_precision_per_class[i][1] += (pred_data[:,i] > 0.5).sum().item() + 1e-8
            avg_recall_per_class[i][0] += ((pred_data[:,i] > 0.5) * targets[:,i]).sum().item()
            avg_recall_per_class[i][1] += (targets[:,i] > 0.5).sum().item() + 1e-8

        # measure accuracy and record loss
        losses.update(loss.item(), inputs.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        # plot progress
        bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Overall Correct: {overall_correct: .4f} | Precision: {precision: .4f} | Recall: {recall: .4f}'.format(
                    batch=batch_idx + 1,
                    size=len(val_loader),
                    data=data_time.avg,
                    bt=batch_time.avg,
                    total=bar.elapsed_td,
                    eta=bar.eta_td,
                    loss=losses.avg,
                    overall_correct=overall_correct.avg,
                    precision=precision.avg,
                    recall=recall.avg,
                    )
        bar.next()
    bar.finish()
    
    if n_train_cls is not None and n_val_cls is not None:
        print("Avg train class accuracy", sum([x[0] for x in avg_accuracy_per_class[:n_train_cls]]) / sum([x[1] for x in avg_accuracy_per_class[:n_train_cls]]))
        print("Avg train class precision", sum([x[0] for x in avg_precision_per_class[:n_train_cls]]) / sum([x[1] for x in avg_precision_per_class[:n_train_cls]]))
        print("Avg train class recall", sum([x[0] for x in avg_recall_per_class[:n_train_cls]]) / sum([x[1] for x in avg_recall_per_class[:n_train_cls]]))
        print("Avg val class accuracy", sum([x[0] for x in avg_accuracy_per_class[-n_val_cls:]]) / sum([x[1] for x in avg_accuracy_per_class[-n_val_cls:]]))
        print("Avg val class precision", sum([x[0] for x in avg_precision_per_class[-n_val_cls:]]) / sum([x[1] for x in avg_precision_per_class[-n_val_cls:]]))
        print("Avg val class recall", sum([x[0] for x in avg_recall_per_class[-n_val_cls:]]) / sum([x[1] for x in avg_recall_per_class[-n_val_cls:]]))
        
    avg_accuracy_per_class = [x[0] / x[1] for x in avg_accuracy_per_class]
    print("Average accuracy per class: {}".format(avg_accuracy_per_class))
    avg_precision_per_class = [x[0] / x[1] for x in avg_precision_per_class]
    print("Average precision per class: {}".format(avg_precision_per_class))
    avg_recall_per_class = [x[0] / x[1] for x in avg_recall_per_class]
    print("Average recall per class: {}".format(avg_recall_per_class))
    
    return (losses.avg, overall_correct.avg, precision.avg, recall.avg)

def save_checkpoint(state, is_best, checkpoint='checkpoint', filename='checkpoint.pth.tar'):
    filepath = os.path.join(checkpoint, filename)
    torch.save(state, filepath)
    if is_best:
        shutil.copyfile(filepath, os.path.join(checkpoint, 'model_best.pth.tar'))

def adjust_learning_rate(optimizer, epoch, scheduler=None):
    global state
    if epoch in args.schedule and scheduler is None:
        state['lr'] *= args.gamma
        for param_group in optimizer.param_groups:
            param_group['lr'] = state['lr']
    if scheduler is not None and epoch > 0:
        scheduler.step()

if __name__ == '__main__':
    main()