Upgrades

requirements.txt

@@ -2,6 +2,7 @@ einops
 faster-coco-eval
 graphviz
 hydra-core
+lazy-loader
 lightning
 loguru
 numpy

train_kwcoco_demo.sh

@@ -0,0 +1,143 @@
+#!/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 vidshapes32-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"
+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
+
+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

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']

yolo/config/config.py

@@ -157,6 +157,7 @@ class Config:
     use_tensorboard: bool
 
     weight: Optional[str]
+    accelerator: str
 
 
 @dataclass

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'

yolo/lazy.py

@@ -2,32 +2,33 @@
 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: Config):
+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",
+        accelerator=cfg.accelerator,
         max_epochs=getattr(cfg.task, "epoch", None),
         precision="16-mixed",
         callbacks=callbacks,
         logger=loggers,
         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,
+        num_sanity_val_steps=0,
     )
 
     if cfg.task.task == "train":