import warnings warnings.filterwarnings("ignore", category=RuntimeWarning) import os os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' from torch.utils.data import DataLoader from tqdm import tqdm import argparse import json import os import torch # Importing from local modules from tools import write2csv, setup_paths, setup_seed, log_metrics, Logger from dataset import get_data from method import AdaCLIP_Trainer setup_seed(111) def train(args): # Configurations epochs = args.epoch learning_rate = args.learning_rate batch_size = args.batch_size image_size = args.image_size device = 'cuda' if torch.cuda.is_available() else 'cpu' save_fig = args.save_fig # Set up paths model_name, image_dir, csv_path, log_path, ckp_path, tensorboard_logger = setup_paths(args) # Logger logger = Logger(log_path) # Print basic information for key, value in sorted(vars(args).items()): logger.info(f'{key} = {value}') logger.info('Model name: {:}'.format(model_name)) config_path = os.path.join('./model_configs', f'{args.model}.json') # Prepare model with open(config_path, 'r') as f: model_configs = json.load(f) # Set up the feature hierarchy n_layers = model_configs['vision_cfg']['layers'] substage = n_layers // 4 features_list = [substage, substage * 2, substage * 3, substage * 4] model = AdaCLIP_Trainer( backbone=args.model, feat_list=features_list, input_dim=model_configs['vision_cfg']['width'], output_dim=model_configs['embed_dim'], learning_rate=learning_rate, device=device, image_size=image_size, prompting_depth=args.prompting_depth, prompting_length=args.prompting_length, prompting_branch=args.prompting_branch, prompting_type=args.prompting_type, use_hsf=args.use_hsf, k_clusters=args.k_clusters ).to(device) train_data_cls_names, train_data, train_data_root = get_data( dataset_type_list=args.training_data, transform=model.preprocess, target_transform=model.transform, training=True) test_data_cls_names, test_data, test_data_root = get_data( dataset_type_list=args.testing_data, transform=model.preprocess, target_transform=model.transform, training=False) logger.info('Data Root: training, {:}; testing, {:}'.format(train_data_root, test_data_root)) train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True) test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False) # Typically, we use MVTec or VisA as the validation set. The best model from this validation # process is then used for zero-shot anomaly detection on novel categories. best_f1 = -1e1 for epoch in tqdm(range(epochs)): loss = model.train_epoch(train_dataloader) # Logs if (epoch + 1) % args.print_freq == 0: logger.info('epoch [{}/{}], loss:{:.4f}'.format(epoch + 1, epochs, loss)) tensorboard_logger.add_scalar('loss', loss, epoch) # Validation if (epoch + 1) % args.valid_freq == 0 or (epoch == epochs - 1): if epoch == epochs - 1: save_fig_flag = save_fig else: save_fig_flag = False logger.info('=============================Testing ====================================') metric_dict = model.evaluation( test_dataloader, test_data_cls_names, save_fig_flag, image_dir, ) log_metrics( metric_dict, logger, tensorboard_logger, epoch ) f1_px = metric_dict['Average']['f1_px'] # Save best if f1_px > best_f1: for k in metric_dict.keys(): write2csv(metric_dict[k], test_data_cls_names, k, csv_path) ckp_path_best = ckp_path + '_best.pth' model.save(ckp_path_best) best_f1 = f1_px def str2bool(v): return v.lower() in ("yes", "true", "t", "1") if __name__ == '__main__': parser = argparse.ArgumentParser("AdaCLIP", add_help=True) # Paths and configurations parser.add_argument("--training_data", type=str, default=["mvtec", "colondb"], nargs='+', help="Datasets for training (default: ['mvtec', 'colondb'])") parser.add_argument("--testing_data", type=str, default="visa", help="Dataset for testing (default: 'visa')") parser.add_argument("--save_path", type=str, default='./workspaces', help="Directory to save results (default: './workspaces')") parser.add_argument("--model", type=str, default="ViT-L-14-336", choices=["ViT-B-16", "ViT-B-32", "ViT-L-14", "ViT-L-14-336"], help="The CLIP model to be used (default: 'ViT-L-14-336')") parser.add_argument("--save_fig", type=str2bool, default=False, help="Save figures for visualizations (default: False)") parser.add_argument("--ckt_path", type=str, default='', help="Path to the pre-trained model (default: '')") # Hyper-parameters parser.add_argument("--exp_indx", type=int, default=0, help="Index of the experiment (default: 0)") parser.add_argument("--epoch", type=int, default=5, help="Number of epochs (default: 5)") parser.add_argument("--learning_rate", type=float, default=0.01, help="Learning rate (default: 0.01)") parser.add_argument("--batch_size", type=int, default=1, help="Batch size (default: 1)") parser.add_argument("--image_size", type=int, default=518, help="Size of the input images (default: 518)") parser.add_argument("--print_freq", type=int, default=1, help="Frequency of print statements (default: 1)") parser.add_argument("--valid_freq", type=int, default=1, help="Frequency of validation (default: 1)") # Prompting parameters parser.add_argument("--prompting_depth", type=int, default=4, help="Depth of prompting (default: 4)") parser.add_argument("--prompting_length", type=int, default=5, help="Length of prompting (default: 5)") parser.add_argument("--prompting_type", type=str, default='SD', choices=['', 'S', 'D', 'SD'], help="Type of prompting. 'S' for Static, 'D' for Dynamic, 'SD' for both (default: 'SD')") parser.add_argument("--prompting_branch", type=str, default='VL', choices=['', 'V', 'L', 'VL'], help="Branch of prompting. 'V' for Visual, 'L' for Language, 'VL' for both (default: 'VL')") parser.add_argument("--use_hsf", type=str2bool, default=True, help="Use HSF for aggregation. If False, original class embedding is used (default: True)") parser.add_argument("--k_clusters", type=int, default=20, help="Number of clusters (default: 20)") args = parser.parse_args() train(args) if args.batch_size != 1: raise NotImplementedError( "Currently, only batch size of 1 is supported due to unresolved bugs. Please set --batch_size to 1.") train(args)