train_upper_unlocked_joint_teacher_student.py

import torch
import torch.nn.functional as F
from torch.optim import SGD, lr_scheduler
from sklearn.metrics.cluster import normalized_mutual_info_score as nmi_score
from sklearn.metrics import adjusted_rand_score as ari_score
from utils.util import cluster_acc, AverageMeter, seed_torch
from utils.logging import Logger
from tqdm import tqdm
import numpy as np
import os
import sys
import copy
import wandb
import math

from utils.sinkhorn_knopp import SinkhornKnopp
from models.build_unlocked_upper import build_unlocked_model
from data.build_dataset import build_data
from data.config_dataset import set_dataset_config

from methods.upper_bound_unlocked_method import UnlockedUpperBound

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='cluster', formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Hyper-parameters Setting
    parser.add_argument('--epochs', default=100, type=int)
    parser.add_argument('--batch_size', default=128, type=int)

    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)

    # Strategy tricks
    # parser.add_argument('--l2_single_cls', action='store_true', default=False,
    #                     help='L2 normalize single classifier weights before forward-prop')
    parser.add_argument('--student_loss', type=str, default='ZP', choices=['CE', 'ZP'], help="CE(cross-entropy loss w/o zero-padding), ZP(w/ zero-padding")

    # UNO knobs
    parser.add_argument("--softmax_temp", default=0.1, type=float, help="softmax temperature")
    parser.add_argument("--threshold", default=0.5, type=float, help="threshold for hard pseudo-labeling")
    parser.add_argument("--num_iters_sk", default=3, type=int, help="number of iters for Sinkhorn")
    parser.add_argument("--epsilon_sk", default=0.05, type=float, help="epsilon for the Sinkhorn")
    parser.add_argument('--alpha', default=0.75, type=float)

    # Dataset Setting
    parser.add_argument('--dataset_name', type=str, default='cifar100', choices=['cifar10', 'cifar100', 'tinyimagenet',
                                                                                 'cub200', 'herb19', 'scars',
                                                                                 'aircraft'])
    # parser.add_argument('--dataset_root', type=str, default='./data/datasets/CIFAR/')
    # parser.add_argument('--num_classes', default=100, type=int)
    parser.add_argument('--aug_type', type=str, default='vit_uno', choices=['vit_frost', 'vit_uno', 'resnet'])
    parser.add_argument('--num_workers', default=2, type=int)

    # Strategy Setting
    parser.add_argument('--num_steps', default=5, type=int)
    parser.add_argument('--current_step', default=0, type=int)
    parser.add_argument('--stage', type=str, default='stage1', choices=['stage1', 'stage2'])

    # Model Config
    parser.add_argument('--mode', type=str, default='train', choices=['train', 'eval'])
    parser.add_argument('--model_name', type=str, default='vit_dino')
    parser.add_argument('--grad_from_block', type=int, default=11)  # 12->do not fine tune backbone at all
    parser.add_argument('--num_mlp_layers', type=int, default=1)  # 12->do not fine tune backbone at all
    parser.add_argument('--dino_pretrain_path', type=str,
                        default='./models/dino_weights/dino_vitbase16_pretrain.pth')
    parser.add_argument('--model_head', type=str, default='LinearHead', choices=['LinearHead', 'DINOHead'])

    # Experimental Setting
    parser.add_argument('--seed', default=10, type=int)
    parser.add_argument('--exp_root', type=str, default='./outputs_upper_unlocked/')
    parser.add_argument('--weights_root', type=str, default='./models/single_weights/')

    parser.add_argument('--exp_marker', type=str, default='warmedup')
    parser.add_argument('--wandb_mode', type=str, default='online', choices=['online', 'offline', 'disabled'])
    parser.add_argument('--wandb_entity', type=str, default='oatmealliu')

    # ----------------------
    # Initial Configurations
    # ----------------------
    args = parser.parse_args()

    # init. dataset config.
    args = set_dataset_config(args)

    # init. config.
    args.cuda = torch.cuda.is_available()
    device = torch.device("cuda" if args.cuda else "cpu")
    args.device = torch.device("cuda" if args.cuda else "cpu")
    seed_torch(args.seed)

    # init. experimental output path
    runner_name = os.path.basename(__file__).split(".")[0]

    # set a dir name which can describe the experiment
    model_dir = os.path.join(args.exp_root, f"{runner_name}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}_{args.exp_marker}")
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    if args.stage == 'stage1':
        # Warmed-up single head path
        warmup_student_head_dir = f"{args.weights_root}{args.model_name}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}"
        # warmup_student_head_dir = f"{args.weights_root}{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}"
        args.warmup_student_head_path = warmup_student_head_dir+f"/SingleHead_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"

        # Single model and head saving path
        args.save_student_model_path = model_dir + f"/single_Backbone_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"
        args.save_student_head_path = model_dir + f"/single_Head_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"

    if args.stage == 'stage2':
        # path to pre-trained teacher heads weights .pth file
        args.pretrained_teacher_backbone_paths_list = []
        args.pretrained_teacher_head_paths_list = []
        for step in range(1+args.current_step):
            this_teacher_backbone_path = model_dir + f"/single_Backbone_S{step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"
            this_teacher_head_path = model_dir + f"/single_Head_S{step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"
            args.pretrained_teacher_backbone_paths_list.append(this_teacher_backbone_path)
            args.pretrained_teacher_head_paths_list.append(this_teacher_head_path)


        # Joint model and head saving path
        args.save_joint_model_path = model_dir + f"/joint_Backbone_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"
        args.save_joint_head_path = model_dir + f"/joint_Head_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}.pth"

    args.log_dir = model_dir + f'/{args.stage}_{args.dataset_name}_S{str(args.current_step)}-{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}_log.txt'
    sys.stdout = Logger(args.log_dir)

    print('log_dir=', args.log_dir)

    # WandB setting
    wandb_tags = [f'TotalStep={args.num_steps}', "Unlocked", args.stage, args.dataset_name, args.model_name,
                  f'Steps={str(args.current_step)}', args.exp_marker, f'Block={args.grad_from_block}']
    wandb_run_name = f'{args.stage}_Upper_unlocked_{args.dataset_name}_S{str(args.current_step)}/{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}_Loss{args.student_loss}_Block{args.grad_from_block}_{args.exp_marker}'
    wandb.init(project='Upper_Joint_Train_Unlock',
               entity=args.wandb_entity,
               tags=wandb_tags,
               name=wandb_run_name,
               mode=args.wandb_mode)

    # ----------------------
    # Experimental Setting Initialization
    # ----------------------
    # Dataset Split Params
    args.num_novel_interval = math.ceil(args.num_classes / args.num_steps)
    args.current_novel_start = args.num_novel_interval * args.current_step
    args.current_novel_end = args.num_novel_interval * (args.current_step + 1) \
        if args.num_novel_interval * (args.current_step + 1) <= args.num_classes \
        else args.num_classes
    args.num_novel_per_step = args.current_novel_end - args.current_novel_start

    # ViT DINO B/16 Params
    # Parameters
    args.image_size = 224
    args.interpolation = 3
    args.crop_pct = 0.875
    args.pretrain_path = args.dino_pretrain_path
    args.feat_dim = 768
    args.mlp_out_dim = args.num_novel_per_step

    # ----------------------
    # Dataloaders Creation for this iNCD step
    # ----------------------
    data_factory = build_data(args)

    # Train loader list for joint training
    ulb_train_loader_list = []
    for s in range(1 + args.current_step):
        if (1 + s) < args.num_steps:
            s_ulb_train_loader = data_factory.get_dataloader(split='train', aug='twice', shuffle=True,
                                                             target_list=range(s * args.num_novel_interval,
                                                                               (1 + s) * args.num_novel_interval))
        else:
            s_ulb_train_loader = data_factory.get_dataloader(split='train', aug='twice', shuffle=True,
                                                             target_list=range(args.current_novel_start,
                                                                               args.current_novel_end))
        ulb_train_loader_list.append(s_ulb_train_loader)


    val_split = args.val_split
    test_split = args.test_split

    # Mixed-val loader
    if args.current_step > 0:
        ulb_all_prev_val_loader = data_factory.get_dataloader(split='train', aug=None, shuffle=False,
                                                              target_list=range(args.current_novel_start))
    else:
        ulb_all_prev_val_loader = None

    ulb_all_val_loader = data_factory.get_dataloader(split='train', aug=None, shuffle=False,
                                                     target_list=range(args.current_novel_end))

    # Mixed-test loader
    if args.current_step > 0:
        ulb_all_prev_test_loader = data_factory.get_dataloader(split='test', aug=None, shuffle=False,
                                                               target_list=range(args.current_novel_start))
    else:
        ulb_all_prev_test_loader = None

    ulb_all_test_loader = data_factory.get_dataloader(split='test', aug=None, shuffle=False,
                                                      target_list=range(args.current_novel_end))

    # Step-wise val/test loader list
    ulb_step_val_loader_list = []
    ulb_step_test_loader_list = []
    for s in range(1 + args.current_step):
        if (1 + s) < args.num_steps:
            this_ulb_val_loader = data_factory.get_dataloader(split=val_split, aug=None, shuffle=False,
                                                              target_list=range(s * args.num_novel_interval,
                                                                                (1 + s) * args.num_novel_interval))
            this_ulb_test_loader = data_factory.get_dataloader(split=test_split, aug=None, shuffle=False,
                                                               target_list=range(s * args.num_novel_interval,
                                                                                 (1 + s) * args.num_novel_interval))
        else:
            this_ulb_val_loader = data_factory.get_dataloader(split=val_split, aug=None, shuffle=False,
                                                              target_list=range(args.current_novel_start,
                                                                                args.current_novel_end))
            this_ulb_test_loader = data_factory.get_dataloader(split=test_split, aug=None, shuffle=False,
                                                               target_list=range(args.current_novel_start,
                                                                                 args.current_novel_end))

        ulb_step_val_loader_list.append(this_ulb_val_loader)
        ulb_step_test_loader_list.append(this_ulb_test_loader)

    if args.mode == 'train' and args.stage == 'stage1':
        student_model, student_head = build_unlocked_model(args)

        print(args)

        # print("------> Single head for this step:")
        # print(student_head)

        # Create SinkhornKnopp pseudo-labeling algo.
        sinkhorn = SinkhornKnopp(args)

        # TeacherStudent learning strategy
        method = UnlockedUpperBound(student_model=student_model, student_head=student_head,
                                    teacher_pair_list=None,
                                    joint_model=None, joint_head=None,
                                    sinkhorn=sinkhorn,
                                    train_loader_list=ulb_train_loader_list,
                                    ulb_step_val_list=ulb_step_val_loader_list,
                                    ulb_all_prev_val=ulb_all_prev_val_loader,
                                    ulb_all_val=ulb_all_val_loader,
                                    ulb_step_test_list=ulb_step_test_loader_list,
                                    ulb_all_prev_test=ulb_all_prev_test_loader,
                                    ulb_all_test=ulb_all_test_loader)

        # Training
        #   |- Stage-1: train a task-specific model and head
        method.train_single(args)
        method.save_student(model_path=args.save_student_model_path, head_path=args.save_student_head_path)

        method.test_single(args)
    elif args.mode == 'train' and args.stage == 'stage2':
        # load and create models
        #   |- teacher_pair_list: pre-trained single models and heads
        #   |- joint_model and joint_head: to-be-trained model and head
        teacher_pair_list, joint_model, joint_head = build_unlocked_model(args)

        print(args)

        print("------> Single heads from previous steps:")
        for (_, teacher_head) in teacher_pair_list:
            print(teacher_head)

        print("------> Joint head:")
        print(joint_head)

        # TeacherStudent learning strategy
        method = UnlockedUpperBound(student_model=None, student_head=None,
                                    teacher_pair_list=teacher_pair_list,
                                    joint_model=joint_model, joint_head=joint_head,
                                    sinkhorn=None,
                                    train_loader_list=ulb_train_loader_list,
                                    ulb_step_val_list=ulb_step_val_loader_list,
                                    ulb_all_prev_val=ulb_all_prev_val_loader,
                                    ulb_all_val=ulb_all_val_loader,
                                    ulb_step_test_list=ulb_step_test_loader_list,
                                    ulb_all_prev_test=ulb_all_prev_test_loader,
                                    ulb_all_test=ulb_all_test_loader)

        # Training
        #   |- Stage-2: train a task-agnostic joint model and head
        method.train_joint(args)
        method.save_joint(model_path=args.save_joint_model_path, head_path=args.save_joint_head_path)
        method.test_joint(args)
    elif args.mode == 'eval':
        raise NotImplementedError
    else:
        raise NotImplementedError