import os os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE' import AnomalyCLIP_lib import torch import argparse import torch.nn.functional as F from training_libs.prompt_ensemble import AnomalyCLIP_PromptLearner from training_libs.loss import FocalLoss, BinaryDiceLoss from training_libs.utils import normalize from training_libs.dataset import Dataset_train from training_libs.logger import get_logger from tqdm import tqdm import numpy as np import random from training_libs.utils import get_transform import matplotlib.pyplot as plt import warnings warnings.filterwarnings("ignore", category=UserWarning) def setup_seed(seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False class RealTimePlotter: # def __init__(self): self.epochs = [] self.loss_list = [] self.image_loss_list = [] self.fig, (self.ax1, self.ax2) = plt.subplots(1, 2, figsize=(14, 6)) plt.ion() self.fig.show() self.fig.canvas.flush_events() def update(self, epoch, loss, image_loss): self.epochs.append(epoch) self.loss_list.append(loss) self.image_loss_list.append(image_loss) self.ax1.clear() self.ax2.clear() self.ax1.plot(self.epochs, self.loss_list, label='Training Loss') self.ax1.set_title('Training Loss') self.ax1.set_xlabel('Epochs') self.ax1.set_ylabel('Loss') self.ax1.legend() self.ax2.plot(self.epochs, self.image_loss_list, label='Image Loss') self.ax2.set_title('Image Loss') self.ax2.set_xlabel('Epochs') self.ax2.set_ylabel('Loss') self.ax2.legend() self.fig.canvas.flush_events() def train(args): logger = get_logger(args.save_path) preprocess, target_transform = get_transform(args) device = "cuda" if torch.cuda.is_available() else "cpu" # device = "cpu" AnomalyCLIP_parameters = {"Prompt_length": args.n_ctx, "learnabel_text_embedding_depth": args.depth, "learnabel_text_embedding_length": args.t_n_ctx} # model, _ = AnomalyCLIP_lib.load("ViT-L/14@336px", device=device, design_details = AnomalyCLIP_parameters) model, _ = AnomalyCLIP_lib.load("pre-trained models/clip/ViT-B-32.pt", device=device, design_details = AnomalyCLIP_parameters) model.eval() train_data = Dataset_train(root=args.train_data_path, transform=preprocess, target_transform=target_transform, dataset_name = args.dataset) train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True) ########################################################################################## prompt_learner = AnomalyCLIP_PromptLearner(model.to(device), AnomalyCLIP_parameters) prompt_learner.to(device) model.to(device) model.visual.DAPM_replace(DPAM_layer = args.dpam) ########################################################################################## optimizer = torch.optim.Adam(list(prompt_learner.parameters()), lr=args.learning_rate, betas=(0.5, 0.999)) # losses loss_focal = FocalLoss() loss_dice = BinaryDiceLoss() model.eval() prompt_learner.train() # plotter = RealTimePlotter() for epoch in tqdm(range(args.epoch)): model.eval() prompt_learner.train() loss_list = [] image_loss_list = [] for items in tqdm(train_dataloader): image = items['img'].to(device) label = items['anomaly'] gt = items['img_mask'].squeeze().to(device) gt[gt > 0.5] = 1 gt[gt <= 0.5] = 0 with torch.no_grad(): # Apply DPAM to the layer from 6 to 24 # DPAM_layer represents the number of layer refined by DPAM from top to bottom # DPAM_layer = 1, no DPAM is used # DPAM_layer = 20 as default image_features, patch_features = model.encode_image(image, args.features_list, DPAM_layer = args.dpam) image_features = image_features / image_features.norm(dim=-1, keepdim=True) #################################### prompts, tokenized_prompts, compound_prompts_text = prompt_learner(cls_id = None) text_features = model.encode_text_learn(prompts, tokenized_prompts, compound_prompts_text).float() text_features = torch.stack(torch.chunk(text_features, dim = 0, chunks = 2), dim = 1) text_features = text_features/text_features.norm(dim=-1, keepdim=True) # Apply DPAM surgery text_probs = image_features.unsqueeze(1) @ text_features.permute(0, 2, 1) text_probs = text_probs[:, 0, ...]/0.07 image_loss = F.cross_entropy(text_probs.squeeze(), label.long().cuda()) #Process with GPU #image_loss = F.cross_entropy(text_probs.squeeze(), label.long()) #Without GPU and using CPU image_loss_list.append(image_loss.item()) ###################################################################### similarity_map_list = [] # similarity_map_list.append(similarity_map) for idx, patch_feature in enumerate(patch_features): if idx >= args.feature_map_layer[0]: patch_feature = patch_feature/ patch_feature.norm(dim = -1, keepdim = True) similarity, _ = AnomalyCLIP_lib.compute_similarity(patch_feature, text_features[0]) similarity_map = AnomalyCLIP_lib.get_similarity_map(similarity[:, 1:, :], args.image_size).permute(0, 3, 1, 2) similarity_map_list.append(similarity_map) loss = 0 for i in range(len(similarity_map_list)): loss += loss_focal(similarity_map_list[i], gt) loss += loss_dice(similarity_map_list[i][:, 1, :, :], gt) loss += loss_dice(similarity_map_list[i][:, 0, :, :], 1-gt) optimizer.zero_grad() (loss+image_loss).backward() optimizer.step() loss_list.append(loss.item()) # logs if (epoch + 1) % args.print_freq == 0: avg_loss = np.mean(loss_list) avg_image_loss = np.mean(image_loss_list) logger.info('epoch [{}/{}], loss:{:.4f}, image_loss:{:.4f}'.format(epoch + 1, args.epoch, avg_loss, avg_image_loss)) # plotter.update(epoch + 1, avg_loss, avg_image_loss) #Realtime training performance monitoring # save model if (epoch + 1) % args.save_freq == 0: ckp_path = os.path.join(args.save_path, 'epoch_' + str(epoch + 1) + '.pth') torch.save({"prompt_learner": prompt_learner.state_dict(),"epoch":epoch+1}, ckp_path) if __name__ == '__main__': parser = argparse.ArgumentParser("AnomalyCLIP", add_help=True) # Initialize the argument parser # Define the path to the training dataset and model checkpoint saving parser.add_argument("--train_data_path", type=str, default="./data/4inlab", help="train dataset path") parser.add_argument("--save_path", type=str, default='./checkpoint/241122_SP_DPAM_13_518', help='path to save results') # Specify the name of the training dataset parser.add_argument("--dataset", type=str, default='4inlab', help="train dataset name") # Set the depth parameter (Note: "image size" in help may be misleading) parser.add_argument("--depth", type=int, default=9, help="image size") # Set the prompt length and learnable text embedding length for "zero-shot" learning parser.add_argument("--n_ctx", type=int, default=12, help="zero shot") parser.add_argument("--t_n_ctx", type=int, default=4, help="zero shot") # Specify layers from which feature maps will be extracted (can pass multiple values) parser.add_argument("--feature_map_layer", type=int, nargs="+", default=[0, 1, 2, 3], help="zero shot") # List of layers whose features will be used parser.add_argument("--features_list", type=int, nargs="+", default=[6, 12, 18, 24], help="features used") # Setting parameters for training parser.add_argument("--epoch", type=int, default=400, help="epochs") parser.add_argument("--learning_rate", type=float, default=0.0001, help="learning rate") parser.add_argument("--batch_size", type=int, default=8, help="batch size") # Size/depth parameter for the DPAM (Deep Prompt Attention Mechanism) parser.add_argument("--dpam", type=int, default=13, help="dpam size") # Define the size of input images used for training parser.add_argument("--image_size", type=int, default=518, help="image size") # Frequency (in epochs) of logging training information and saving parser.add_argument("--print_freq", type=int, default=1, help="print frequency") parser.add_argument("--save_freq", type=int, default=1, help="save frequency") parser.add_argument("--seed", type=int, default=111, help="random seed") args = parser.parse_args() # Parse the command-line arguments and store them in the 'args' object setup_seed(args.seed) # Set the random seed for reproducibility using the provided seed value train(args) # Call the training function with the parsed arguments