From 6daf202509b897740a0bd287f32c53baccfb8c1f Mon Sep 17 00:00:00 2001 From: Rokas Date: Sat, 4 May 2024 21:45:38 +0300 Subject: [PATCH] Customized Yolov8 training --- CHANGELOG.md | 11 + README.md | 3 +- Tutorials/11_Yolov8/README.md | 176 ++++++++- Tutorials/11_Yolov8/requirements.txt | 2 +- Tutorials/11_Yolov8/run_pretrained.py | 3 +- Tutorials/11_Yolov8/test_yolov8.py | 33 ++ .../11_Yolov8/train_yolov8.py | 7 +- mltu/__init__.py | 2 +- mltu/augmentors.py | 2 + mltu/dataProvider.py | 4 + mltu/torch/yolo/preprocessors.py | 11 +- mltu/torch/yolo/train.py | 348 ------------------ 12 files changed, 240 insertions(+), 362 deletions(-) create mode 100644 Tutorials/11_Yolov8/test_yolov8.py rename mltu/torch/yolo/train_yolo.py => Tutorials/11_Yolov8/train_yolov8.py (95%) delete mode 100644 mltu/torch/yolo/train.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b51785..b9988a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,14 @@ +## [1.2.5] - 2024-05-04 +### Added +- Added exception in `mltu.dataProvider.DataProvider` to raise ValueError when dataset is not iterable +- Added custom training code for YoloV8 object detector: `Tutorials\11_Yolov8\train_yolov8.py` +- Added custom trained inference code for YoloV8 object detector:`Tutorials\11_Yolov8\test_yolov8.py` + +### Changed +- Fixed `RandomElasticTransform` in `mltu.augmentors` to handle elastic transformation not to exceed image boundaries +- Modified `YoloPreprocessor` in `mltu.torch.yolo.preprocessors` to output dictionary with np.arrays istead of lists + + ## [1.2.4] - 2024-03-21 ### Added - Added `RandomElasticTransform` to `mltu.augmentors` to work with `Image` objects diff --git a/README.md b/README.md index f9ac84e..7d3c00e 100644 --- a/README.md +++ b/README.md @@ -26,4 +26,5 @@ Each tutorial has its own requirements.txt file for a specific mltu version. As 8. [Handwriting words recognition with PyTorch](https://pylessons.com/handwriting-recognition-pytorch), code in ```Tutorials\08_handwriting_recognition_torch``` folder; 9. [Transformer training with TensorFlow for Translation task](https://pylessons.com/transformers-training), code in ```Tutorials\09_translation_transformer``` folder; 10. [Speech Recognition in Python | finetune wav2vec2 model for a custom ASR model](https://youtu.be/h6ooEGzjkj0), code in ```Tutorials\10_wav2vec2_torch``` folder; -11. [YOLOv8: Real-Time Object Detection Simplified](https://youtu.be/vegL__weCxY), code in ```Tutorials\11_Yolov8``` folder; \ No newline at end of file +11. [YOLOv8: Real-Time Object Detection Simplified](https://youtu.be/vegL__weCxY), code in ```Tutorials\11_Yolov8``` folder; +12. [YOLOv8: Customizing Object Detector training](https://youtu.be/ysYiV1CbCyY), code in ```Tutorials\11_Yolov8\train_yolov8.py``` folder; \ No newline at end of file diff --git a/Tutorials/11_Yolov8/README.md b/Tutorials/11_Yolov8/README.md index 9efe9e5..76ec0e3 100644 --- a/Tutorials/11_Yolov8/README.md +++ b/Tutorials/11_Yolov8/README.md @@ -1,10 +1,12 @@ # Run Ultralytics YOLOv8 pretrained model -YouTube tutorial link: [YOLOv8: Real-Time Object Detection Simplified](https://youtu.be/vegL__weCxY) +YouTube tutorial link: +- [YOLOv8: Real-Time Object Detection Simplified](https://youtu.be/vegL__weCxY); +- [YOLOv8: Customizing Object Detector training](https://youtu.be/ysYiV1CbCyY); First, I recommend you to install the required packages in a virtual environment: ```bash -mltu==1.2.3 +mltu==1.2.5 ultralytics==8.1.28 torch==2.0.0 torchvision==0.15.1 @@ -134,5 +136,175 @@ while True: break cap.release() +cv2.destroyAllWindows() +``` + +## Customize YoloV8 Object Detector training: +```python +import os +import time +import torch +from mltu.preprocessors import ImageReader +from mltu.annotations.images import CVImage +from mltu.transformers import ImageResizer, ImageShowCV2, ImageNormalizer +from mltu.augmentors import RandomBrightness, RandomRotate, RandomErodeDilate, RandomSharpen, \ + RandomMirror, RandomFlip, RandomGaussianBlur, RandomSaltAndPepper, RandomDropBlock, RandomMosaic, RandomElasticTransform +from mltu.torch.model import Model +from mltu.torch.dataProvider import DataProvider +from mltu.torch.yolo.annotation import VOCAnnotationReader +from mltu.torch.yolo.preprocessors import YoloPreprocessor +from mltu.torch.yolo.loss import v8DetectionLoss +from mltu.torch.yolo.metrics import YoloMetrics +from mltu.torch.yolo.optimizer import build_optimizer, AccumulativeOptimizer +from mltu.torch.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, Model2onnx, WarmupCosineDecay + +from ultralytics.nn.tasks import DetectionModel +from ultralytics.engine.model import Model as BaseModel + +# https://www.kaggle.com/datasets/andrewmvd/car-plate-detection +annotations_path = "Datasets/car-plate-detection/annotations" + +# Create a dataset from the annotations, the dataset is a list of lists where each list contains the [image path, annotation path] +dataset = [[None, os.path.join(annotations_path, f)] for f in os.listdir(annotations_path)] + +# Make sure torch can see GPU device, it is not recommended to train with CPU +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +img_size = 416 +labels = {0: "licence"} + +# Create a data provider for the dataset +data_provider = DataProvider( + dataset=dataset, + skip_validation=True, + batch_size=16, + data_preprocessors=[ + VOCAnnotationReader(labels=labels), + ImageReader(CVImage), + ], + transformers=[ + # ImageShowCV2(), + ImageResizer(img_size, img_size), + ImageNormalizer(transpose_axis=True), + ], + batch_postprocessors=[ + YoloPreprocessor(device, img_size) + ], + numpy=False, +) + +# split the dataset into train and test +train_data_provider, val_data_provider = data_provider.split(0.9, shuffle=False) + +# Attaach augmentation to the train data provider +train_data_provider.augmentors = [ + RandomBrightness(), + RandomErodeDilate(), + RandomSharpen(), + RandomMirror(), + RandomFlip(), + RandomElasticTransform(), + RandomGaussianBlur(), + RandomSaltAndPepper(), + RandomRotate(angle=10), + RandomDropBlock(), + RandomMosaic(), +] + +base_model = BaseModel("yolov8n.pt") +# Create a YOLO model +model = DetectionModel('yolov8n.yaml', nc=len(labels)) + +# Load the weight from base model +try: model.load_state_dict(base_model.model.state_dict(), strict=False) +except: pass + +model.to(device) + +for k, v in model.named_parameters(): + if any(x in k for x in [".dfl"]): + print("freezing", k) + v.requires_grad = False + elif not v.requires_grad: + v.requires_grad = True + +lr = 1e-3 +optimizer = build_optimizer(model.model, name="AdamW", lr=lr, weight_decay=0.0, momentum=0.937, decay=0.0005) +optimizer = AccumulativeOptimizer(optimizer, 16, 64) + +# create model object that will handle training and testing of the network +model = Model( + model, + optimizer, + v8DetectionLoss(model), + metrics=[YoloMetrics(nc=len(labels))], + log_errors=False, + output_path=f"Models/11_Yolov8/{int(time.time())}", + clip_grad_norm=10.0, + ema=True, +) + +modelCheckpoint = ModelCheckpoint(monitor="val_fitness", mode="max", save_best_only=True, verbose=True) +tensorBoard = TensorBoard() +earlyStopping = EarlyStopping(monitor="val_fitness", mode="max", patience=31, verbose=True) +model2onnx = Model2onnx(input_shape=(1, 3, img_size, img_size), verbose=True, opset_version=14, + dynamic_axes = {"input": {0: "batch_size", 2: "height", 3: "width"}, + "output": {0: "batch_size", 2: "anchors"}}, + metadata={"classes": labels}) +warmupCosineDecayBias = WarmupCosineDecay(lr_after_warmup=lr, final_lr=lr, initial_lr=0.1, + warmup_steps=len(train_data_provider), warmup_epochs=10, ignore_param_groups=[1, 2]) # lr0 +warmupCosineDecay = WarmupCosineDecay(lr_after_warmup=lr, final_lr=lr/10, initial_lr=1e-7, + warmup_steps=len(train_data_provider), warmup_epochs=10, decay_epochs=190, ignore_param_groups=[0]) # lr1 and lr2 + +# Train the model +history = model.fit( + train_data_provider, + test_dataProvider=val_data_provider, + epochs=200, + callbacks=[ + modelCheckpoint, + tensorBoard, + earlyStopping, + model2onnx, + warmupCosineDecayBias, + warmupCosineDecay + ] + ) +``` + +## Test Custom trained YoloV8 Object Detector: +```python +import os +import cv2 +from mltu.annotations.detections import Detections +from mltu.torch.yolo.detectors.onnx_detector import Detector as OnnxDetector + +# https://www.kaggle.com/datasets/andrewmvd/car-plate-detection +images_path = "Datasets/car-plate-detection/images" + +input_width, input_height = 416, 416 +confidence_threshold = 0.5 +iou_threshold = 0.5 + +detector = OnnxDetector("Models/11_Yolov8/1714135287/model.onnx", input_width, input_height, confidence_threshold, iou_threshold, force_cpu=False) + +for image_path in os.listdir(images_path): + + frame = cv2.imread(os.path.join(images_path, image_path)) + + # Perform Yolo object detection + detections: Detections = detector(frame) + + # Apply the detections to the frame + frame = detections.applyToFrame(frame) + + # Print the FPS + print(detector.fps) + + # Display the output image + cv2.imshow("Object Detection", frame) + if cv2.waitKey(0) & 0xFF == ord('q'): + break + cv2.destroyAllWindows() ``` \ No newline at end of file diff --git a/Tutorials/11_Yolov8/requirements.txt b/Tutorials/11_Yolov8/requirements.txt index 63a530d..98c16d6 100644 --- a/Tutorials/11_Yolov8/requirements.txt +++ b/Tutorials/11_Yolov8/requirements.txt @@ -1,4 +1,4 @@ -mltu==1.2.3 +mltu==1.2.5 ultralytics==8.1.28 torch==2.0.0 torchvision==0.15.1 diff --git a/Tutorials/11_Yolov8/run_pretrained.py b/Tutorials/11_Yolov8/run_pretrained.py index 110ec4c..b160910 100644 --- a/Tutorials/11_Yolov8/run_pretrained.py +++ b/Tutorials/11_Yolov8/run_pretrained.py @@ -1,5 +1,6 @@ import cv2 from ultralytics.engine.model import Model as BaseModel +from mltu.annotations.detections import Detections from mltu.torch.yolo.detectors.torch_detector import Detector as TorchDetector from mltu.torch.yolo.detectors.onnx_detector import Detector as OnnxDetector @@ -18,7 +19,7 @@ break # Perform Yolo object detection - detections = detector(frame) + detections: Detections = detector(frame) # Apply the detections to the frame frame = detections.applyToFrame(frame) diff --git a/Tutorials/11_Yolov8/test_yolov8.py b/Tutorials/11_Yolov8/test_yolov8.py new file mode 100644 index 0000000..6b22303 --- /dev/null +++ b/Tutorials/11_Yolov8/test_yolov8.py @@ -0,0 +1,33 @@ +import os +import cv2 +from mltu.annotations.detections import Detections +from mltu.torch.yolo.detectors.onnx_detector import Detector as OnnxDetector + +# https://www.kaggle.com/datasets/andrewmvd/car-plate-detection +images_path = "Datasets/car-plate-detection/images" + +input_width, input_height = 416, 416 +confidence_threshold = 0.5 +iou_threshold = 0.5 + +detector = OnnxDetector("Models/11_Yolov8/1714135287/model.onnx", input_width, input_height, confidence_threshold, iou_threshold, force_cpu=False) + +for image_path in os.listdir(images_path): + + frame = cv2.imread(os.path.join(images_path, image_path)) + + # Perform Yolo object detection + detections: Detections = detector(frame) + + # Apply the detections to the frame + frame = detections.applyToFrame(frame) + + # Print the FPS + print(detector.fps) + + # Display the output image + cv2.imshow("Object Detection", frame) + if cv2.waitKey(0) & 0xFF == ord('q'): + break + +cv2.destroyAllWindows() \ No newline at end of file diff --git a/mltu/torch/yolo/train_yolo.py b/Tutorials/11_Yolov8/train_yolov8.py similarity index 95% rename from mltu/torch/yolo/train_yolo.py rename to Tutorials/11_Yolov8/train_yolov8.py index 044537c..2fe9cdf 100644 --- a/mltu/torch/yolo/train_yolo.py +++ b/Tutorials/11_Yolov8/train_yolov8.py @@ -13,12 +13,12 @@ from mltu.torch.yolo.loss import v8DetectionLoss from mltu.torch.yolo.metrics import YoloMetrics from mltu.torch.yolo.optimizer import build_optimizer, AccumulativeOptimizer -from mltu.torch.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard, Model2onnx, WarmupCosineDecay +from mltu.torch.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, Model2onnx, WarmupCosineDecay from ultralytics.nn.tasks import DetectionModel from ultralytics.engine.model import Model as BaseModel - +# https://www.kaggle.com/datasets/andrewmvd/car-plate-detection annotations_path = "Datasets/car-plate-detection/annotations" # Create a dataset from the annotations, the dataset is a list of lists where each list contains the [image path, annotation path] @@ -72,6 +72,7 @@ # Create a YOLO model model = DetectionModel('yolov8n.yaml', nc=len(labels)) +# Load the weight from base model try: model.load_state_dict(base_model.model.state_dict(), strict=False) except: pass @@ -95,7 +96,7 @@ v8DetectionLoss(model), metrics=[YoloMetrics(nc=len(labels))], log_errors=False, - output_path=f"Models/detector/{int(time.time())}", + output_path=f"Models/11_Yolov8/{int(time.time())}", clip_grad_norm=10.0, ema=True, ) diff --git a/mltu/__init__.py b/mltu/__init__.py index 80ae7a5..015215d 100644 --- a/mltu/__init__.py +++ b/mltu/__init__.py @@ -1,4 +1,4 @@ -__version__ = "1.2.4" +__version__ = "1.2.5" from .annotations.images import Image from .annotations.images import CVImage diff --git a/mltu/augmentors.py b/mltu/augmentors.py index 0da0772..7bfb9c0 100644 --- a/mltu/augmentors.py +++ b/mltu/augmentors.py @@ -916,6 +916,8 @@ def __call__(self, image: Image, annotation: typing.Any) -> typing.Tuple[Image, detections = [] for detection in annotation: x_min, y_min, x_max, y_max = detection.xyxy_abs + x_max = min(x_max, dx.shape[1] - 1) + y_max = min(y_max, dy.shape[0] - 1) new_x_min = min(max(0, x_min + dx[y_min, x_min]), image.width - 1) new_y_min = min(max(0, y_min + dy[y_min, x_min]), image.height - 1) new_x_max = min(max(0, x_max + dx[y_max, x_max]), image.width - 1) diff --git a/mltu/dataProvider.py b/mltu/dataProvider.py index 994d8a2..06f708d 100644 --- a/mltu/dataProvider.py +++ b/mltu/dataProvider.py @@ -70,6 +70,10 @@ def __init__( else: self.logger.info("Skipping Dataset validation...") + # Check if dataset has length + if not len(dataset): + raise ValueError("Dataset must be iterable") + if limit: self.logger.info(f"Limiting dataset to {limit} samples.") self._dataset = self._dataset[:limit] diff --git a/mltu/torch/yolo/preprocessors.py b/mltu/torch/yolo/preprocessors.py index 7518c4a..87d1981 100644 --- a/mltu/torch/yolo/preprocessors.py +++ b/mltu/torch/yolo/preprocessors.py @@ -1,12 +1,13 @@ import torch +import typing import numpy as np class YoloPreprocessor: - def __init__(self, device, imgsz=640): + def __init__(self, device: torch.device, imgsz: int=640): self.device = device self.imgsz = imgsz - def __call__(self, images, annotations): + def __call__(self, images, annotations) -> typing.Tuple[np.ndarray, dict]: batch = { "ori_shape": [], "resized_shape": [], @@ -23,8 +24,8 @@ def __call__(self, images, annotations): batch["bboxes"].append(detection.xywh) batch["batch_idx"].append(i) - batch["cls"] = torch.tensor(batch["cls"]).to(self.device) - batch["bboxes"] = torch.tensor(batch["bboxes"]).to(self.device) - batch["batch_idx"] = torch.tensor(batch["batch_idx"]).to(self.device) + batch["cls"] = torch.tensor(np.array(batch["cls"])).to(self.device) + batch["bboxes"] = torch.tensor(np.array(batch["bboxes"])).to(self.device) + batch["batch_idx"] = torch.tensor(np.array(batch["batch_idx"])).to(self.device) return np.array(images), batch \ No newline at end of file diff --git a/mltu/torch/yolo/train.py b/mltu/torch/yolo/train.py deleted file mode 100644 index eb695fe..0000000 --- a/mltu/torch/yolo/train.py +++ /dev/null @@ -1,348 +0,0 @@ -import os -import cv2 -import time -import torch -import numpy as np -from tqdm import tqdm -from pathlib import Path -from ultralytics.nn.tasks import DetectionModel -# from ultralytics.data.dataset import YOLODataset -from mltu.torch.dataProvider import DataProvider -from mltu.preprocessors import ImageReader -from mltu.annotations.images import CVImage -from mltu.torch.model import Model -# model = YOLO('yolov8n.pt') -from metrics import YoloMetrics -from mltu.transformers import ImageResizer, ImageShowCV2 -from mltu.annotations.detections import Detections, Detection -from mltu.augmentors import RandomBrightness, RandomRotate, RandomErodeDilate, RandomSharpen, \ - RandomMirror, RandomFlip, RandomGaussianBlur, RandomSaltAndPepper, RandomDropBlock, RandomMosaic -import torch.nn as nn - -from mltu.torch.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard, Model2onnx, WarmupCosineDecay -# from loss import v8DetectionLoss - - -train_path = "/home/rokbal/Personal/ultralytics/generated_data/train/images" -train_labels = "/home/rokbal/Personal/ultralytics/generated_data/train/labels" - -val_path = "/home/rokbal/Personal/ultralytics/generated_data/valid/images" -val_labels = "/home/rokbal/Personal/ultralytics/generated_data/valid/labels" - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -dataset = [] -for f in tqdm(os.listdir(train_path)): - img_path = Path(os.path.join(train_path, f)) - txt_path = Path(os.path.join(train_labels, img_path.stem + ".txt")) - if txt_path.exists(): - dataset.append([str(img_path), str(txt_path)]) - -val_dataset = [] -for f in tqdm(os.listdir(val_path)): - img_path = Path(os.path.join(val_path, f)) - txt_path = Path(os.path.join(val_labels, img_path.stem + ".txt")) - if txt_path.exists(): - val_dataset.append([str(img_path), str(txt_path)]) - -def annotation_preprocessor(image, annotation): - with open(annotation, "r") as f: - data = f.readlines() - - dets = [Detection( - bbox=[float(x) for x in line.split()[1:]], - label="ID", - confidence=1.0, - relative=True - ) for line in data] - - detections = Detections( - labels={0: "ID"}, - width=image.width, - height=image.height, - image_path=image.path, - detections=dets - ) - - return image, detections - - -def image_normalizer(image, annotation): - img = image.numpy().transpose(2, 0, 1) / 255.0 - return img, annotation - - -class YoloPreprocessor: - def __init__(self, imgsz=640): - self.imgsz = imgsz - - def __call__(self, images, annotations): - batch = { - "ori_shape": [], - "resized_shape": [], - "cls": [], - "bboxes": [], - "batch_idx": [], - } - - for i, (image, detections) in enumerate(zip(images, annotations)): - batch["ori_shape"].append([detections.height, detections.width]) - batch["resized_shape"].append([self.imgsz, self.imgsz]) - for detection in detections: - batch["cls"].append([detection.labelId]) - batch["bboxes"].append(detection.xywh) - batch["batch_idx"].append(i) - - batch["cls"] = torch.tensor(batch["cls"]).to(device) - batch["bboxes"] = torch.tensor(batch["bboxes"]).to(device) - batch["batch_idx"] = torch.tensor(batch["batch_idx"]).to(device) - - return np.array(images), batch - - - - - -img_size = 320 -# Create a data provider for the dataset -train_data_provider = DataProvider( - dataset=dataset, - skip_validation=True, - batch_size=16, - data_preprocessors=[ - ImageReader(CVImage), - annotation_preprocessor, - ], - transformers=[ - ImageShowCV2(), - ImageResizer(img_size, img_size), - image_normalizer, - ], - use_cache=True, - batch_postprocessors=[ - YoloPreprocessor(img_size) - ], - use_multiprocessing=False, - max_queue_size=1, - workers=1, - numpy=False, - augmentors=[ - # RandomBrightness(), - # RandomErodeDilate(), - # RandomSharpen(), - # RandomMirror(0.5), - # RandomFlip(0.5), - # RandomGaussianBlur(), - # RandomSaltAndPepper(), - # RandomRotate(0.5, angle=10), - # RandomDropBlock(0.5), - RandomMosaic(random_chance=1.0) - ] -) - -for d in tqdm(train_data_provider): - pass - - -# train_data_provider, val_data_provider = data_provider.split(0.8) - -# Augment training data with random brightness, rotation and erode/dilate -train_data_provider.augmentors = [ - RandomBrightness(), - RandomErodeDilate(), - RandomSharpen(), - RandomMirror(0.5), - RandomFlip(0.5), - RandomGaussianBlur(), - RandomSaltAndPepper(), - - # RandomRotate(angle=10), - # ImageShowCV2(), -] - -val_data_provider = DataProvider( - dataset=val_dataset, - skip_validation=True, - batch_size=16, - data_preprocessors=[ - ImageReader(CVImage), - annotation_preprocessor, - ], - transformers=[ - # ImageShowCV2(), - ImageResizer(img_size, img_size), - image_normalizer, - ], - use_cache=True, - batch_postprocessors=[ - YoloPreprocessor(img_size) - ], - use_multiprocessing=False, - max_queue_size=32, - # workers=32, - numpy=False, -) - - -# while True: -# start = time.time() -# for batch in tqdm(train_data_provider): -# pass -# end = time.time() -# fps = len(data_provider) / (end - start) -# print(f"FPS: {fps:.2f}") - -# for i in range(1000): -# pbar = tqdm(train_data_provider) -# for d in pbar: -# pass - - -# # model_y = YOLO('yolov8n.yaml') -# model = DetectionModel('yolov8n.yaml', nc=1) # YOLO('yolov8n.pt') -# model.to(device) -# model.load_state_dict(torch.load("Models/detector/1707665153/model.pt")) - -# for k, v in model.named_parameters(): -# if any(x in k for x in [".dfl"]): -# print("freezing", k) -# v.requires_grad = False -# elif not v.requires_grad: -# v.requires_grad = True - -# # model.model.nc = 1 -# # optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.937, 0.999)) - -# pg0, pg1, pg2 = [], [], [] # optimizer parameter groups -# for k, v in model.named_modules(): -# if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): -# pg2.append(v.bias) # biases -# if isinstance(v, nn.BatchNorm2d): -# pg0.append(v.weight) # no decay -# elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): -# pg1.append(v.weight) # apply decay - -# optimizer = torch.optim.AdamW(pg0, lr=1e-2, betas=(0.937, 0.999)) -# optimizer.add_param_group({'params': pg1, 'weight_decay': 0.0005}) # add pg1 with weight_decay -# optimizer.add_param_group({'params': pg2, 'weight_decay': 0.0}) # add pg2 (biases) -# # del pg0, pg1, pg2 - - - -# loss = v8DetectionLoss(model) -# # create model object that will handle training and testing of the network -# model = Model( -# model, -# optimizer, -# loss, -# metrics=[YoloMetrics(nc=1)], -# log_errors=False, -# mixed_precision=False, -# output_path=f"Models/detector/{int(time.time())}", -# clip_grad_norm=10.0, -# ema=True, -# ) - -# # ouput_path = f"Models/detector/{int(time.time())}" -# # save_model_path = ouput_path + "/model.pt" -# # modelCheckpoint = ModelCheckpoint(save_model_path, monitor="val_fitness", mode="max", save_best_only=True, verbose=True) -# # tensorBoard = TensorBoard(log_dir=ouput_path) -# # earlyStopping = EarlyStopping(monitor="val_fitness", mode="max", patience=21, verbose=True) -# # reduceLROnPlateau = ReduceLROnPlateau(monitor="val_fitness", factor=0.9, mode="max", patience=5, verbose=True) -# # model2onnx = Model2onnx(modelCheckpoint.filepath, input_shape=(1, 3, img_size, img_size), verbose=True, opset_version=14, dynamic_axes=None) -# # warmupCosineDecay = WarmupCosineDecay(lr_after_warmup=1e-3, final_lr=1e-3, warmup_epochs=5, decay_epochs=0, initial_lr=1e-5) - -# # history = model.fit( -# # train_data_provider, -# # test_dataProvider=val_data_provider, -# # epochs=1, -# # callbacks=[modelCheckpoint, tensorBoard, earlyStopping, model2onnx, warmupCosineDecay] -# # ) -# # for _ in data_provider: -# # pass - - -# # data = { -# # "train": train_path, -# # } - -# # dataset = YOLODataset( -# # img_path=train_path, -# # imgsz=320, -# # batch_size=4, -# # ) - - -# import math -# from mltu.torch.yolo.tolov8_utils import replace_c2f_with_c2f_v2 -# import torch_pruning as tp -# from ultralytics.nn.modules import Detect - - -# example_inputs = torch.randn(1, 3, img_size, img_size).to(model.device) -# macs_list, nparams_list, map_list, pruned_map_list = [], [], [], [] -# replace_c2f_with_c2f_v2(model.model) -# model.model = model.model.to(model.device) -# model.model.train() -# base_macs, base_nparams = tp.utils.count_ops_and_params(model.model, example_inputs) -# macs_list.append(base_macs) - -# # prune same ratio of filter based on initial size -# target_prune_rate = 0.5 -# iterative_steps = 16 -# pruning_ratio = 1 - math.pow((1 - target_prune_rate), 1 / iterative_steps) - -# init_map = 0 -# print(f"Before Pruning: MACs={base_macs / 1e9: .5f} G, #Params={base_nparams / 1e6: .5f} M, mAP={init_map: .5f}") - -# ignored_layers = [m for m in model.model.modules() if isinstance(m, (Detect,))] -# output_path = model.output_path -# for i in range(iterative_steps): - -# model.model.train() -# for name, param in model.model.named_parameters(): -# param.requires_grad = True - - -# example_inputs = example_inputs.to(model.device) -# pruner = tp.pruner.GroupNormPruner( -# model.model, -# example_inputs, -# importance=tp.importance.GroupNormImportance(), # L2 norm pruning, -# iterative_steps=1, -# pruning_ratio=pruning_ratio, -# ignored_layers=ignored_layers, -# unwrapped_parameters=[] -# ) - -# pruner.step() - -# pruned_macs, pruned_nparams = tp.utils.count_ops_and_params(pruner.model, example_inputs.to(model.device)) -# macs_list.append(pruned_macs) -# current_speed_up = float(macs_list[0]) / pruned_macs -# print(f"After pruning iter {i + 1}: MACs={pruned_macs / 1e9} G, #Params={pruned_nparams / 1e6} M, speed up={current_speed_up}") - -# # os.makedirs(ouput_path, exist_ok=True) -# #save_model_path = ouput_path + f"/model_prune{i}.pth" -# model.output_path = Path(os.path.join(output_path.parent, f"prune_{i}", output_path.name)) -# history = model.fit( -# train_data_provider, -# test_dataProvider=val_data_provider, -# epochs=10, -# callbacks=[ -# TensorBoard(), -# ModelCheckpoint(monitor="val_fitness", mode="max", save_best_only=True, verbose=True), -# Model2onnx(input_shape=(1, 3, img_size, img_size), verbose=True, opset_version=14, dynamic_axes=None), -# # WarmupCosineDecay(lr_after_warmup=1e-3, final_lr=1e-4, warmup_epochs=3, decay_epochs=7, initial_lr=1e-5) -# #modelCheckpoint, tensorBoard, earlyStopping, model2onnx, warmupCosineDecay -# ], -# ) - -# # reset model optimizer -# # model.optimizer = torch.optim.AdamW(model.model.parameters(), lr=1e-3, betas=(0.937, 0.999)) -# # model.optimizer.add_param_group({'params': pg1, 'weight_decay': 0.0005}) # add pg1 with weight_decay -# # model.optimizer.add_param_group({'params': pg2, 'weight_decay': 0.0}) # add pg2 (biases) - -# # torch.save(model.model, save_model_path) -# model.model.to(model.device) -# del pruner \ No newline at end of file