From ddda7144bef41be4efd57f60b47d651e7b949f57 Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:16:24 -0500 Subject: [PATCH 01/24] feat: Basic support for kwcoco files --- yolo/tools/data_loader.py | 112 +++++++++++++++++++++++++----------- yolo/utils/dataset_utils.py | 8 +++ 2 files changed, 88 insertions(+), 32 deletions(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index c44f00c68..a9c4a648f 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -81,47 +81,95 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor. """ images_path = dataset_path / "images" / phase_name + labels_path, data_type = locate_label_paths(dataset_path, phase_name) - images_list = sorted([p.name for p in Path(images_path).iterdir() if p.is_file()]) - if data_type == "json": - annotations_index, image_info_dict = create_image_metadata(labels_path) - data = [] - valid_inputs = 0 - for image_name in track(images_list, description="Filtering data"): - if not image_name.lower().endswith((".jpg", ".jpeg", ".png")): - continue - image_id = Path(image_name).stem + if data_type == 'kwcoco': + """ + More robust data handling that only depends on paths within the + specified manifest file. - if data_type == "json": - image_info = image_info_dict.get(image_id, None) - if image_info is None: - continue - annotations = annotations_index.get(image_info["id"], []) - image_seg_annotations = scale_segmentation(annotations, image_info) - elif data_type == "txt": - label_path = labels_path / f"{image_id}.txt" - if not label_path.is_file(): - continue - with open(label_path, "r") as file: - image_seg_annotations = [list(map(float, line.strip().split())) for line in file] - else: - image_seg_annotations = [] + Principles: - labels = self.load_valid_labels(image_id, image_seg_annotations) + * Dont glob for the images, let the dataset tell you where they are. + + * A Dataset should be referenced as a single URI to a manifest. + The manifest should either contain relevant data or point to + paths for everything. + """ + import kwcoco + coco_dset = kwcoco.CocoDataset(labels_path) + + total_images = coco_dset.n_images - img_path = images_path / image_name if sort_image: - with Image.open(img_path) as img: - width, height = img.size - else: - width, height = 0, 1 - data.append((img_path, labels, width / height)) - valid_inputs += 1 + # Ensure all images have populated sizes + coco_dset._ensure_imgsize() + + ALLOW_EMPTY_IMAGES = 0 + + # Build the expected output + data = [] + valid_inputs = 0 + for coco_img in coco_dset.images().coco_images_iter(): + image_info = coco_img.img + img_path = coco_img.primary_image_filepath() + + if sort_image: + width, height = coco_img['width'], coco_img['height'] + else: + width, height = 0, 1 + + annotations = coco_img.annots().objs + if ALLOW_EMPTY_IMAGES or len(annotations): + image_seg_annotations = scale_segmentation(annotations, image_info) + labels = self.load_valid_labels(None, image_seg_annotations) + + data.append((img_path, labels, width / height)) + valid_inputs += 1 + + else: + images_list = sorted([p.name for p in Path(images_path).iterdir() if p.is_file()]) + if data_type == "json": + annotations_index, image_info_dict = create_image_metadata(labels_path) + + data = [] + valid_inputs = 0 + for image_name in track(images_list, description="Filtering data"): + if not image_name.lower().endswith((".jpg", ".jpeg", ".png")): + continue + image_id = Path(image_name).stem + + if data_type == "json": + image_info = image_info_dict.get(image_id, None) + if image_info is None: + continue + annotations = annotations_index.get(image_info["id"], []) + image_seg_annotations = scale_segmentation(annotations, image_info) + elif data_type == "txt": + label_path = labels_path / f"{image_id}.txt" + if not label_path.is_file(): + continue + with open(label_path, "r") as file: + image_seg_annotations = [list(map(float, line.strip().split())) for line in file] + else: + image_seg_annotations = [] + + labels = self.load_valid_labels(image_id, image_seg_annotations) + + img_path = images_path / image_name + if sort_image: + with Image.open(img_path) as img: + width, height = img.size + else: + width, height = 0, 1 + data.append((img_path, labels, width / height)) + valid_inputs += 1 + total_images = len(images_list) data = sorted(data, key=lambda x: x[2], reverse=True) - logger.info(f"Recorded {valid_inputs}/{len(images_list)} valid inputs") + logger.info(f"Recorded {valid_inputs}/{total_images} valid inputs") return data def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Tensor, None]: diff --git a/yolo/utils/dataset_utils.py b/yolo/utils/dataset_utils.py index dd9a66abc..98699e617 100644 --- a/yolo/utils/dataset_utils.py +++ b/yolo/utils/dataset_utils.py @@ -34,6 +34,14 @@ def locate_label_paths(dataset_path: Path, phase_name: Path) -> Tuple[Path, Path if txt_files: return txt_labels_path, "txt" + HANDLE_KWCOCO_FILES = 1 + if HANDLE_KWCOCO_FILES: + candidate = dataset_path / phase_name + if candidate.is_file(): + labels_path = dataset_path / phase_name + data_type = 'kwcoco' + return labels_path, data_type + logger.warning("No labels found in the specified dataset path and phase name.") return [], None From 0e0de752f540293442d50875ae8e82fccb11f55b Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:16:55 -0500 Subject: [PATCH 02/24] refactor: cleanup code golf --- yolo/tools/data_loader.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index a9c4a648f..0993ba751 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -36,7 +36,10 @@ def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()] self.transform = AugmentationComposer(transforms, self.image_size, self.base_size) self.transform.get_more_data = self.get_more_data - self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name)) + + dataset_path = Path(dataset_cfg.path) + data = self.load_data(dataset_path, phase_name) + self.img_paths, self.bboxes, self.ratios = tensorlize(data) def load_data(self, dataset_path: Path, phase_name: str): """ From 087cd8b27330653dc2bc2985c6896d93f79437a9 Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:17:37 -0500 Subject: [PATCH 03/24] change: disable determinism by default --- yolo/lazy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/lazy.py b/yolo/lazy.py index 0f1cc55b6..7c617c328 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -25,7 +25,7 @@ def main(cfg: Config): log_every_n_steps=1, gradient_clip_val=10, gradient_clip_algorithm="value", - deterministic=True, + # deterministic=True, enable_progress_bar=not getattr(cfg, "quite", False), default_root_dir=save_path, ) From 4133c08a0c2c5f4dcd30964267d343304185b994 Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:21:34 -0500 Subject: [PATCH 04/24] docs: add fixme note --- yolo/lazy.py | 1 + 1 file changed, 1 insertion(+) diff --git a/yolo/lazy.py b/yolo/lazy.py index 7c617c328..8bcd1a020 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -4,6 +4,7 @@ import hydra from lightning import Trainer +# FIXME: messing with sys.path is a bad idea. Factor this out. project_root = Path(__file__).resolve().parent.parent sys.path.append(str(project_root)) From 41f207692eead7eb276dd17ccd335123f4a6b9ae Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:22:00 -0500 Subject: [PATCH 05/24] change: other deterministic disable --- yolo/utils/logging_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/utils/logging_utils.py b/yolo/utils/logging_utils.py index 28a536222..4234ec35b 100644 --- a/yolo/utils/logging_utils.py +++ b/yolo/utils/logging_utils.py @@ -48,7 +48,7 @@ def set_seed(seed): if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. - torch.backends.cudnn.deterministic = True + # torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False From 41c8a66de2d33128fd4028858b76f948d9090123 Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:24:42 -0500 Subject: [PATCH 06/24] refactor: Remove import *, and use getattr to avoid an unsafe eval --- yolo/tools/data_loader.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index 0993ba751..e638e8790 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -12,7 +12,7 @@ from torch.utils.data import DataLoader, Dataset from yolo.config.config import DataConfig, DatasetConfig -from yolo.tools.data_augmentation import * +from yolo.tools import data_augmentation from yolo.tools.data_augmentation import AugmentationComposer from yolo.tools.dataset_preparation import prepare_dataset from yolo.utils.dataset_utils import ( @@ -33,13 +33,10 @@ def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str self.dynamic_shape = getattr(data_cfg, "dynamic_shape", False) self.base_size = mean(self.image_size) - transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()] + transforms = [getattr(data_augmentation, aug)(prob) for aug, prob in augment_cfg.items()] self.transform = AugmentationComposer(transforms, self.image_size, self.base_size) self.transform.get_more_data = self.get_more_data - - dataset_path = Path(dataset_cfg.path) - data = self.load_data(dataset_path, phase_name) - self.img_paths, self.bboxes, self.ratios = tensorlize(data) + self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name)) def load_data(self, dataset_path: Path, phase_name: str): """ From f0c92c4ba581b80dfcf3fc4f623c4b6642c2bcbf Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:25:45 -0500 Subject: [PATCH 07/24] fix: handle case where classes is not in epoch_metrics --- yolo/tools/solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/tools/solver.py b/yolo/tools/solver.py index 8246a66d2..f631ab3fa 100644 --- a/yolo/tools/solver.py +++ b/yolo/tools/solver.py @@ -55,7 +55,7 @@ def validation_step(self, batch, batch_idx): def on_validation_epoch_end(self): epoch_metrics = self.metric.compute() - del epoch_metrics["classes"] + epoch_metrics.pop("classes", None) self.log_dict(epoch_metrics, prog_bar=True, sync_dist=True, rank_zero_only=True) self.log_dict( {"PyCOCO/AP @ .5:.95": epoch_metrics["map"], "PyCOCO/AP @ .5": epoch_metrics["map_50"]}, From ef948a688cd12f40c48a4819e06556d7c2074817 Mon Sep 17 00:00:00 2001 From: joncrall Date: Sun, 5 Jan 2025 13:28:39 -0500 Subject: [PATCH 08/24] fix: error when v_num is not in the loss dict --- yolo/utils/logging_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/utils/logging_utils.py b/yolo/utils/logging_utils.py index 4234ec35b..42e0ac4ff 100644 --- a/yolo/utils/logging_utils.py +++ b/yolo/utils/logging_utils.py @@ -107,7 +107,7 @@ def on_train_batch_end(self, trainer, pl_module, outputs, batch: Any, batch_idx: epoch_descript = "[cyan]Train [white]|" batch_descript = "[green]Train [white]|" metrics = self.get_metrics(trainer, pl_module) - metrics.pop("v_num") + metrics.pop("v_num", None) for metrics_name, metrics_val in metrics.items(): if "Loss_step" in metrics_name: epoch_descript += f"{metrics_name.removesuffix('_step').split('/')[1]: ^9}|" From 9de3de6c425c16f7444e749645e0a5acc675b18c Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 14:03:57 -0500 Subject: [PATCH 09/24] lint: remove unused f-string --- yolo/model/yolo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/model/yolo.py b/yolo/model/yolo.py index cc9ce20ba..ccf456748 100644 --- a/yolo/model/yolo.py +++ b/yolo/model/yolo.py @@ -32,7 +32,7 @@ def __init__(self, model_cfg: ModelConfig, class_num: int = 80): def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]): self.layer_index = {} output_dim, layer_idx = [3], 1 - logger.info(f":tractor: Building YOLO") + logger.info(":tractor: Building YOLO") for arch_name in model_arch: if model_arch[arch_name]: logger.info(f" :building_construction: Building {arch_name}") From 8733d4932225b557bbf9757a793a87f7859bde64 Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 14:43:41 -0500 Subject: [PATCH 10/24] fix: type error in main --- yolo/lazy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/yolo/lazy.py b/yolo/lazy.py index 8bcd1a020..d6091ebbd 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -3,6 +3,7 @@ import hydra from lightning import Trainer +from omegaconf.dictconfig import DictConfig # FIXME: messing with sys.path is a bad idea. Factor this out. project_root = Path(__file__).resolve().parent.parent @@ -14,7 +15,7 @@ @hydra.main(config_path="config", config_name="config", version_base=None) -def main(cfg: Config): +def main(cfg: DictConfig): callbacks, loggers, save_path = setup(cfg) trainer = Trainer( From 1338bd1c77c741d50551e0b40be0b744c589288c Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 16:02:46 -0500 Subject: [PATCH 11/24] test: add doctest for DualLoss with helper config_utils --- yolo/tools/loss_functions.py | 15 +++++++++++++++ yolo/utils/config_utils.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) create mode 100644 yolo/utils/config_utils.py diff --git a/yolo/tools/loss_functions.py b/yolo/tools/loss_functions.py index 79fe1cf96..379a42ce1 100644 --- a/yolo/tools/loss_functions.py +++ b/yolo/tools/loss_functions.py @@ -107,6 +107,21 @@ def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Ten class DualLoss: + """ + Example: + >>> import torch + >>> from yolo.tools.loss_functions import DualLoss + >>> from yolo.utils.bounding_box_utils import Vec2Box + >>> from yolo.utils.config_utils import build_config + >>> cfg = build_config(overrides=['task=train']) + >>> device = 'cpu' + >>> vec2box = Vec2Box(model=None, anchor_cfg=cfg.model.anchor, image_size=cfg.image_size, device=device) + >>> self = DualLoss(cfg, vec2box) + >>> targets = torch.zeros(1, 20, 5, device=device) + >>> aux_predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]] + >>> main_predicts = [torch.zeros(1, 8400, *cn, device=device) for cn in [(80,), (4, 16), (4,)]] + >>> loss, loss_dict = self(aux_predicts, main_predicts, targets) + """ def __init__(self, cfg: Config, vec2box) -> None: loss_cfg = cfg.task.loss self.loss = YOLOLoss(loss_cfg, vec2box, class_num=cfg.dataset.class_num, reg_max=cfg.model.anchor.reg_max) diff --git a/yolo/utils/config_utils.py b/yolo/utils/config_utils.py new file mode 100644 index 000000000..465373c4e --- /dev/null +++ b/yolo/utils/config_utils.py @@ -0,0 +1,36 @@ +import omegaconf +from typing import List + + +def build_config(overrides: List[str] = []) -> omegaconf.DictConfig: + """ + Creates an explicit config for testing. + + Example: + >>> from yolo.utils.config_utils import build_config + >>> cfg = build_config(overrides=['task=train']) + >>> cfg = build_config(overrides=['task=validation']) + >>> cfg = build_config(overrides=['task=inference']) + """ + import yolo + import os + import pathlib + from hydra import compose, initialize + + # This is annoying that we cant just specify an absolute path when it is + # robustly built. Furthermore, the relative path seems like it isn't even + # from the cwd, but the module that is currently being run. + + # Find the path that we need to be relative to in a somewhat portable + # manner (i.e. will work in a Jupyter snippet). + try: + path_base = pathlib.Path(__file__).parent + except NameError: + path_base = pathlib.Path.cwd() + yolo_path = pathlib.Path(yolo.__file__).parent + rel_yolo_path = pathlib.Path(os.path.relpath(yolo_path, path_base)) + # rel_yolo_path = yolo_path.relative_to(path_base, walk_up=True) # requires Python 3.12 + config_path = os.fspath(rel_yolo_path / 'config') + with initialize(config_path=config_path, version_base=None): + cfg = compose(config_name="config", overrides=overrides) + return cfg From 549ca26d882dfa1755b8a579b8fdfd3e9ca17d8a Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 16:33:47 -0500 Subject: [PATCH 12/24] feat: add lazy imports for faster startup time --- requirements.txt | 1 + yolo/__init__.py | 109 +++++++++++++++++++++++++++++------------ yolo/lazy.py | 8 ++- yolo/model/__init__.py | 0 4 files changed, 81 insertions(+), 37 deletions(-) create mode 100644 yolo/model/__init__.py diff --git a/requirements.txt b/requirements.txt index f6d336cbb..5651e7e0b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ einops faster-coco-eval graphviz hydra-core +lazy-loader lightning loguru numpy diff --git a/yolo/__init__.py b/yolo/__init__.py index b4b98d7f8..77da671a6 100644 --- a/yolo/__init__.py +++ b/yolo/__init__.py @@ -1,33 +1,78 @@ -from yolo.config.config import Config, NMSConfig -from yolo.model.yolo import create_model -from yolo.tools.data_loader import AugmentationComposer, create_dataloader -from yolo.tools.drawer import draw_bboxes -from yolo.tools.solver import TrainModel -from yolo.utils.bounding_box_utils import Anc2Box, Vec2Box, bbox_nms, create_converter -from yolo.utils.deploy_utils import FastModelLoader -from yolo.utils.logging_utils import ( - ImageLogger, - YOLORichModelSummary, - YOLORichProgressBar, +""" +The MIT YOLO rewrite +""" + +__autogen__ = """ +mkinit ~/code/YOLO-v9/yolo/__init__.py --nomods --write --lazy-loader + +# Check to see how long it takes to run a simple help command +time python -m yolo.lazy --help +""" + +__submodules__ = { + 'config.config': ['Config', 'NMSConfig'], + 'model.yolo': ['create_model'], + 'tools.data_loader': ['AugmentationComposer', 'create_dataloader'], + 'tools.drawer': ['draw_bboxes'], + 'tools.solver': ['TrainModel'], + 'utils.bounding_box_utils': ['Anc2Box', 'Vec2Box', 'bbox_nms', 'create_converter'], + 'utils.deploy_utils': ['FastModelLoader'], + 'utils.logging_utils': [ + 'ImageLogger', 'YOLORichModelSummary', + 'YOLORichProgressBar', + 'validate_log_directory' + ], + 'utils.model_utils': ['PostProcess'], +} + + +import lazy_loader + + +__getattr__, __dir__, __all__ = lazy_loader.attach( + __name__, + submodules={}, + submod_attrs={ + 'config.config': [ + 'Config', + 'NMSConfig', + ], + 'model.yolo': [ + 'create_model', + ], + 'tools.data_loader': [ + 'AugmentationComposer', + 'create_dataloader', + ], + 'tools.drawer': [ + 'draw_bboxes', + ], + 'tools.solver': [ + 'TrainModel', + ], + 'utils.bounding_box_utils': [ + 'Anc2Box', + 'Vec2Box', + 'bbox_nms', + 'create_converter', + ], + 'utils.deploy_utils': [ + 'FastModelLoader', + ], + 'utils.logging_utils': [ + 'ImageLogger', + 'YOLORichModelSummary', + 'YOLORichProgressBar', + 'validate_log_directory', + ], + 'utils.model_utils': [ + 'PostProcess', + ], + }, ) -from yolo.utils.model_utils import PostProcess - -all = [ - "create_model", - "Config", - "YOLORichProgressBar", - "NMSConfig", - "YOLORichModelSummary", - "validate_log_directory", - "draw_bboxes", - "Vec2Box", - "Anc2Box", - "bbox_nms", - "create_converter", - "AugmentationComposer", - "ImageLogger", - "create_dataloader", - "FastModelLoader", - "TrainModel", - "PostProcess", -] + +__all__ = ['Anc2Box', 'AugmentationComposer', 'Config', 'FastModelLoader', + 'ImageLogger', 'NMSConfig', 'PostProcess', 'TrainModel', 'Vec2Box', + 'YOLORichModelSummary', 'YOLORichProgressBar', 'bbox_nms', + 'create_converter', 'create_dataloader', 'create_model', + 'draw_bboxes', 'validate_log_directory'] diff --git a/yolo/lazy.py b/yolo/lazy.py index d6091ebbd..a44268c9a 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -2,22 +2,20 @@ from pathlib import Path import hydra -from lightning import Trainer from omegaconf.dictconfig import DictConfig # FIXME: messing with sys.path is a bad idea. Factor this out. project_root = Path(__file__).resolve().parent.parent sys.path.append(str(project_root)) -from yolo.config.config import Config -from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel -from yolo.utils.logging_utils import setup - @hydra.main(config_path="config", config_name="config", version_base=None) def main(cfg: DictConfig): + from yolo.utils.logging_utils import setup callbacks, loggers, save_path = setup(cfg) + from lightning import Trainer + from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel trainer = Trainer( accelerator="auto", max_epochs=getattr(cfg.task, "epoch", None), diff --git a/yolo/model/__init__.py b/yolo/model/__init__.py new file mode 100644 index 000000000..e69de29bb From 636f153dbb19400b1559e000e3d4379bab1c8f1d Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 16:55:42 -0500 Subject: [PATCH 13/24] feat: allow user to specify accelerator --- yolo/config/config.py | 1 + yolo/config/general.yaml | 1 + yolo/lazy.py | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/yolo/config/config.py b/yolo/config/config.py index 9dd85b07a..a0a3f2512 100644 --- a/yolo/config/config.py +++ b/yolo/config/config.py @@ -157,6 +157,7 @@ class Config: use_tensorboard: bool weight: Optional[str] + accelerator: str @dataclass diff --git a/yolo/config/general.yaml b/yolo/config/general.yaml index c3380a799..0a4af6ca0 100644 --- a/yolo/config/general.yaml +++ b/yolo/config/general.yaml @@ -11,3 +11,4 @@ use_wandb: True use_tensorboard: False weight: True # Path to weight or True for auto, False for no pretrained weight +accelerator: 'auto' diff --git a/yolo/lazy.py b/yolo/lazy.py index a44268c9a..21df73339 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -17,7 +17,7 @@ def main(cfg: DictConfig): from lightning import Trainer from yolo.tools.solver import InferenceModel, TrainModel, ValidateModel trainer = Trainer( - accelerator="auto", + accelerator=cfg.accelerator, max_epochs=getattr(cfg.task, "epoch", None), precision="16-mixed", callbacks=callbacks, From 39ce158d7fbc8e6b9ebf7bbc26a1c94fdd391f1e Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 16:59:39 -0500 Subject: [PATCH 14/24] feat: allow user to simplify output with environ --- yolo/utils/logger.py | 11 +++++++++-- yolo/utils/logging_utils.py | 9 ++++++--- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/yolo/utils/logger.py b/yolo/utils/logger.py index 286027205..4880cf07c 100644 --- a/yolo/utils/logger.py +++ b/yolo/utils/logger.py @@ -3,9 +3,16 @@ from lightning.pytorch.utilities.rank_zero import rank_zero_only from rich.console import Console from rich.logging import RichHandler +import os logger = logging.getLogger("yolo") logger.setLevel(logging.DEBUG) logger.propagate = False -if rank_zero_only.rank == 0 and not logger.hasHandlers(): - logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True, markup=True)) + +# allow the user to get a simpler output +# TODO: needs to be better integrated +DISABLE_RICH_HANDLER = bool(os.environ.get('DISABLE_RICH_HANDLER', '')) + +if not DISABLE_RICH_HANDLER: + if rank_zero_only.rank == 0 and not logger.hasHandlers(): + logger.addHandler(RichHandler(console=Console(), show_level=True, show_path=True, show_time=True, markup=True)) diff --git a/yolo/utils/logging_utils.py b/yolo/utils/logging_utils.py index 42e0ac4ff..bdc872c3c 100644 --- a/yolo/utils/logging_utils.py +++ b/yolo/utils/logging_utils.py @@ -264,9 +264,12 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru logger.setLevel(logging.ERROR) return progress, loggers, save_path - progress.append(YOLORichProgressBar()) - progress.append(YOLORichModelSummary()) - progress.append(ImageLogger()) + from yolo.utils.logger import DISABLE_RICH_HANDLER + if not DISABLE_RICH_HANDLER: + progress.append(YOLORichProgressBar()) + progress.append(YOLORichModelSummary()) + progress.append(ImageLogger()) + if cfg.use_tensorboard: loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path)) if cfg.use_wandb: From 2e15f0a21817e2c13b551ef14a2dd5790e387c5e Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 17:14:47 -0500 Subject: [PATCH 15/24] refactor: disable validation sanity check for faster training response time --- yolo/lazy.py | 1 + yolo/utils/model_utils.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/yolo/lazy.py b/yolo/lazy.py index 21df73339..c0b0d6a1b 100644 --- a/yolo/lazy.py +++ b/yolo/lazy.py @@ -28,6 +28,7 @@ def main(cfg: DictConfig): # deterministic=True, enable_progress_bar=not getattr(cfg, "quite", False), default_root_dir=save_path, + num_sanity_val_steps=0, ) if cfg.task.task == "train": diff --git a/yolo/utils/model_utils.py b/yolo/utils/model_utils.py index 9d6c0ce59..bb78096eb 100644 --- a/yolo/utils/model_utils.py +++ b/yolo/utils/model_utils.py @@ -57,6 +57,14 @@ def on_validation_start(self, trainer: "Trainer", pl_module: "LightningModule"): self.ema_state_dict = deepcopy(pl_module.model.state_dict()) pl_module.ema.load_state_dict(self.ema_state_dict) + @no_grad() + def on_train_batch_start(self, trainer: "Trainer", pl_module: "LightningModule", *args, **kwargs) -> None: + if self.ema_state_dict is None: + # If validation sanity checks are disabled, then we need to + # initialize the ema state before training starts. + self.ema_state_dict = deepcopy(pl_module.model.state_dict()) + pl_module.ema.load_state_dict(self.ema_state_dict) + @no_grad() def on_train_batch_end(self, trainer: "Trainer", pl_module: "LightningModule", *args, **kwargs) -> None: self.step += 1 From c87f8c8ff7de2862f7ab82265a02846ad284619c Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 17:50:38 -0500 Subject: [PATCH 16/24] fix: ensure categories are remapped with kwcoco --- yolo/tools/data_loader.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index e638e8790..0850899e6 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -36,7 +36,9 @@ def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, phase: str transforms = [getattr(data_augmentation, aug)(prob) for aug, prob in augment_cfg.items()] self.transform = AugmentationComposer(transforms, self.image_size, self.base_size) self.transform.get_more_data = self.get_more_data - self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name)) + + data = self.load_data(Path(dataset_cfg.path), phase_name) + self.img_paths, self.bboxes, self.ratios = tensorlize(data) def load_data(self, dataset_path: Path, phase_name: str): """ @@ -100,6 +102,9 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa import kwcoco coco_dset = kwcoco.CocoDataset(labels_path) + from yolo.tools.data_conversion import discretize_categories + id_to_idx = discretize_categories(coco_dset.dataset.get("categories", [])) if "categories" in coco_dset.dataset else None + total_images = coco_dset.n_images if sort_image: @@ -121,7 +126,20 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa width, height = 0, 1 annotations = coco_img.annots().objs + + # Handle filtering as done in + # :func:`dataset_utils.organize_annotations_by_image` + modified_annotations = [] + for anno in annotations: + if id_to_idx: + anno["category_id"] = id_to_idx[anno["category_id"]] + if anno["iscrowd"]: + continue + modified_annotations.append(anno) + annotations = modified_annotations + if ALLOW_EMPTY_IMAGES or len(annotations): + image_seg_annotations = scale_segmentation(annotations, image_info) labels = self.load_valid_labels(None, image_seg_annotations) From 0f2d723476244a986237809328c2c2e524cb30cb Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 17:54:14 -0500 Subject: [PATCH 17/24] doc: add todo about data / classes --- yolo/tools/solver.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/yolo/tools/solver.py b/yolo/tools/solver.py index f631ab3fa..8590d3aeb 100644 --- a/yolo/tools/solver.py +++ b/yolo/tools/solver.py @@ -69,6 +69,10 @@ class TrainModel(ValidateModel): def __init__(self, cfg: Config): super().__init__(cfg) self.cfg = cfg + + # TODO: if we defer creating the model until the dataset is loaded, we + # can introspect the number of categories and other things to make user + # configuration have less interdependencies and thus be more robust. self.train_loader = create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task) def setup(self, stage): From 65df5afb196aeb4a73c9f4916af93588c11eafa2 Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 18:17:14 -0500 Subject: [PATCH 18/24] fix: dont assume iscrowd exists --- yolo/tools/data_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index 0850899e6..d7dfc7386 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -133,7 +133,7 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa for anno in annotations: if id_to_idx: anno["category_id"] = id_to_idx[anno["category_id"]] - if anno["iscrowd"]: + if anno.get("iscrowd", False): # TODO: make configurable continue modified_annotations.append(anno) annotations = modified_annotations From b48f2740003dbf34dd4501d5dddf6c4f041ce836 Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 18:17:51 -0500 Subject: [PATCH 19/24] fix: valid points check was incorrect --- yolo/tools/data_loader.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index d7dfc7386..dde3030dd 100644 --- a/yolo/tools/data_loader.py +++ b/yolo/tools/data_loader.py @@ -205,8 +205,18 @@ def load_valid_labels(self, label_path: str, seg_data_one_img: list) -> Union[Te bboxes = [] for seg_data in seg_data_one_img: cls = seg_data[0] - points = np.array(seg_data[1:]).reshape(-1, 2) - valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2) + # This seems like an incorrect check. Putting my fix inside an if + # in case I don't understand why it is this way. + FIX_INCORRECT_CHECK = 1 + if FIX_INCORRECT_CHECK: + points = np.array(seg_data[1:]).reshape(-1, 2) + # This probably should just be a clamp / clip operation + # but I'm keeping it similar to the original + flags = (points >= 0).all(axis=1) & (points <= 1).all(axis=1) + valid_points = points[flags] + else: + points = np.array(seg_data[1:]).reshape(-1, 2) + valid_points = points[(points >= 0) & (points <= 1)].reshape(-1, 2) if valid_points.size > 1: bbox = torch.tensor([cls, *valid_points.min(axis=0), *valid_points.max(axis=0)]) bboxes.append(bbox) From 6e868a599574ef67a2621f21329da481c9fd95a0 Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 18:19:00 -0500 Subject: [PATCH 20/24] feat: use kwimage to handle more polygon reprs --- yolo/utils/dataset_utils.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/yolo/utils/dataset_utils.py b/yolo/utils/dataset_utils.py index 98699e617..b794b8c55 100644 --- a/yolo/utils/dataset_utils.py +++ b/yolo/utils/dataset_utils.py @@ -105,12 +105,22 @@ def scale_segmentation( if annotations is None: return None + try: + import kwimage + except ImportError: + kwimage = None + seg_array_with_cat = [] h, w = image_dimensions["height"], image_dimensions["width"] for anno in annotations: category_id = anno["category_id"] if "segmentation" in anno: - seg_list = [item for sublist in anno["segmentation"] for item in sublist] + if kwimage is None: + # original fallback code + seg_list = [item for sublist in anno["segmentation"] for item in sublist] + else: + # Convert to original coco representation + seg_list = kwimage.MultiPolygon.coerce(anno["segmentation"]).to_coco('orig') elif "bbox" in anno: x, y, width, height = anno["bbox"] seg_list = [x, y, x + width, y, x + width, y + height, x, y + height] From 0c92ec15a4438d141a06322aa86322bf1dfe07db Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 18:21:35 -0500 Subject: [PATCH 21/24] add: kwcoco training tutorial --- train_kwcoco_demo.sh | 102 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 train_kwcoco_demo.sh diff --git a/train_kwcoco_demo.sh b/train_kwcoco_demo.sh new file mode 100644 index 000000000..53e785bb1 --- /dev/null +++ b/train_kwcoco_demo.sh @@ -0,0 +1,102 @@ +#!/bin/bash +__doc__=" +YOLO Training Tutorial with KWCOCO DemoData +=========================================== + +This demonstrates an end-to-end YOLO pipeline on toydata generated with kwcoco. +" + +# Define where we will store results +BUNDLE_DPATH=$HOME/demo-yolo-kwcoco-train +mkdir -p "$BUNDLE_DPATH" + +echo " +Generate Toy Data +----------------- + +Now that we know where the data and our intermediate files will go, lets +generate the data we will use to train and evaluate with. + +The kwcoco package comes with a commandline utility called 'kwcoco toydata' to +accomplish this. +" + +# Define the names of the kwcoco files to generate +TRAIN_FPATH=$BUNDLE_DPATH/vidshapes_rgb_train/data.kwcoco.json +VALI_FPATH=$BUNDLE_DPATH/vidshapes_rgb_vali/data.kwcoco.json +TEST_FPATH=$BUNDLE_DPATH/vidshapes_rgb_test/data.kwcoco.json + +# Generate toy datasets using the "kwcoco toydata" tool +kwcoco toydata vidshapes2-frames10 --dst "$TRAIN_FPATH" +kwcoco toydata vidshapes4-frames10 --dst "$VALI_FPATH" +kwcoco toydata vidshapes2-frames6 --dst "$TEST_FPATH" + +# Ensure legacy COCO structure for now +kwcoco conform "$TRAIN_FPATH" --inplace --legacy=True +kwcoco conform "$VALI_FPATH" --inplace --legacy=True +kwcoco conform "$TEST_FPATH" --inplace --legacy=True + + +echo " +Create the YOLO Configuration +----------------------------- + +Constructing the YOLO configuration is not entirely kwcoco aware +so we need to set +" +# In the current version we need to write configs to the repo itself. +# Its a bit gross, but this should be somewhat robust. +# Find where the yolo repo is installed (we need to be careful that this is the +# our fork of the WongKinYiu variant +REPO_DPATH=$(python -c "import yolo, pathlib; print(pathlib.Path(yolo.__file__).parent.parent)") +MODULE_DPATH=$(python -c "import yolo, pathlib; print(pathlib.Path(yolo.__file__).parent)") +CONFIG_DPATH=$(python -c "import yolo.config, pathlib; print(pathlib.Path(yolo.config.__file__).parent / 'dataset')") +echo "REPO_DPATH = $REPO_DPATH" +echo "MODULE_DPATH = $MODULE_DPATH" +echo "CONFIG_DPATH = $CONFIG_DPATH" + +DATASET_CONFIG_FPATH=$CONFIG_DPATH/kwcoco-demo.yaml + +# Hack to construct the class part of the YAML +CLASS_YAML=$(python -c "if 1: + import kwcoco + train_fpath = kwcoco.CocoDataset('$TRAIN_FPATH') + categories = train_fpath.categories().objs + # It would be nice to have better class introspection, but in the meantime + # do the same sorting as yolo.tools.data_conversion.discretize_categories + categories = sorted(categories, key=lambda cat: cat['id']) + class_num = len(categories) + class_list = [c['name'] for c in categories] + print(f'class_num: {class_num}') + print(f'class_list: {class_list}') +") + + +CONFIG_YAML=" +path: $BUNDLE_DPATH +train: $TRAIN_FPATH +validation: $VALI_FPATH + +$CLASS_YAML +" + +echo "$CONFIG_YAML" > "$DATASET_CONFIG_FPATH" + + +# This might only work in development mode, otherwise we will get site packages +# That still might be fine, but we do want to fix this to run anywhere. +cd "$REPO_DPATH" +python -m yolo.lazy \ + task=train \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + accelerator=auto \ + task.data.batch_size=2 \ + "image_size=[224,224]" \ + task.optimizer.args.lr=0.003 + +#--help From d7fa4478339efe634c3b99003f9b428b02eaedf7 Mon Sep 17 00:00:00 2001 From: joncrall Date: Thu, 23 Jan 2025 19:34:14 -0500 Subject: [PATCH 22/24] refactor: improve on-disk batch viz --- train_kwcoco_demo.sh | 49 ++++++++++++++++++++-- yolo/utils/logging_utils.py | 83 ++++++++++++++++++++++++++++++++++++- 2 files changed, 127 insertions(+), 5 deletions(-) diff --git a/train_kwcoco_demo.sh b/train_kwcoco_demo.sh index 53e785bb1..e631203a9 100644 --- a/train_kwcoco_demo.sh +++ b/train_kwcoco_demo.sh @@ -27,7 +27,7 @@ VALI_FPATH=$BUNDLE_DPATH/vidshapes_rgb_vali/data.kwcoco.json TEST_FPATH=$BUNDLE_DPATH/vidshapes_rgb_test/data.kwcoco.json # Generate toy datasets using the "kwcoco toydata" tool -kwcoco toydata vidshapes2-frames10 --dst "$TRAIN_FPATH" +kwcoco toydata vidshapes32-frames10 --dst "$TRAIN_FPATH" kwcoco toydata vidshapes4-frames10 --dst "$VALI_FPATH" kwcoco toydata vidshapes2-frames6 --dst "$TEST_FPATH" @@ -86,7 +86,7 @@ echo "$CONFIG_YAML" > "$DATASET_CONFIG_FPATH" # This might only work in development mode, otherwise we will get site packages # That still might be fine, but we do want to fix this to run anywhere. cd "$REPO_DPATH" -python -m yolo.lazy \ +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ task=train \ dataset=kwcoco-demo \ use_wandb=False \ @@ -97,6 +97,47 @@ python -m yolo.lazy \ accelerator=auto \ task.data.batch_size=2 \ "image_size=[224,224]" \ - task.optimizer.args.lr=0.003 + task.optimizer.args.lr=0.0003 -#--help +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ + task=train \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + accelerator=auto \ + task.data.batch_size=2 \ + "image_size=[224,224]" \ + task.optimizer.args.lr=0.0003 + + +### TODO: show how to validate + +# Grab a checkpoint +CKPT_FPATH=$(python -c "if 1: + import pathlib + ckpt_dpath = pathlib.Path('$BUNDLE_DPATH') / 'training/train/kwcoco-demo/checkpoints' + checkpoints = sorted(ckpt_dpath.glob('*')) + print(checkpoints[-1]) +") +echo "CKPT_FPATH = $CKPT_FPATH" + + +#DISABLE_RICH_HANDLER=1 +LOG_BATCH_VIZ_TO_DISK=1 python -m yolo.lazy \ + task=validation \ + dataset=kwcoco-demo \ + use_wandb=False \ + out_path="$BUNDLE_DPATH"/training \ + name=kwcoco-demo \ + cpu_num=0 \ + device=0 \ + weight="'$CKPT_FPATH'" \ + accelerator=auto \ + "task.data.batch_size=2" \ + "image_size=[224,224]" + + +### TODO: show how to run inference diff --git a/yolo/utils/logging_utils.py b/yolo/utils/logging_utils.py index bdc872c3c..9ca14d25e 100644 --- a/yolo/utils/logging_utils.py +++ b/yolo/utils/logging_utils.py @@ -216,12 +216,92 @@ def on_validation_batch_end(self, trainer: Trainer, pl_module, outputs, batch, b pred_boxes = outputs[0] if isinstance(outputs, list) else outputs images = [images[0]] step = trainer.current_epoch + for logger in trainer.loggers: if isinstance(logger, WandbLogger): + # FIXME: not robust to configured image sizes, need to know + # that info. logger.log_image("Input Image", images, step=step) logger.log_image("Ground Truth", images, step=step, boxes=[log_bbox(gt_boxes)]) logger.log_image("Prediction", images, step=step, boxes=[log_bbox(pred_boxes)]) + # TODO: better config + import os + LOG_BATCH_VIZ_TO_DISK = bool(os.environ.get('LOG_BATCH_VIZ_TO_DISK', '')) + if LOG_BATCH_VIZ_TO_DISK: + import einops + import kwimage + + # TODO: + # get a batter output path + import pathlib + root_dpath = pathlib.Path(trainer.default_root_dir) + out_dpath = root_dpath / 'debug_images' / trainer.state.stage.name + out_dpath.mkdir(exist_ok=True, parents=True) + epoch = trainer.current_epoch + + for bx in range(len(images)): + image_chw = images[bx].data.cpu().numpy() + image_hwc = einops.rearrange(image_chw, 'c h w -> h w c') + image_hwc = kwimage.ensure_uint255(image_hwc) + + assert bx == 0, 'not handling multiple per batch' + true_dets = tensor_to_kwimage(gt_boxes).numpy() + pred_dets = tensor_to_kwimage(pred_boxes).numpy() + pred_dets = pred_dets.non_max_supress(thresh=0.3) + # pred_dets = pred_dets.compress(pred_dets.scores > 0.1) + + raw_canvas = image_hwc.copy() + true_canvas = true_dets.draw_on(raw_canvas.copy(), color='green') + pred_canvas = pred_dets.draw_on(raw_canvas.copy(), color='blue') + + raw_canvas = kwimage.draw_header_text(raw_canvas, 'raw') + true_canvas = kwimage.draw_header_text(true_canvas, f'true, n={len(true_dets)}') + pred_canvas = kwimage.draw_header_text(pred_canvas, f'pred, n={len(pred_dets)}') + canvas = kwimage.stack_images([ + raw_canvas, true_canvas, pred_canvas + ], axis=1, pad=3) + + fname = f'img_{epoch:04d}_{batch_idx:04d}.jpg' + fpath = out_dpath / fname + kwimage.imwrite(fpath, canvas) + + +def tensor_to_kwimage(yolo_annot_tensor): + import kwimage + class_idxs = yolo_annot_tensor[:, 0].int() + boxes = kwimage.Boxes(yolo_annot_tensor[:, 1:5], format='xyxy') + dets = kwimage.Detections( + boxes=boxes, + class_idxs=class_idxs + ) + + if yolo_annot_tensor.shape[1] > 5: + scores = yolo_annot_tensor[:, 5] + dets.data['scores'] = scores + return dets + + +def wandb_to_kwimage(wand_annots): + import numpy as np + import kwimage + box_list = [] + class_idxs = [] + for row in wand_annots['predictions']['box_data']: + pos = row['position'] + class_idx = row['class_id'] + xyxy = [pos['minX'], pos['minY'], pos['maxX'], pos['maxY']] + box_list.append(xyxy) + class_idxs.append(class_idx) + + boxes = kwimage.Boxes(np.array(box_list), format='xyxy') + dets = kwimage.Detections( + boxes=boxes, + class_idxs=np.array(class_idxs) + ) + dets = dets.compress(dets.class_idxs > -1) + return dets + def setup_logger(logger_name, quite=False): class EmojiFormatter(logging.Formatter): @@ -268,7 +348,8 @@ def custom_wandb_log(string="", level=int, newline=True, repeat=True, prefix=Tru if not DISABLE_RICH_HANDLER: progress.append(YOLORichProgressBar()) progress.append(YOLORichModelSummary()) - progress.append(ImageLogger()) + + progress.append(ImageLogger()) if cfg.use_tensorboard: loggers.append(TensorBoardLogger(log_graph="all", save_dir=save_path)) From 3eb8a68086c1f87584fb910ea685fe47764cae7f Mon Sep 17 00:00:00 2001 From: joncrall Date: Mon, 27 Jan 2025 19:50:08 -0500 Subject: [PATCH 23/24] feat: add weights loading log statement --- yolo/model/yolo.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/yolo/model/yolo.py b/yolo/model/yolo.py index ccf456748..e38db1629 100644 --- a/yolo/model/yolo.py +++ b/yolo/model/yolo.py @@ -164,6 +164,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, OmegaConf.set_struct(model_cfg, False) model = YOLO(model_cfg, class_num) if weight_path: + logger.info('🏋 Initializing weights') if weight_path == True: weight_path = Path("weights") / f"{model_cfg.name}.pt" elif isinstance(weight_path, str): @@ -173,8 +174,9 @@ def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, logger.info(f"🌐 Weight {weight_path} not found, try downloading") prepare_weight(weight_path=weight_path) if weight_path.exists(): + logger.info(f'🏋 Loading weights from {weight_path}') model.save_load_weights(weight_path) logger.info(":white_check_mark: Success load model & weight") else: - logger.info(":white_check_mark: Success load model") + logger.info(":white_check_mark: Success load model without weights") return model From b5f729ef34b8ba6367ab5ead254b75b9f8663c6f Mon Sep 17 00:00:00 2001 From: joncrall Date: Mon, 27 Jan 2025 20:33:51 -0500 Subject: [PATCH 24/24] fix: workaround weight loading issue at inference time --- yolo/model/yolo.py | 112 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 92 insertions(+), 20 deletions(-) diff --git a/yolo/model/yolo.py b/yolo/model/yolo.py index e38db1629..bb43016f7 100644 --- a/yolo/model/yolo.py +++ b/yolo/model/yolo.py @@ -129,27 +129,99 @@ def save_load_weights(self, weights: Union[Path, OrderedDict]): weights = torch.load(weights, map_location=torch.device("cpu"), weights_only=False) if "model_state_dict" in weights: weights = weights["model_state_dict"] + if "state_dict" in weights: + weights = weights["state_dict"] + + if 0: + # Debug the state of the model and the loaded weights + import networkx as nx + graph_src = nx.DiGraph() + for key in list(weights.keys()): + graph_src.add_node(key) + graph_src.add_node('__root__') + for key in list(weights.keys()): + parts = key.split('.') + graph_src.add_edge('__root__', parts[0]) + for i in range(1, len(parts)): + parent = '.'.join(parts[:i - 1]) + child = '.'.join(parts[:i]) + graph_src.add_edge(parent, child) + nx.write_network_text(graph_src, max_depth=4, sources=['__root__']) + + graph_dst = nx.DiGraph() + dst_weights = self.state_dict() + for key in list(dst_weights.keys()): + graph_dst.add_node(key) + graph_dst.add_node('__root__') + for key in list(dst_weights.keys()): + parts = key.split('.') + graph_dst.add_edge('__root__', parts[0]) + for i in range(1, len(parts)): + parent = '.'.join(parts[:i - 1]) + child = '.'.join(parts[:i]) + graph_dst.add_edge(parent, child) + nx.write_network_text(graph_dst, max_depth=3, sources=['__root__']) + + USE_TORCH_LIBERATOR = False + if USE_TORCH_LIBERATOR: + + # Torch liberator will figure out the mapping in most cases but it + # is slow. + HACK_DONT_LOAD_EMA_WEIGHTS = True + if HACK_DONT_LOAD_EMA_WEIGHTS: + for key in list(weights.keys()): + if key.startswith('ema.model'): + weights.pop(key) + from torch_liberator.initializer import load_partial_state + load_partial_state(self, weights, verbose=3) - model_state_dict = self.model.state_dict() - - # TODO1: autoload old version weight - # TODO2: weight transform if num_class difference - - error_dict = {"Mismatch": set(), "Not Found": set()} - for model_key, model_weight in model_state_dict.items(): - if model_key not in weights: - error_dict["Not Found"].add(tuple(model_key.split(".")[:-2])) - continue - if model_weight.shape != weights[model_key].shape: - error_dict["Mismatch"].add(tuple(model_key.split(".")[:-2])) - continue - model_state_dict[model_key] = weights[model_key] - - for error_name, error_set in error_dict.items(): - for weight_name in error_set: - logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}") - - self.model.load_state_dict(model_state_dict) + else: + # TODO1: autoload old version weight + # TODO2: weight transform if num_class difference + + model_state_dict = self.model.state_dict() + + CHECK_FOR_WEIGHT_MUNGING = True + if CHECK_FOR_WEIGHT_MUNGING: + # Handle the simple case of weight munging ourselves + src_keys = list(weights.keys()) + dst_keys = list(model_state_dict.keys()) + src_roots = {p.split('.')[0] for p in src_keys} + dst_roots = {p.split('.')[0] for p in dst_keys} + if len(src_roots & dst_roots) == 0: + src_prefixes = {tuple(p.split('.')[0:2]) for p in src_keys} + if src_prefixes == {('ema', 'model'), ('model', 'model')}: + logger.warning(":warning: Munging weights") + munged_weights = {} + for key in list(weights.keys()): + prefix = 'model.model.' + if key.startswith(prefix): + new_key = key[len(prefix):] + munged_weights[new_key] = weights[key] + logger.warning(f":warning: Munged {len(munged_weights)} / {len(weights)} tensors") + weights = munged_weights + + error_dict = {"Mismatch": set(), "Not Found": set()} + for model_key, model_weight in model_state_dict.items(): + if model_key not in weights: + error_dict["Not Found"].add(tuple(model_key.split(".")[:-2])) + continue + if model_weight.shape != weights[model_key].shape: + error_dict["Mismatch"].add(tuple(model_key.split(".")[:-2])) + continue + model_state_dict[model_key] = weights[model_key] + + for error_name, error_set in error_dict.items(): + for weight_name in error_set: + logger.warning(f":warning: Weight {error_name} for key: {'.'.join(weight_name)}") + + for error_name, error_set in error_dict.items(): + if len(error_set) == 0: + logger.info(f":white_check_mark: Num: weight {error_name}: {len(error_set)}") + else: + logger.warning(f":warning: Num: weight {error_name}: {len(error_set)}") + + self.model.load_state_dict(model_state_dict) def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, class_num: int = 80) -> YOLO: