train_sota_scan.py

import torch

from utils.util import seed_torch
from utils.logging import Logger
import os
import sys
import wandb
import math

from utils.sinkhorn_knopp import SinkhornKnopp

from models.build_scan import build_scan
from data.scan_build_dataset import build_data
from methods.scan import SCAN
from utils.scan_memory import MemoryBank
from data.config_dataset import set_dataset_config


if __name__ == "__main__":
    import argparse

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

    # Hyper-parameters Setting
    parser.add_argument('--epochs_scan', default=200, type=int)
    parser.add_argument('--epochs_selflabel', default=200, type=int)

    parser.add_argument('--batch_size', default=256, 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)

    # SCAN knobs
    # parser.add_argument('--stage', type=str, default='scan', choices=['scan', 'selflabel'])
    parser.add_argument("--softmax_temp", default=0.1, type=float, help="softmax temperature")
    parser.add_argument("--w_entropy", default=1.0, type=float, help="entropy loss temperature")
    parser.add_argument("--confidence_threshold", default=0.3, type=float,
                        help="confidence_threshold for self-label training")
    parser.add_argument('--apply_class_balancing', action='store_true', default=False)
    parser.add_argument('--use_norm', action='store_true', default=False)

    # 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',
                                                                            'vit_uno_clip'])
    parser.add_argument('--num_workers', default=8, type=int)

    # Strategy Setting
    parser.add_argument('--num_steps', default=10, type=int)
    parser.add_argument('--current_step', default=0, type=int)

    # Model Config
    parser.add_argument('--mode', type=str, default='train', choices=['train', 'eval'])
    parser.add_argument('--model_name', type=str, default='vit_dino', choices=['vit_dino', 'clip', 'resnet50_dino'])
    parser.add_argument('--grad_from_block', type=int, default=12)  # 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_scan/')

    parser.add_argument('--exp_marker', type=str, default='nonsense_expt')
    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]

    # Experimental Dir.
    model_dir = os.path.join(args.exp_root, f"{runner_name}_CosNorm={args.use_norm}_{args.model_name}_{args.dataset_name}_Steps{args.num_steps}")
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    # path to pre-trained single heads weights .pth file
    args.learned_single_head_path_dict = {'scan': [], 'selflabel': []}
    for step in range(args.current_step):
        this_single_path_scan = model_dir + f"/scan_SingleHead_S{step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"
        this_single_path_selflabel = model_dir + f"/selflabel_SingleHead_S{step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"
        args.learned_single_head_path_dict['scan'].append(this_single_path_scan)
        args.learned_single_head_path_dict['selflabel'].append(this_single_path_selflabel)

    # path to save single head selflabel
    args.save_single_path_scan = model_dir + f"/scan_SingleHead_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"
    args.save_joint_path_scan = model_dir + f"/scan_JointHead_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"

    args.save_single_path_selflabel = model_dir + f"/selflabel_SingleHead_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"
    args.save_joint_path_selflabel = model_dir + f"/selflabel_JointHead_S{args.current_step}_{args.dataset_name}_Steps{args.num_steps}_{args.model_head}_Mlp{args.num_mlp_layers}.pth"

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

    print('log_dir=', args.log_dir)

    # WandB setting
    if args.mode == 'train':
        wandb_run_name = f'SCAN_CosNorm={args.use_norm}_{args.dataset_name}_{args.model_name}_S{str(args.current_step)}-{args.num_steps}_{args.model_head}'
        wandb.init(project='SOTA_scan',
                   entity=args.wandb_entity,
                   tags=[f'TotalStep={args.num_steps}', args.dataset_name, f'Steps={str(args.current_step)}',
                         args.model_name, f'CosNorm={args.use_norm}', f'w_entropy={args.w_entropy}',
                         f'conf_thresh={args.confidence_threshold}', f'class_balancing={args.apply_class_balancing}'],
                   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
    # Parameters
    if 'vit' in args.model_name:
        args.image_size = 224
    elif 'resnet' in args.model_name:
        args.image_size = 64
    else:
        raise NotImplementedError

    # args.image_size = 224
    args.interpolation = 3
    args.crop_pct = 0.875
    args.pretrain_path = args.dino_pretrain_path
    args.mlp_out_dim = args.num_novel_per_step

    args.feat_dim = 768

    # ----------------------
    # Model + Classifier Creation for this iNCD step
    # ----------------------
    encoder, single_head, single_heads_dict, joint_head_dict = build_scan(args)

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

    # NN-finding dataloader
    memory_loader, N_train = data_factory.get_dataloader(split='train', aug=None, shuffle=False,
                                                         target_list=range(args.current_novel_start,
                                                                           args.current_novel_end))
    # ----------------------
    # Memory Bank creation
    # ----------------------
    #   |- init MemBank
    memory_bank_train = MemoryBank(n=N_train, dim=args.feat_dim, num_classes=args.num_novel_per_step, temperature=0.1)
    memory_bank_train.to(args.device)
    #   |- fill memory bank (feature extraction)
    memory_bank_train.fill_memory_bank(memory_loader, encoder)
    #   |- mine the K-nearest_neighbors and generate indices
    topk = 20 if 'cifar' in args.dataset_name else 50
    to_neighbors_dataset_indices = memory_bank_train.mine_nearest_neighbors(topk=topk)

    # Train loader

    #   |- training loader for SCAN-Stage
    ulb_train_loader_scan, _ = data_factory.get_dataloader(
        split='train', aug='scan', shuffle=True,
        target_list=range(args.current_novel_start, args.current_novel_end),
        to_neighbors_dataset_indices=to_neighbors_dataset_indices)

    #   |- training loader for Self-label-Stage
    ulb_train_loader_selflabel, _ = data_factory.get_dataloader(
        split='train', aug='selflabel', shuffle=True,
        target_list=range(args.current_novel_start, args.current_novel_end),
        to_neighbors_dataset_indices=None)

    # 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))

    val_split = args.val_split
    test_split = args.test_split

    # 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)

    print(args)

    if args.mode == 'train':
        method = SCAN(encoder=encoder,
                      single_head=single_head,
                      learned_single_heads_dict=single_heads_dict, joint_head_dict=joint_head_dict,
                      train_loader_scan=ulb_train_loader_scan, train_loader_selflabel=ulb_train_loader_selflabel,
                      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
        #   |- SCAN stage
        #       |- train
        method.train_scan(args)
        #       |- test
        method.test(args, stage='scan')
        #       |- show
        method.print_eval_results(args, stage='scan')
        #       |- save
        method.save_single(path=args.save_single_path_scan)
        method.save_joint_head(path=args.save_joint_path_scan)

        #   |- Self-label stage
        #       |- train
        method.train_selflabel(args)
        #       |- test
        method.test(args, stage='selflabel')
        #       |- show
        method.print_eval_results(args, stage='scan')
        method.print_eval_results(args, stage='selflabel')
        #       |- save
        method.save_single(path=args.save_single_path_selflabel)
        method.save_joint_head(path=args.save_joint_path_selflabel)

    elif args.mode == 'eval':
        raise NotImplementedError
    else:
        raise NotImplementedError