import json import gc import psutil import yaml import wget import json import os, sys from glob import glob from yolo_cam.eigen_cam import EigenCAM from yolo_cam.utils.image import show_cam_on_image import torch from torchvision.utils import make_grid import wandb import cv2 import numpy as np import argparse def parse_args(): parser = argparse.ArgumentParser(description="Transfer learning script.") parser.add_argument("--dataset", type=str, required=True, help='Dataset name to be used') parser.add_argument("--epochs", type=int, default=1000, help="Number of epochs") parser.add_argument("--batch", type=int, default=16, help="Batch size") parser.add_argument("--imgsz", type=int, default=640, help="Image size") parser.add_argument("--patience", type=int, default=30, help="Patience for early stopping") parser.add_argument("--cache", type=str, default="ram", help="Cache option") parser.add_argument("--pretrained", type=bool, default=False, help="Use pretrained weights") parser.add_argument("--cos_lr", type=bool, default=False, help="Use cosine learning rate") parser.add_argument("--profile", type=bool, default=False, help="Profile the training") parser.add_argument("--plots", type=bool, default=True, help="Generate plots") parser.add_argument("--resume", type=bool, default=False, help="Resume run") parser.add_argument("--augment", type=bool, default=False, help="Apply augmentation techniques during training") parser.add_argument("--model", type=str, required=True, help="Model name") parser.add_argument("--run", type=str, required=True, help="Run mode") return parser.parse_args() args = parse_args() dict_to_freeze = {"Finetuning": 0, "freeze_[P1-P3]": 4, "freeze_Backbone": 9, "freeze_[P1-23]": 23 } layers_to_freeze = dict_to_freeze[args.run] ##--------------------------------------------/USER DEFINE mem_torch = 0 RAM_used = 0 RAM_available = 0 grad_norm = [] cam_images = {} first_batch_paths = None args_model = { "epochs": args.epochs, "batch": args.batch, "imgsz": args.imgsz, "patience": args.patience, "cache": args.cache, "pretrained": args.pretrained, "cos_lr": args.cos_lr, "profile": args.profile, "plots": args.plots} def freeze_layer(trainer): model = trainer.model num_freeze = layers_to_freeze print(f"Freezing {num_freeze} layers") if num_freeze: freeze = [f'model.{x}.' for x in range(num_freeze)] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers if any(x in k for x in freeze): print(f'freezing {k}') v.requires_grad = False # Non trainable layer print(f"{num_freeze} layers are freezed.") model.info(detailed=True) def get_gpu_usage(param): global mem_torch global RAM_used global RAM_available mem_torch = float(torch.cuda.memory_reserved() / 1E6 if torch.cuda.is_available() else 0) #(MB) RAM_used = float(psutil.virtual_memory().used / 1e9) RAM_available = float(psutil.virtual_memory().available / 1e9) def compute_gradients_L2_norm(trainer): model = trainer.model global grad_norm temp_grad = 0.0 for _, params in model.named_parameters(): if params.grad is not None: temp_grad += params.grad.data.norm(2).item() ** 2 grad_norm.append(float(temp_grad ** 0.5)) def save_val_images_paths(trainer): global first_batch_paths first_batch_paths = next(iter(trainer.validator.dataloader))['im_file'][:args.batch] def compute_grad_CAM(trainer): global global_step global cam_images global first_batch_paths if ((global_step in [1, 2, 3, 4, 5] or global_step %10 == 0) and global_step <= args.epochs): try: model_copy = YOLO(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/weights/last.pt") except FileNotFoundError: if(global_step!=0): print(f"No such file or directory: /experiment/runs/{args.dataset}/{args.model}/train_{run_name}/weights/last.pt") global_step += 1 return iterator = [("C2fCIB [-2]", -2),("Conv [-4]", -4), ("SPPF", 9), ("PSA", 10)] if "v10" in args.model else [("Conv [-2]", -2),("Conv [-4]", -4), ("SPPF", 9)] json_preds_gradCAM = {layer[0] : None for layer in iterator} for layer_name, layer_index in iterator: cam_images[layer_name] = [] target_layers = [model_copy.model.model[layer_index]] for path_to_image in first_batch_paths: image_name = path_to_image.split("/")[-1].split(".")[0] os.makedirs(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/outs/gradCAMs_inferences/images/{image_name}/{layer_name}", exist_ok=True) img = cv2.cvtColor(cv2.imread(path_to_image), cv2.COLOR_BGR2RGB) #img = cv2.resize(img, (640, 640)) rgb_img = img.copy() img = np.float32(img) / 255 cam = EigenCAM(model_copy, target_layers, task='od') grayscale_cam = cam(rgb_img)[0, :, :] temp_cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True) cam = EigenCAM(model_copy, target_layers, task='od') grayscale_cam = cam(rgb_img)[0, :, :] temp_cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True) if (not json_preds_gradCAM[layer_name]): json_preds_gradCAM[layer_name] = [] boxes = model_copy.predictor.results[0].boxes.xyxy.detach().cpu().numpy() # Get boxes as numpy array confs = model_copy.predictor.results[0].boxes.conf.detach().cpu().numpy() # Get boxes as numpy array cam_image_annotated = temp_cam_image.copy() cv2.imwrite(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/outs/gradCAMs_inferences/images/{image_name}/{layer_name}/{image_name}_gs_{global_step}.jpg", cv2.cvtColor(temp_cam_image, cv2.COLOR_BGR2RGB)) for i, box in enumerate(boxes): cam_image_annotated = cv2.rectangle(cam_image_annotated, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2) # Prepare confidence text conf_text = f"{confs[i]:.2f}" # Choose the position for the text (top-left corner of the rectangle) text_position = (int(box[0]), int(box[1]) - 5) # Slightly above the top-left corner # Get text size (text_width, text_height), baseline = cv2.getTextSize(conf_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) # Draw a filled rectangle as background for the text background_tl = (text_position[0], text_position[1] - text_height - baseline) background_br = (text_position[0] + text_width, text_position[1] + baseline) cam_image_annotated = cv2.rectangle(cam_image_annotated, background_tl, background_br, (0, 255, 0), cv2.FILLED) # Add text to the image cam_image_annotated = cv2.putText(cam_image_annotated, conf_text, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, # Font scale (0, 0, 0), # White color 1, # Thickness cv2.LINE_AA) # Anti-aliased line json_preds_gradCAM[layer_name].append({"image_name": image_name, "cls" : model_copy.predictor.results[0].boxes.cls.detach().clone().tolist(), "conf" : model_copy.predictor.results[0].boxes.conf.detach().clone().tolist(), "boxes(xywhn)" : model_copy.predictor.results[0].boxes.xywhn.detach().clone().tolist(), "orig_shape": model_copy.predictor.results[0].boxes.orig_shape}) cam_images[layer_name].append(torch.from_numpy(cv2.resize(temp_cam_image, (640,640))).permute(2, 0, 1)) cv2.imwrite(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/outs/gradCAMs_inferences/images/{image_name}/{layer_name}/{image_name}_annotated_gs_{global_step}.jpg", cv2.cvtColor(cam_image_annotated, cv2.COLOR_BGR2RGB)) model_copy = None cam = None with open(f'./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/outs/gradCAMs_inferences/gradCAMS_inference_step_{global_step}.json', 'w') as f: # Serializing json json_object = json.dumps(json_preds_gradCAM, indent=4) f.write(json_object) gc.collect() torch.cuda.empty_cache() global_step += 1 def log_data(trainer): global grad_norm global mem_torch global RAM_used global RAM_available global global_step global cam_images grad_norm = np.array(grad_norm) grad_norm[grad_norm == 0] = np.nan grad_norm = np.nanmean(grad_norm, axis=0) if(cam_images and global_step <= args.epochs): if "v10" in args.model: cam_images_C2fCIB = torch.stack(cam_images["C2fCIB [-2]"]) cam_images_PSA = torch.stack(cam_images["PSA"]) cam_images_SPPF = torch.stack(cam_images["SPPF"]) cam_images_Conv = torch.stack(cam_images["Conv [-4]"]) grid_C2fCIB = wandb.Image(make_grid(cam_images_C2fCIB, nrow=int(args.batch/4)).float(), caption="C2fCIB layer") grid_PSA = wandb.Image(make_grid(cam_images_PSA, nrow=int(args.batch/4)).float(), caption="PSA layer") grid_SPPF = wandb.Image(make_grid(cam_images_SPPF, nrow=int(args.batch/4)).float(), caption="SPPF layer") grid_Conv = wandb.Image(make_grid(cam_images_Conv, nrow=int(args.batch/4)).float(), caption="Conv [-2] layer") wandb.log({"Gradients/L2 Gradients Norm": grad_norm, "GPU/GPU usage Ultralytics (MB)": mem_torch, "Memory/Memory used (GB):": RAM_used, "Memory/Memory available (GB):": RAM_available, "GradCAM/C2fCIB": grid_C2fCIB, "GradCAM/SPPF": grid_SPPF, "GradCAM/PSA": grid_PSA, "GradCAM/Conv": grid_Conv, } ) elif "v8" in args.model: cam_images_Conv2 = torch.stack(cam_images["Conv [-2]"]) cam_images_Conv4 = torch.stack(cam_images["Conv [-4]"]) cam_images_SPPF = torch.stack(cam_images["SPPF"]) grid_Conv2 = wandb.Image(make_grid(cam_images_Conv2, nrow=int(args.batch/4)).float(), caption="Conv [-2] layer") grid_Conv4 = wandb.Image(make_grid(cam_images_Conv4, nrow=int(args.batch/4)).float(), caption="Conv [-4] layer") grid_SPPF = wandb.Image(make_grid(cam_images_SPPF, nrow=int(args.batch/4)).float(), caption="SPPF layer") wandb.log({"Gradients/L2 Gradients Norm": grad_norm, "GPU/GPU usage Ultralytics (MB)": mem_torch, "Memory/Memory used (GB):": RAM_used, "Memory/Memory available (GB):": RAM_available, "GradCAM/Conv2": grid_Conv2, "GradCAM/Conv4": grid_Conv4, "GradCAM/SPPF": grid_SPPF, } ) elif(global_step <= args.epochs): wandb.log({"Gradients/L2 Gradients Norm": grad_norm, "GPU/GPU usage Ultralytics (MB)": mem_torch, "Memory/Memory used (GB):": RAM_used, "Memory/Memory available (GB):": RAM_available } ) cam_images = {} grad_norm = [] mem_torch = None RAM_used = None RAM_available = None ultralytics_augmentation_args_disabled = { "hsv_h": 0.0, "hsv_s": 0.0, "hsv_v": 0.0, "degrees": 0.0, "translate": 0.0, "scale": 0.0, # Setting scale to 0.0 keeps the original size "shear": 0.0, "perspective": 0.0, "flipud": 0.0, "fliplr": 0.0, "mosaic": 0.0, "mixup": 0.0, "copy_paste": 0.0, "augment": False, # Setting to 'none' disables auto augmentation } os.makedirs("./experiment/pretrained_weights", exist_ok=True) pretrained_weights_list = [weights_path.split("/")[-1][:-3]for weights_path in glob("./experiment/pretrained_weights/*.pt")] if args.model not in pretrained_weights_list: if "v10" in args.model: wget.download(f"https://github.com/THU-MIG/yolov10/releases/download/v1.1/{args.model}.pt", out="./experiment/pretrained_weights") elif "v8" in args.model: wget.download(f"https://github.com/ultralytics/assets/releases/download/v8.2.0/{args.model}.pt", out="./experiment/pretrained_weights") global global_step global_step = 0 gc.collect() torch.cuda.empty_cache() wandb.login(key="c5c2f8d387804338825114c4133a31016c9ebf87") api = wandb.Api() # Step 1: Initialize a Weights & Biases run run_name = args.model.split('yolo')[-1]+"_"+args.run wandb.init(project=f"transfer_learning_{args.dataset}", dir="./experiment", name=run_name, # run name job_type="training", notes=f"Finetuning of {args.model} model on {args.dataset} dataset", tags=["object detection", "FaRADAI", "AI4TES", "Finetuning", args.dataset], resume="allow") username = wandb.run.entity project = wandb.run.project run_id = wandb.run.id if("From_Scratch" in run_name): flag_pretrained = False args.pretrained = flag_pretrained else: flag_pretrained = True args.pretrained = flag_pretrained # Step 2: Define the YOLO Model if "v10" in args.model: from ultralytics import YOLOv10 as YOLO elif "v8" in args.model: from ultralytics import YOLO if(not args.pretrained): model = YOLO(f"{args.model}.yaml") #Necessary to doesn't initialize pretrained weights elif args.resume: model = YOLO(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/weights/last.pt") else: model = YOLO(f'./experiment/pretrained_weights/{args.model}.pt') # Step 3: Add trainer & validator callbacks model.add_callback("on_train_start", freeze_layer) model.add_callback("on_train_start", save_val_images_paths) model.add_callback("on_batch_end", compute_gradients_L2_norm) #After scheduler step model.add_callback("on_train_epoch_end", get_gpu_usage) model.add_callback("on_train_epoch_end", compute_grad_CAM) model.add_callback("on_train_epoch_end", log_data) if(not args.augment): args_model.update(ultralytics_augmentation_args_disabled) if args.resume: try: model.train(model=f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/weights/last.pt", resume=True) except: print("Exception catched, proceeding with validation:") print(model.info(detailed=True)) pass else: model.train(data=f"./datasets/{args.dataset}/dataset.yaml", project=f"./experiment/runs/{args.dataset}/{args.model}", name=f"train_{run_name}", **args_model) model.data = None model = None # Necessary to Free RAM gc.collect() torch.cuda.empty_cache() with open(f'./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/args.yaml', 'r') as config_file: config = yaml.safe_load(config_file) config['model'] = args.model if (layers_to_freeze!=0): config["freeze"] = {"number_of_layers": layers_to_freeze, "layers": run_name.split("_")[1]} else: config["freeze"] = {"number_of_layers": layers_to_freeze, "layers": "None"} wandb.init(entity=username,dir="./experiment", project=project, id=run_id, resume="must") wandb.config.update(config) model = YOLO(f"./experiment/runs/{args.dataset}/{args.model}/train_{run_name}/weights/best.pt") metrics = model.val(data=f"../datasets/{args.dataset}/dataset.yaml", imgsz=args.imgsz, batch=args.batch, save_json=True, save_txt=True, split='test', plots=True, conf=0.5, iou=0.7, project=f"./experiment/runs/{args.dataset}/{args.model}", name=f"test_{run_name}" ) with open(f"./experiment/runs/{args.dataset}/{args.model}/test_{run_name}/metrics.json", "w") as f: f.write(json.dumps(metrics.results_dict, indent=4)) wandb.log( {"test/precision(B)": metrics.results_dict["metrics/precision(B)"], "test/recall(B)": metrics.results_dict["metrics/recall(B)"], "test/mAP50(B)": metrics.results_dict["metrics/mAP50(B)"], "test/mAP50-95(B)": metrics.results_dict["metrics/mAP50-95(B)"], "test/fitness": metrics.results_dict["fitness"] } ) # Step 7: Finalize the W&B Run wandb.finish() model.data = None model = None # Necessary to Free RAM gc.collect() torch.cuda.empty_cache()