diff --git a/configs/README.md b/configs/README.md index 96444f66..ca95b08c 100644 --- a/configs/README.md +++ b/configs/README.md @@ -148,7 +148,6 @@ Here you can change everything related to actual training of the model. | use_weighted_sampler | bool | False | bool if use WeightedRandomSampler for training, only works with classification tasks | | epochs | int | 100 | number of training epochs | | n_workers | int | 2 | number of workers for data loading | -| train_metrics_interval | int | -1 | frequency of computing metrics on train data, -1 if don't perform | | validation_interval | int | 1 | frequency of computing metrics on validation data | | n_log_images | int | 4 | maximum number of images to visualize and log | | skip_last_batch | bool | True | whether to skip last batch while training | diff --git a/configs/detection_model.yaml b/configs/classification_heavy_model.yaml similarity index 50% rename from configs/detection_model.yaml rename to configs/classification_heavy_model.yaml index 7bc87eef..9feda041 100644 --- a/configs/detection_model.yaml +++ b/configs/classification_heavy_model.yaml @@ -1,25 +1,25 @@ -# Example configuration for training a predefined detection model +# Example configuration for training a predefined heavy classification model model: - name: coco_detection + name: classification_light predefined_model: - name: DetectionModel + name: ClassificationModel params: - use_neck: True + variant: heavy loader: params: - dataset_name: coco_test + dataset_name: cifar10_test trainer: preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False + train_image_size: [384, 512] + keep_aspect_ratio: True normalize: active: True - batch_size: 4 - epochs: &epochs 200 + batch_size: 8 + epochs: 200 n_workers: 4 validation_interval: 10 n_log_images: 8 @@ -29,9 +29,9 @@ trainer: - name: TestOnTrainEnd optimizer: - name: SGD + name: Adam params: - lr: 0.02 + lr: 0.001 scheduler: name: ConstantLR diff --git a/configs/classification_light_model.yaml b/configs/classification_light_model.yaml new file mode 100644 index 00000000..32f7d96b --- /dev/null +++ b/configs/classification_light_model.yaml @@ -0,0 +1,37 @@ +# Example configuration for training a predefined light classification model + +model: + name: classification_light + predefined_model: + name: ClassificationModel + params: + variant: light + +loader: + params: + dataset_name: cifar10_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: 200 + n_workers: 4 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd + - name: TestOnTrainEnd + + optimizer: + name: Adam + params: + lr: 0.003 + + scheduler: + name: ConstantLR diff --git a/configs/classification_model.yaml b/configs/classification_model.yaml deleted file mode 100644 index 4db7a9b1..00000000 --- a/configs/classification_model.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Example configuration for training a predefined segmentation model - -model: - name: cifar10_classification - predefined_model: - name: ClassificationModel - params: - backbone: MicroNet - visualizer_params: - font_scale: 0.5 - color: [255, 0, 0] - thickness: 2 - include_plot: True - -loader: - params: - dataset_name: cifar10_test - -trainer: - preprocessing: - train_image_size: [&height 128, &width 128] - keep_aspect_ratio: False - normalize: - active: True - - batch_size: 4 - epochs: &epochs 200 - n_workers: 4 - validation_interval: 10 - n_log_images: 8 - - callbacks: - - name: ExportOnTrainEnd - - name: TestOnTrainEnd - - optimizer: - name: SGD - params: - lr: 0.02 - - scheduler: - name: ConstantLR diff --git a/configs/coco_model.yaml b/configs/complex_model.yaml similarity index 78% rename from configs/coco_model.yaml rename to configs/complex_model.yaml index dd0256b8..1ba11dd1 100644 --- a/configs/coco_model.yaml +++ b/configs/complex_model.yaml @@ -2,37 +2,21 @@ model: - name: coco_test + name: complex_model nodes: - name: EfficientRep - params: - channels_list: [64, 128, 256, 512, 1024] - n_repeats: [1, 6, 12, 18, 6] - depth_mul: 0.33 - width_mul: 0.33 - name: RepPANNeck inputs: - EfficientRep - params: - channels_list: [256, 128, 128, 256, 256, 512] - n_repeats: [12, 12, 12, 12] - depth_mul: 0.33 - width_mul: 0.33 - - name: ImplicitKeypointBBoxHead + - name: EfficientKeypointBBoxHead inputs: - RepPANNeck - params: - conf_thres: 0.25 - iou_thres: 0.45 + losses: - name: ImplicitKeypointBBoxLoss - params: - keypoint_regression_loss_weight: 0.5 - keypoint_visibility_loss_weight: 0.7 - bbox_loss_weight: 0.05 - objectness_loss_weight: 0.2 + name: EfficientKeypointBboxLoss + metrics: - name: ObjectKeypointSimilarity is_main_metric: true @@ -40,7 +24,6 @@ model: visualizers: name: MultiVisualizer - attached_to: ImplicitKeypointBBoxHead params: visualizers: - name: KeypointVisualizer @@ -86,7 +69,6 @@ tracker: save_directory: output is_tensorboard: True is_wandb: False - wandb_entity: luxonis is_mlflow: False loader: @@ -105,11 +87,10 @@ trainer: n_sanity_val_steps: 1 profiler: null verbose: True - batch_size: 4 + batch_size: 8 accumulate_grad_batches: 1 epochs: &epochs 200 n_workers: 8 - train_metrics_interval: -1 validation_interval: 10 n_log_images: 8 skip_last_batch: True @@ -117,8 +98,8 @@ trainer: save_top_k: 3 preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False + train_image_size: [&height 384, &width 384] + keep_aspect_ratio: True train_rgb: True normalize: active: True diff --git a/configs/ddrnet_segmentation_model.yaml b/configs/ddrnet_segmentation_model.yaml deleted file mode 100644 index 2bd3b7e8..00000000 --- a/configs/ddrnet_segmentation_model.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# DDRNet-23-slim model for segmentation -# Refer to here for optimal hyperparameters for this model: https://github.com/Deci-AI/super-gradients/blob/4797c974c7c445d12e2575c468848d9c3e04becd/src/super_gradients/recipes/cityscapes_ddrnet.yaml#L4 - -model: - name: ddrnet_segmentation - predefined_model: - name: DDRNetSegmentationModel - params: - task: binary - backbone_params: - use_aux_heads: True # set to False to disable auxiliary heads (for export) - variant: '23-slim' - -loader: - params: - dataset_name: coco_test - -trainer: - preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False - normalize: - active: True - - batch_size: 4 - epochs: &epochs 500 - num_workers: 4 - validation_interval: 10 - num_log_images: 8 - - callbacks: - - name: TestOnTrainEnd - - name: ExportOnTrainEnd - - optimizer: - name: SGD - params: - lr: 0.01 - momentum: 0.9 - weight_decay: 0.0005 - - scheduler: - name: CosineAnnealingLR - params: - T_max: *epochs diff --git a/configs/detection_heavy_model.yaml b/configs/detection_heavy_model.yaml new file mode 100644 index 00000000..f35c1ed3 --- /dev/null +++ b/configs/detection_heavy_model.yaml @@ -0,0 +1,45 @@ +# Example configuration for training a predefined heavy detection model + +model: + name: detection_heavy + predefined_model: + name: DetectionModel + params: + variant: heavy + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + n_workers: 4 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd + - name: TestOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.01 + momentum: 0.937 + weight_decay: 0.0005 + dampening: 0.0 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.0001 + last_epoch: -1 diff --git a/configs/detection_light_model.yaml b/configs/detection_light_model.yaml new file mode 100644 index 00000000..1f982d92 --- /dev/null +++ b/configs/detection_light_model.yaml @@ -0,0 +1,45 @@ +# Example configuration for training a predefined light detection model + +model: + name: detection_light + predefined_model: + name: DetectionModel + params: + variant: light + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + n_workers: 8 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd + - name: TestOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.02 + momentum: 0.937 + weight_decay: 0.0005 + dampening: 0.0 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.0002 + last_epoch: -1 diff --git a/configs/efficient_coco_model.yaml b/configs/efficient_coco_model.yaml deleted file mode 100644 index f2c9db5d..00000000 --- a/configs/efficient_coco_model.yaml +++ /dev/null @@ -1,114 +0,0 @@ - -model: - name: coco_test - nodes: - - name: EfficientRep - params: - channels_list: [64, 128, 256, 512, 1024] - n_repeats: [1, 6, 12, 18, 6] - depth_mul: 0.33 - width_mul: 0.33 - - - name: RepPANNeck - inputs: - - EfficientRep - params: - channels_list: [256, 128, 128, 256, 256, 512] - n_repeats: [12, 12, 12, 12] - depth_mul: 0.33 - width_mul: 0.33 - - - name: EfficientKeypointBBoxHead - inputs: - - RepPANNeck - params: - conf_thres: 0.25 - iou_thres: 0.45 - - - name: SegmentationHead - inputs: - - RepPANNeck - - - name: EfficientBBoxHead - inputs: - - RepPANNeck - params: - conf_thres: 0.75 - iou_thres: 0.45 - - losses: - - name: AdaptiveDetectionLoss - attached_to: EfficientBBoxHead - - name: BCEWithLogitsLoss - attached_to: SegmentationHead - - name: EfficientKeypointBBoxLoss - attached_to: EfficientKeypointBBoxHead - - metrics: - - name: ObjectKeypointSimilarity - is_main_metric: true - attached_to: EfficientKeypointBBoxHead - - name: MeanAveragePrecisionKeypoints - attached_to: EfficientKeypointBBoxHead - - name: MeanAveragePrecision - attached_to: EfficientBBoxHead - - name: F1Score - attached_to: SegmentationHead - params: - task: binary - - name: JaccardIndex - attached_to: SegmentationHead - params: - task: binary - - visualizers: - - name: MultiVisualizer - attached_to: EfficientKeypointBBoxHead - params: - visualizers: - - name: KeypointVisualizer - params: - nonvisible_color: blue - - name: BBoxVisualizer - params: - colors: - person: "#FF5055" - - name: SegmentationVisualizer - attached_to: SegmentationHead - params: - colors: "#FF5055" - - name: BBoxVisualizer - attached_to: EfficientBBoxHead - -tracker: - project_name: coco_test - save_directory: output - is_tensorboard: True - -loader: - params: - dataset_name: coco_test - -trainer: - - n_sanity_val_steps: 1 - batch_size: 4 - accumulate_grad_batches: 1 - epochs: 200 - n_workers: 4 - train_metrics_interval: -1 - validation_interval: 10 - n_log_images: 8 - save_top_k: 3 - - preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False - train_rgb: True - normalize: - active: True - - callbacks: - - name: ExportOnTrainEnd - - name: TestOnTrainEnd - diff --git a/configs/example_export.yaml b/configs/example_export.yaml index 51f768dc..78f1c650 100644 --- a/configs/example_export.yaml +++ b/configs/example_export.yaml @@ -1,13 +1,12 @@ -# Example configuration for exporting a predefined segmentation model +# Example configuration for exporting a predefined light detection model model: - name: coco_segmentation - weights: null # specify a path to the weights here + name: detection_light + weights: null # TODO: Specify a path to the weights here predefined_model: - name: SegmentationModel + name: DetectionModel params: - backbone: MicroNet - task: binary + variant: light loader: params: @@ -15,12 +14,12 @@ loader: trainer: preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False + train_image_size: [384, 512] + keep_aspect_ratio: True normalize: active: True - batch_size: 4 + batch_size: 8 epochs: &epochs 200 n_workers: 4 validation_interval: 10 @@ -28,11 +27,22 @@ trainer: optimizer: name: SGD + params: + lr: 0.02 + momentum: 0.937 + weight_decay: 0.0005 + dampening: 0.0 + nesterov: true scheduler: - name: ConstantLR + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.0002 + last_epoch: -1 exporter: + output_names: [output1_yolov6r2, output2_yolov6r2, output3_yolov6r2] onnx: opset_version: 11 blobconverter: diff --git a/configs/example_tuning.yaml b/configs/example_tuning.yaml old mode 100644 new mode 100755 index d8c9027d..9e63c877 --- a/configs/example_tuning.yaml +++ b/configs/example_tuning.yaml @@ -1,12 +1,12 @@ -# Example configuration for tuning a predefined segmentation model +# Example configuration for tuning a predefined light segmentation model + model: - name: coco_segmentation + name: segmentation_light predefined_model: name: SegmentationModel params: - backbone: MicroNet - task: binary + variant: light loader: params: @@ -14,8 +14,8 @@ loader: trainer: preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False + train_image_size: [384, 512] + keep_aspect_ratio: True normalize: active: True augmentations: @@ -27,8 +27,8 @@ trainer: p: 0.1 - name: Flip - batch_size: 4 - epochs: &epochs 100 + batch_size: 8 + epochs: &epochs 200 validation_interval: 10 n_log_images: 8 @@ -38,7 +38,12 @@ trainer: T_max: *epochs eta_min: 0 + tuner: + study_name: seg_study + n_trials: 10 + storage: + storage_type: local params: trainer.optimizer.name_categorical: ["Adam", "SGD"] trainer.optimizer.params.lr_float: [0.0001, 0.001] diff --git a/configs/keypoint_bbox_heavy_model.yaml b/configs/keypoint_bbox_heavy_model.yaml new file mode 100644 index 00000000..c6b22f35 --- /dev/null +++ b/configs/keypoint_bbox_heavy_model.yaml @@ -0,0 +1,45 @@ +# Example configuration for training a predefined heavy keypoint-detection model + +model: + name: keypoint_detection_heavy + predefined_model: + name: KeypointDetectionModel + params: + variant: heavy + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + n_workers: 4 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd + - name: TestOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.006 + momentum: 0.937 + weight_decay: 0.0005 + dampening: 0.0 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.00001 + last_epoch: -1 diff --git a/configs/keypoint_bbox_light_model.yaml b/configs/keypoint_bbox_light_model.yaml new file mode 100644 index 00000000..a095a551 --- /dev/null +++ b/configs/keypoint_bbox_light_model.yaml @@ -0,0 +1,45 @@ +# Example configuration for training a predefined light keypoint-detection model + +model: + name: keypoint_detection_light + predefined_model: + name: KeypointDetectionModel + params: + variant: light + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + n_workers: 4 + validation_interval: 10 + n_log_images: 8 + + callbacks: + - name: ExportOnTrainEnd + - name: TestOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.006 + momentum: 0.937 + weight_decay: 0.0005 + dampening: 0.0 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.00001 + last_epoch: -1 diff --git a/configs/keypoint_bbox_model.yaml b/configs/keypoint_bbox_model.yaml deleted file mode 100644 index 51554f73..00000000 --- a/configs/keypoint_bbox_model.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# Example configuration for training a predefined keypoint-detection model - -model: - name: coco_keypoints - predefined_model: - name: KeypointDetectionModel - -loader: - params: - dataset_name: coco_test - -trainer: - preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False - normalize: - active: True - - batch_size: 4 - epochs: &epochs 200 - n_workers: 4 - validation_interval: 10 - n_log_images: 8 - - callbacks: - - name: ExportOnTrainEnd - - name: TestOnTrainEnd - - optimizer: - name: SGD - params: - lr: 0.02 - - scheduler: - name: ConstantLR diff --git a/configs/resnet_model.yaml b/configs/resnet_model.yaml deleted file mode 100644 index fd0b66cd..00000000 --- a/configs/resnet_model.yaml +++ /dev/null @@ -1,54 +0,0 @@ - -model: - name: resnet50_classification - nodes: - - name: ResNet - params: - variant: "50" - download_weights: True - - - name: ClassificationHead - - losses: - name: CrossEntropyLoss - - metrics: - name: Accuracy - is_main_metric: true - - visualizers: - name: ClassificationVisualizer - params: - font_scale: 0.5 - color: [255, 0, 0] - thickness: 2 - include_plot: True - -loader: - params: - dataset_name: cifar10_test - -trainer: - batch_size: 4 - epochs: &epochs 200 - n_workers: 4 - validation_interval: 10 - n_log_images: 8 - - preprocessing: - train_image_size: [&height 224, &width 224] - keep_aspect_ratio: False - normalize: - active: True - - callbacks: - - name: ExportOnTrainEnd - - name: TestOnTrainEnd - - optimizer: - name: SGD - params: - lr: 0.02 - - scheduler: - name: ConstantLR diff --git a/configs/segmentation_heavy_model.yaml b/configs/segmentation_heavy_model.yaml new file mode 100644 index 00000000..e9bc16d6 --- /dev/null +++ b/configs/segmentation_heavy_model.yaml @@ -0,0 +1,42 @@ +# Example configuration for training a predefined heavy segmentation model + +model: + name: segmentation_heavy + predefined_model: + name: SegmentationModel + params: + variant: heavy + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + num_workers: 4 + validation_interval: 10 + num_log_images: 8 + + callbacks: + - name: TestOnTrainEnd + - name: ExportOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0005 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs diff --git a/configs/segmentation_light_model.yaml b/configs/segmentation_light_model.yaml new file mode 100644 index 00000000..c03703f4 --- /dev/null +++ b/configs/segmentation_light_model.yaml @@ -0,0 +1,43 @@ +# Example configuration for training a predefined light segmentation model + +model: + name: segmentation_light + predefined_model: + name: SegmentationModel + params: + variant: light + +loader: + params: + dataset_name: coco_test + +trainer: + preprocessing: + train_image_size: [384, 512] + keep_aspect_ratio: True + normalize: + active: True + + batch_size: 8 + epochs: &epochs 200 + num_workers: 4 + validation_interval: 10 + num_log_images: 8 + + callbacks: + - name: TestOnTrainEnd + - name: ExportOnTrainEnd + + optimizer: + name: SGD + params: + lr: 0.02 + momentum: 0.9 + weight_decay: 0.0005 + nesterov: true + + scheduler: + name: CosineAnnealingLR + params: + T_max: *epochs + eta_min: 0.0002 diff --git a/configs/segmentation_model.yaml b/configs/segmentation_model.yaml deleted file mode 100644 index b403a75e..00000000 --- a/configs/segmentation_model.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# Example configuration for training a predefined segmentation model - -model: - name: coco_segmentation - predefined_model: - name: SegmentationModel - params: - backbone: MicroNet - task: binary - -loader: - params: - dataset_name: coco_test - -trainer: - preprocessing: - train_image_size: [&height 256, &width 320] - keep_aspect_ratio: False - normalize: - active: True - - batch_size: 4 - epochs: &epochs 200 - n_workers: 4 - validation_interval: 10 - n_log_images: 8 - - callbacks: - - name: ExportOnTrainEnd - - name: TestOnTrainEnd - - optimizer: - name: SGD - params: - lr: 0.02 - - scheduler: - name: ConstantLR diff --git a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py index 30369784..7d25a8d2 100644 --- a/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py +++ b/luxonis_train/attached_modules/losses/efficient_keypoint_bbox_loss.py @@ -33,10 +33,10 @@ def __init__( n_warmup_epochs: int = 4, iou_type: IoUType = "giou", reduction: Literal["sum", "mean"] = "mean", - class_loss_weight: float = 1.0, - iou_loss_weight: float = 2.5, + class_loss_weight: float = 0.5, + iou_loss_weight: float = 7.5, viz_pw: float = 1.0, - regr_kpts_loss_weight: float = 1.5, + regr_kpts_loss_weight: float = 12, vis_kpts_loss_weight: float = 1.0, sigmas: list[float] | None = None, area_factor: float | None = None, diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py index 509a95bd..a37f1da9 100644 --- a/luxonis_train/config/config.py +++ b/luxonis_train/config/config.py @@ -365,7 +365,6 @@ class TrainerConfig(BaseModelExtraForbid): NonNegativeInt, Field(validation_alias=AliasChoices("n_workers", "num_workers")), ] = 4 - train_metrics_interval: Literal[-1] | PositiveInt = -1 validation_interval: Literal[-1] | PositiveInt = 5 n_log_images: Annotated[ NonNegativeInt, diff --git a/luxonis_train/config/predefined_models/README.md b/luxonis_train/config/predefined_models/README.md index bdf49178..3733534d 100644 --- a/luxonis_train/config/predefined_models/README.md +++ b/luxonis_train/config/predefined_models/README.md @@ -23,73 +23,81 @@ models which can be used instead. ## SegmentationModel -See an example configuration file using this predefined model [here](../../../configs/segmentation_model.yaml) +The `SegmentationModel` allows for both "light" and "heavy" variants, where the "heavy" variant is more accurate, and the "light" variant is faster. + +See an example configuration file using this predefined model [here](../../../configs/segmentation_light_model.yaml) for the "light" variant, and [here](../../../configs/segmentation_heavy_model.yaml) for the "heavy" variant. **Components** -| Name | Alias | Function | -| --------------------------------------------------------------------------------------------- | -------------------------- | ----------------------------------------------------------------------- | -| [MicroNet](../../nodes/README.md#micronet) | segmentation_backbone | Backbone of the model. Can be changed | -| [SegmentationHead](../../nodes/README.md#segmentationhead) | segmentation_head | Head of the model. | -| [BCEWithLogitsLoss](../../attached_modules/losses/README.md#bcewithlogitsloss) | segmentation_loss | Loss of the model when the task is set to "binary". | -| [CrossEntropyLoss](../../attached_modules/losses/README.md#crossentropyloss) | segmentation_loss | Loss of the model when the task is set to "multiclass" or "multilabel". | -| [JaccardIndex](../../attached_modules/metrics/README.md#torchmetrics) | segmentation_jaccard_index | Main metric of the model. | -| [F1Score](../../attached_modules/metrics/README.md#torchmetrics) | segmentation_f1_score | Secondary metric of the model. | -| [SegmentationVisualizer](../../attached_modules/visualizers/README.md#segmentationvisualizer) | segmentation_visualizer | Visualizer of the `SegmentationHead`. | +| Name | Alias | Function | +| --------------------------------------------------------------------------------------------- | -------------------------- | -------------------------------------------------------------------------------------------- | +| [DDRNet](../../nodes/README.md#ddrnet) | segmentation_backbone | Backbone of the model. Available variants: "light" (DDRNet-23-slim) and "heavy" (DDRNet-23). | +| [SegmentationHead](../../nodes/README.md#segmentationhead) | segmentation_head | Head of the model. | +| [BCEWithLogitsLoss](../../attached_modules/losses/README.md#bcewithlogitsloss) | segmentation_loss | Loss of the model when the task is set to "binary". | +| [CrossEntropyLoss](../../attached_modules/losses/README.md#crossentropyloss) | segmentation_loss | Loss of the model when the task is set to "multiclass" or "multilabel". | +| [JaccardIndex](../../attached_modules/metrics/README.md#torchmetrics) | segmentation_jaccard_index | Main metric of the model. | +| [F1Score](../../attached_modules/metrics/README.md#torchmetrics) | segmentation_f1_score | Secondary metric of the model. | +| [SegmentationVisualizer](../../attached_modules/visualizers/README.md#segmentationvisualizer) | segmentation_visualizer | Visualizer of the `SegmentationHead`. | **Params** -| Key | Type | Default value | Description | -| ----------------- | --------------------------------- | ------------- | ------------------------------------------ | -| task | Literal\["binary", "multiclass"\] | "binary" | Type of the task of the model. | -| task_name | str \| None | None | Custom task name for the head. | -| backbone | str | "MicroNet" | Name of the node to be used as a backbone. | -| backbone_params | dict | {} | Additional parameters to the backbone. | -| head_params | dict | {} | Additional parameters to the head. | -| loss_params | dict | {} | Additional parameters to the loss. | -| visualizer_params | dict | {} | Additional parameters to the visualizer. | +| Key | Type | Default value | Description | +| ----------------- | --------------------------------- | ------------- | ------------------------------------------------------------------------------------------------ | +| variant | Literal\["light", "heavy"\] | "light" | Defines the variant of the model. "light" uses DDRNet-23-slim, "heavy" uses DDRNet-23. | +| backbone | str | "DDRNet" | Name of the node to be used as a backbone. | +| backbone_params | dict | {} | Additional parameters for the backbone. If not provided, variant-specific defaults will be used. | +| head_params | dict | {} | Additional parameters for the head. | +| aux_head_params | dict | {} | Additional parameters for auxiliary heads. | +| loss_params | dict | {} | Additional parameters for the loss. | +| visualizer_params | dict | {} | Additional parameters for the visualizer. | +| task | Literal\["binary", "multiclass"\] | "binary" | Type of the task of the model. | +| task_name | str \| None | None | Custom task name for the head. | ## DetectionModel -See an example configuration file using this predefined model [here](../../../configs/detection_model.yaml) +The `DetectionModel` allows for both "light" and "heavy" variants, where the "heavy" variant is more accurate, and the "light" variant is faster. + +See an example configuration file using this predefined model [here](../../../configs/detection_light_model.yaml) for the "light" variant, and [here](../../../configs/detection_heavy_model.yaml) for the "heavy" variant. **Components** -| Name | Alias | Function | -| -------------------------------------------------------------------------------------- | -------------------- | ----------------------------------- | -| [EfficientRep](../../nodes/README.md#efficientrep) | detection_backbone | Backbone of the model. | -| [RepPANNeck](../../nodes/README.md#reppanneck) | detection_neck | Neck of the model. | -| [EfficientBBoxHead](../../nodes/README.md#efficientbboxhead) | detection_head | Head of the model. | -| [AdaptiveDetectionLoss](../../attached_modules/losses/README.md#adaptivedetectionloss) | detection_loss | Loss of the model. | -| [MeanAveragePrecision](../../attached_modules/metrics/README.md#meanaverageprecision) | detection_map | Main metric of the model. | -| [BBoxVisualizer](../../attached_modules/visualizers/README.md#bboxvisualizer) | detection_visualizer | Visualizer of the `detection_head`. | +| Name | Alias | Function | +| -------------------------------------------------------------------------------------- | -------------------- | ------------------------------------------------------------------------------------------------- | +| [EfficientRep](../../nodes/README.md#efficientrep) | detection_backbone | Backbone of the model. Available variants: "light" (EfficientRep-N) and "heavy" (EfficientRep-L). | +| [RepPANNeck](../../nodes/README.md#reppanneck) | detection_neck | Neck of the model. | +| [EfficientBBoxHead](../../nodes/README.md#efficientbboxhead) | detection_head | Head of the model. | +| [AdaptiveDetectionLoss](../../attached_modules/losses/README.md#adaptivedetectionloss) | detection_loss | Loss of the model. | +| [MeanAveragePrecision](../../attached_modules/metrics/README.md#meanaverageprecision) | detection_map | Main metric of the model. | +| [BBoxVisualizer](../../attached_modules/visualizers/README.md#bboxvisualizer) | detection_visualizer | Visualizer of the `detection_head`. | **Params** -| Key | Type | Default value | Description | -| ----------------- | ----------- | ------------- | ----------------------------------------- | -| task_name | str \| None | None | Custom task name for the head. | -| use_neck | bool | True | Whether to include the neck in the model. | -| backbone_params | dict | {} | Additional parameters to the backbone. | -| neck_params | dict | {} | Additional parameters to the neck. | -| head_params | dict | {} | Additional parameters to the head. | -| loss_params | dict | {} | Additional parameters to the loss. | -| visualizer_params | dict | {} | Additional parameters to the visualizer. | +| Key | Type | Default value | Description | +| ----------------- | --------------------------- | -------------- | ------------------------------------------------------------------------------------------- | +| variant | Literal\["light", "heavy"\] | "light" | Defines the variant of the model. "light" uses EfficientRep-N, "heavy" uses EfficientRep-L. | +| use_neck | bool | True | Whether to include the neck in the model. | +| backbone | str | "EfficientRep" | Name of the node to be used as a backbone. | +| backbone_params | dict | {} | Additional parameters to the backbone. | +| neck_params | dict | {} | Additional parameters to the neck. | +| head_params | dict | {} | Additional parameters to the head. | +| loss_params | dict | {} | Additional parameters to the loss. | +| visualizer_params | dict | {} | Additional parameters to the visualizer. | +| task_name | str \| None | None | Custom task name for the head. | ## KeypointDetectionModel -See an example configuration file using this predefined model [here](../../../configs/keypoint_bbox_model.yaml) +The `KeypointDetectionModel` allows for both "light" and "heavy" variants, where the "heavy" variant is more accurate, and the "light" variant is faster. + +See an example configuration file using this predefined model [here](../../../configs/keypoint_bbox_light_model.yaml) for the "light" variant, and [here](../../../configs/keypoint_bbox_heavy_model.yaml) for the "heavy" variant. **Components** | Name | Alias | Function | | ------------------------------------------------------------------------------------------------------- | ---------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | -| [EfficientRep](../../nodes/README.md#efficientrep) | kpt_detection_backbone | Backbone of the model. | +| [EfficientRep](../../nodes/README.md#efficientrep) | kpt_detection_backbone | Backbone of the model.. Available variants: "light" (EfficientRep-N) and "heavy" (EfficientRep-L). | | [RepPANNeck](../../nodes/README.md#reppanneck) | kpt_detection_neck | Neck of the model. | -| [ImplicitKeypointBBoxHead](../../nodes/README.md#implicitkeypointbboxhead) | kpt_detection_head | Possible head of the model, changes depending on the value of `head_type` argument. | -| [EfficientKeypointBBoxHead](../../nodes/README.md#efficientkeypointbboxhead) | kpt_detection_head | Possible head of the model, changes depending on the value of `head_type` argument | -| [ImplicitKeypointBBoxLoss](../../attached_modules/losses/README.md#implicitkeypointbboxloss) | kpt_detection_loss | Loss of the model if the `head_type` is set to "ImplicitKeypointBBoxHead" | -| [EfficientKeypointBBoxLoss](../../attached_modules/losses/README.md#efficientkeypointbboxloss) | kpt_detection_loss | Loss of the model if `head_type` is set to "EfficientKeypointBBoxHead". | +| [EfficientKeypointBBoxHead](../../nodes/README.md#efficientkeypointbboxhead) | kpt_detection_head | Head of the model. | +| [EfficientKeypointBBoxLoss](../../attached_modules/losses/README.md#efficientkeypointbboxloss) | kpt_detection_loss | Loss of the model. | | [ObjectKeypointSimilarity](../../attached_modules/metrics/README.md#objectkeypointsimilarity) | kpt_detection_oks | Main metric of the model. | | [MeanAveragePrecisionKeypoints](../../attached_modules/metrics/README.md#meanaverageprecisionkeypoints) | kpt_detection_map | Secondary metric of the model. | | [BBoxVisualizer](../../attached_modules/visualizers/README.md#bboxvisualizer) | | Visualizer for bounding boxes. Combined with keypoint visualizer in [MultiVisualizer](../../attached_modules/visualizers/README.md#multivisualizer). | @@ -97,44 +105,46 @@ See an example configuration file using this predefined model [here](../../../co **Params** -| Key | Type | Default value | Description | -| ---------------------- | ----------- | --------------------------------------------------------- | ------------------------------------------------- | -| use_neck | bool | True | Whether to include the neck in the model. | -| backbone_params | dict | {} | Additional parameters to the backbone. | -| neck_params | dict | {} | Additional parameters to the neck. | -| head_params | dict | {} | Additional parameters to the head. | -| head_type | str | "ImplicitKeypointBBoxHead" \| "EfficientKeypointBBoxHead" | Type of the head. | -| loss_params | dict | {} | Additional parameters to the loss. | -| kpt_visualizer_params | dict | {} | Additional parameters to the keypoint visualizer. | -| bbox_visualizer_params | dict | {} | Additional parameters to the bbox visualizer. | -| bbox_task_name | str \| None | None | Custom task name for the detection head. | -| kpt_task_name | str \| None | None | Custom task name for the keypoint head. | +| Key | Type | Default value | Description | +| ---------------------- | --------------------------- | -------------- | ------------------------------------------------------------------------------------------- | +| variant | Literal\["light", "heavy"\] | "light" | Defines the variant of the model. "light" uses EfficientRep-N, "heavy" uses EfficientRep-L. | +| use_neck | bool | True | Whether to include the neck in the model. | +| backbone | str | "EfficientRep" | Name of the node to be used as a backbone. | +| backbone_params | dict | {} | Additional parameters to the backbone. | +| neck_params | dict | {} | Additional parameters to the neck. | +| head_params | dict | {} | Additional parameters to the head. | +| loss_params | dict | {} | Additional parameters to the loss. | +| kpt_visualizer_params | dict | {} | Additional parameters to the keypoint visualizer. | +| bbox_visualizer_params | dict | {} | Additional parameters to the bbox visualizer. | +| bbox_task_name | str \| None | None | Custom task name for the detection head. | +| kpt_task_name | str \| None | None | Custom task name for the keypoint head. | ## ClassificationModel -Basic model for classification. Can be used for multiclass and multilabel tasks. +The `ClassificationModel` allows for both "light" and "heavy" variants, where the "heavy" variant is more accurate, and the "light" variant is faster. Can be used for multiclass and multilabel tasks. -See an example configuration file using this predefined model [here](../../../configs/classification_model.yaml) +See an example configuration file using this predefined model [here](../../../configs/classification_light_model.yaml) for the "light" variant, and [here](../../../configs/classification_heavy_model.yaml) for the "heavy" variant. **Components** -| Name | Alias | Function | -| ---------------------------------------------------------------------------- | ----------------------- | ------------------------------------- | -| [MicroNet](../../nodes/README.md#micronet) | classification_backbone | Backbone of the model. Can be changed | -| [ClassificationHead](../../nodes/README.md#classificationhead) | classification_head | Head of the model. | -| [CrossEntropyLoss](../../attached_modules/losses/README.md#crossentropyloss) | classification_loss | Loss of the model. | -| [F1Score](../../attached_modules/metrics/README.md#torchmetrics) | classification_f1_score | Main metric of the model. | -| [Accuracy](../../attached_modules/metrics/README.md#torchmetrics) | classification_accuracy | Secondary metric of the model. | -| [Recall](../../attached_modules/metrics/README.md#torchmetrics) | classification_recall | Secondary metric of the model. | +| Name | Alias | Function | +| ---------------------------------------------------------------------------- | ----------------------- | ----------------------------------------------------------------------------------------------------- | +| [ResNet](../../nodes/README.md#resnet) | classification_backbone | Backbone of the model. The "light" variant uses ResNet-18, while the "heavy" variant uses ResNet-101. | +| [ClassificationHead](../../nodes/README.md#classificationhead) | classification_head | Head of the model. | +| [CrossEntropyLoss](../../attached_modules/losses/README.md#crossentropyloss) | classification_loss | Loss of the model. | +| [F1Score](../../attached_modules/metrics/README.md#torchmetrics) | classification_f1_score | Main metric of the model. | +| [Accuracy](../../attached_modules/metrics/README.md#torchmetrics) | classification_accuracy | Secondary metric of the model. | +| [Recall](../../attached_modules/metrics/README.md#torchmetrics) | classification_recall | Secondary metric of the model. | **Params** -| Key | Type | Default value | Description | -| ----------------- | ------------------------------------- | ------------- | ------------------------------------------ | -| task | Literal\["multiclass", "multilabel"\] | "multiclass" | Type of the task of the model. | -| task_name | str \| None | None | Custom task name for the head. | -| backbone | str | "MicroNet" | Name of the node to be used as a backbone. | -| backbone_params | dict | {} | Additional parameters to the backbone. | -| head_params | dict | {} | Additional parameters to the head. | -| loss_params | dict | {} | Additional parameters to the loss. | -| visualizer_params | dict | {} | Additional parameters to the visualizer. | +| Key | Type | Default value | Description | +| ----------------- | ------------------------------------- | ------------- | ----------------------------------------------------------------------------------- | +| variant | Literal\["light", "heavy"\] | "light" | Defines the variant of the model. "light" uses ResNet-18, "heavy" uses ResNet-101. | +| backbone | str | "ResNet" | Name of the node to be used as a backbone. | +| backbone_params | dict | {} | Additional parameters to the backbone. | +| head_params | dict | {} | Additional parameters to the head. | +| loss_params | dict | {} | Additional parameters to the loss. | +| visualizer_params | dict | {} | Additional parameters to the visualizer. | +| task | Literal\["multiclass", "multilabel"\] | "multiclass" | Type of the task of the model. | +| task_name | str \| None | None | Custom task name for the head. | diff --git a/luxonis_train/config/predefined_models/__init__.py b/luxonis_train/config/predefined_models/__init__.py index c52e359d..0fc4e6da 100644 --- a/luxonis_train/config/predefined_models/__init__.py +++ b/luxonis_train/config/predefined_models/__init__.py @@ -1,15 +1,13 @@ from .base_predefined_model import BasePredefinedModel from .classification_model import ClassificationModel -from .ddrnet_segmentation_model import DDRNetSegmentationModel from .detection_model import DetectionModel from .keypoint_detection_model import KeypointDetectionModel from .segmentation_model import SegmentationModel __all__ = [ "BasePredefinedModel", - "SegmentationModel", "DetectionModel", "KeypointDetectionModel", "ClassificationModel", - "DDRNetSegmentationModel", + "SegmentationModel", ] diff --git a/luxonis_train/config/predefined_models/classification_model.py b/luxonis_train/config/predefined_models/classification_model.py index d76df5c0..37225207 100644 --- a/luxonis_train/config/predefined_models/classification_model.py +++ b/luxonis_train/config/predefined_models/classification_model.py @@ -1,5 +1,6 @@ -from dataclasses import dataclass, field -from typing import Literal +from typing import Literal, TypeAlias + +from pydantic import BaseModel from luxonis_train.config import ( AttachedModuleConfig, @@ -11,19 +12,65 @@ from .base_predefined_model import BasePredefinedModel +VariantLiteral: TypeAlias = Literal["light", "heavy"] + + +class ClassificationVariant(BaseModel): + backbone: str + backbone_params: Params + + +def get_variant(variant: VariantLiteral) -> ClassificationVariant: + """Returns the specific variant configuration for the + ClassificationModel.""" + variants = { + "light": ClassificationVariant( + backbone="ResNet", + backbone_params={"variant": "18"}, + ), + "heavy": ClassificationVariant( + backbone="ResNet", + backbone_params={"variant": "101"}, + ), + } + + if variant not in variants: + raise ValueError( + f"Classification variant should be one of {list(variants.keys())}, got '{variant}'." + ) + + return variants[variant] + -@dataclass class ClassificationModel(BasePredefinedModel): - backbone: str = "MicroNet" - task: Literal["multiclass", "multilabel"] = "multiclass" - backbone_params: Params = field(default_factory=dict) - head_params: Params = field(default_factory=dict) - loss_params: Params = field(default_factory=dict) - visualizer_params: Params = field(default_factory=dict) - task_name: str | None = None + def __init__( + self, + variant: VariantLiteral = "light", + backbone: str | None = None, + backbone_params: Params | None = None, + head_params: Params | None = None, + loss_params: Params | None = None, + visualizer_params: Params | None = None, + task: Literal["multiclass", "multilabel"] = "multiclass", + task_name: str | None = None, + ): + var_config = get_variant(variant) + + self.backbone = backbone or var_config.backbone + self.backbone_params = ( + backbone_params + if backbone is not None or backbone_params is not None + else var_config.backbone_params + ) or {} + self.head_params = head_params or {} + self.loss_params = loss_params or {} + self.visualizer_params = visualizer_params or {} + self.task = task + self.task_name = task_name or "classification" @property def nodes(self) -> list[ModelNodeConfig]: + """Defines the model nodes, including backbone and head.""" return [ ModelNodeConfig( name=self.backbone, @@ -43,6 +90,7 @@ def nodes(self) -> list[ModelNodeConfig]: @property def losses(self) -> list[LossModuleConfig]: + """Defines the loss module for the classification task.""" return [ LossModuleConfig( name="CrossEntropyLoss", @@ -55,6 +103,7 @@ def losses(self) -> list[LossModuleConfig]: @property def metrics(self) -> list[MetricModuleConfig]: + """Defines the metrics used for evaluation.""" return [ MetricModuleConfig( name="F1Score", @@ -79,6 +128,7 @@ def metrics(self) -> list[MetricModuleConfig]: @property def visualizers(self) -> list[AttachedModuleConfig]: + """Defines the visualizer used for the classification task.""" return [ AttachedModuleConfig( name="ClassificationVisualizer", diff --git a/luxonis_train/config/predefined_models/ddrnet_segmentation_model.py b/luxonis_train/config/predefined_models/ddrnet_segmentation_model.py deleted file mode 100644 index e991208b..00000000 --- a/luxonis_train/config/predefined_models/ddrnet_segmentation_model.py +++ /dev/null @@ -1,73 +0,0 @@ -from dataclasses import dataclass, field - -from luxonis_train.config import LossModuleConfig, ModelNodeConfig, Params - -from .segmentation_model import SegmentationModel - - -@dataclass -class DDRNetSegmentationModel(SegmentationModel): - backbone: str = "DDRNet" - aux_head_params: Params = field(default_factory=dict) - - @property - def nodes(self) -> list[ModelNodeConfig]: - self.head_params.update({"attach_index": -1}) - - self.aux_head_params.update({"attach_index": -2}) - - node_list = [ - ModelNodeConfig( - name=self.backbone, - alias="ddrnet_backbone", - freezing=self.backbone_params.pop("freezing", {}), - params=self.backbone_params, - ), - ModelNodeConfig( - name="DDRNetSegmentationHead", - alias="segmentation_head", - inputs=["ddrnet_backbone"], - freezing=self.head_params.pop("freezing", {}), - params=self.head_params, - task=self.task_name, - ), - ] - if self.backbone_params.get("use_aux_heads", False): - node_list.append( - ModelNodeConfig( - name="DDRNetSegmentationHead", - alias="aux_segmentation_head", - inputs=["ddrnet_backbone"], - freezing=self.aux_head_params.pop("freezing", {}), - params=self.aux_head_params, - task=self.task_name, - ) - ) - return node_list - - @property - def losses(self) -> list[LossModuleConfig]: - loss_list = [ - LossModuleConfig( - name="BCEWithLogitsLoss" - if self.task == "binary" - else "CrossEntropyLoss", - alias="segmentation_loss", - attached_to="segmentation_head", - params=self.loss_params, - weight=1.0, - ), - ] - if self.backbone_params.get("use_aux_heads", False): - loss_list.append( - LossModuleConfig( - name="BCEWithLogitsLoss" - if self.task == "binary" - else "CrossEntropyLoss", - alias="aux_segmentation_loss", - attached_to="aux_segmentation_head", - params=self.loss_params, - weight=0.4, - ) - ) - return loss_list diff --git a/luxonis_train/config/predefined_models/detection_model.py b/luxonis_train/config/predefined_models/detection_model.py index f205a993..66ca00aa 100644 --- a/luxonis_train/config/predefined_models/detection_model.py +++ b/luxonis_train/config/predefined_models/detection_model.py @@ -1,4 +1,6 @@ -from dataclasses import dataclass, field +from typing import Literal, TypeAlias + +from pydantic import BaseModel from luxonis_train.config import ( AttachedModuleConfig, @@ -10,22 +12,71 @@ from .base_predefined_model import BasePredefinedModel +VariantLiteral: TypeAlias = Literal["light", "heavy"] + + +class DetectionVariant(BaseModel): + backbone: str + backbone_params: Params + + +def get_variant(variant: VariantLiteral) -> DetectionVariant: + """Returns the specific variant configuration for the + DetectionModel.""" + variants = { + "light": DetectionVariant( + backbone="EfficientRep", + backbone_params={"variant": "n"}, + ), + "heavy": DetectionVariant( + backbone="EfficientRep", + backbone_params={"variant": "l"}, + ), + } + + if variant not in variants: + raise ValueError( + f"Detection variant should be one of {list(variants.keys())}, got '{variant}'." + ) + + return variants[variant] + -@dataclass class DetectionModel(BasePredefinedModel): - use_neck: bool = True - backbone_params: Params = field(default_factory=dict) - neck_params: Params = field(default_factory=dict) - head_params: Params = field(default_factory=dict) - loss_params: Params = field(default_factory=dict) - visualizer_params: Params = field(default_factory=dict) - task_name: str | None = None + def __init__( + self, + variant: VariantLiteral = "light", + use_neck: bool = True, + backbone: str | None = None, + backbone_params: Params | None = None, + neck_params: Params | None = None, + head_params: Params | None = None, + loss_params: Params | None = None, + visualizer_params: Params | None = None, + task_name: str | None = None, + ): + var_config = get_variant(variant) + + self.use_neck = use_neck + self.backbone_params = ( + backbone_params + if backbone is not None or backbone_params is not None + else var_config.backbone_params + ) or {} + self.backbone = backbone or var_config.backbone + self.neck_params = neck_params or {} + self.head_params = head_params or {} + self.loss_params = loss_params or {"n_warmup_epochs": 0} + self.visualizer_params = visualizer_params or {} + self.task_name = task_name or "boundingbox" @property def nodes(self) -> list[ModelNodeConfig]: + """Defines the model nodes, including backbone, neck, and + head.""" nodes = [ ModelNodeConfig( - name="EfficientRep", + name=self.backbone, alias="detection_backbone", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, @@ -58,6 +109,7 @@ def nodes(self) -> list[ModelNodeConfig]: @property def losses(self) -> list[LossModuleConfig]: + """Defines the loss module for the detection task.""" return [ LossModuleConfig( name="AdaptiveDetectionLoss", @@ -70,6 +122,7 @@ def losses(self) -> list[LossModuleConfig]: @property def metrics(self) -> list[MetricModuleConfig]: + """Defines the metrics used for evaluation.""" return [ MetricModuleConfig( name="MeanAveragePrecision", @@ -81,6 +134,7 @@ def metrics(self) -> list[MetricModuleConfig]: @property def visualizers(self) -> list[AttachedModuleConfig]: + """Defines the visualizer used for the detection task.""" return [ AttachedModuleConfig( name="BBoxVisualizer", diff --git a/luxonis_train/config/predefined_models/keypoint_detection_model.py b/luxonis_train/config/predefined_models/keypoint_detection_model.py index c7206ab2..7ddf6d6d 100644 --- a/luxonis_train/config/predefined_models/keypoint_detection_model.py +++ b/luxonis_train/config/predefined_models/keypoint_detection_model.py @@ -1,5 +1,6 @@ -from dataclasses import dataclass, field -from typing import Literal +from typing import Literal, TypeAlias + +from pydantic import BaseModel from luxonis_train.config import ( AttachedModuleConfig, @@ -11,27 +12,75 @@ from .base_predefined_model import BasePredefinedModel +VariantLiteral: TypeAlias = Literal["light", "heavy"] + + +class KeypointDetectionVariant(BaseModel): + backbone: str + backbone_params: Params + + +def get_variant(variant: VariantLiteral) -> KeypointDetectionVariant: + """Returns the specific variant configuration for the + KeypointDetectionModel.""" + variants = { + "light": KeypointDetectionVariant( + backbone="EfficientRep", + backbone_params={"variant": "n"}, + ), + "heavy": KeypointDetectionVariant( + backbone="EfficientRep", + backbone_params={"variant": "l"}, + ), + } + + if variant not in variants: + raise ValueError( + f"KeypointDetection variant should be one of {list(variants.keys())}, got '{variant}'." + ) + + return variants[variant] + -@dataclass class KeypointDetectionModel(BasePredefinedModel): - use_neck: bool = True - backbone_params: Params = field(default_factory=dict) - neck_params: Params = field(default_factory=dict) - head_params: Params = field(default_factory=dict) - loss_params: Params = field(default_factory=dict) - head_type: Literal[ - "ImplicitKeypointBBoxHead", "EfficientKeypointBBoxHead" - ] = "EfficientKeypointBBoxHead" - kpt_visualizer_params: Params = field(default_factory=dict) - bbox_visualizer_params: Params = field(default_factory=dict) - bbox_task_name: str | None = None - kpt_task_name: str | None = None + def __init__( + self, + variant: VariantLiteral = "light", + use_neck: bool = True, + backbone: str | None = None, + backbone_params: Params | None = None, + neck_params: Params | None = None, + head_params: Params | None = None, + loss_params: Params | None = None, + kpt_visualizer_params: Params | None = None, + bbox_visualizer_params: Params | None = None, + bbox_task_name: str | None = None, + kpt_task_name: str | None = None, + ): + var_config = get_variant(variant) + + self.use_neck = use_neck + self.backbone = backbone or var_config.backbone + self.backbone_params = ( + backbone_params + if backbone is not None or backbone_params is not None + else var_config.backbone_params + ) or {} + self.neck_params = neck_params or {} + self.head_params = head_params or {} + self.loss_params = loss_params or {"n_warmup_epochs": 0} + self.kpt_visualizer_params = kpt_visualizer_params or {} + self.bbox_visualizer_params = bbox_visualizer_params or {} + self.bbox_task_name = bbox_task_name + self.kpt_task_name = kpt_task_name @property def nodes(self) -> list[ModelNodeConfig]: + """Defines the model nodes, including backbone, neck, and + head.""" nodes = [ ModelNodeConfig( - name="EfficientRep", + name=self.backbone, alias="kpt_detection_backbone", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, @@ -56,11 +105,13 @@ def nodes(self) -> list[ModelNodeConfig]: nodes.append( ModelNodeConfig( - name=self.head_type, + name="EfficientKeypointBBoxHead", alias="kpt_detection_head", - inputs=["kpt_detection_neck"] - if self.use_neck - else ["kpt_detection_backbone"], + inputs=( + ["kpt_detection_neck"] + if self.use_neck + else ["kpt_detection_backbone"] + ), freezing=self.head_params.pop("freezing", {}), params=self.head_params, task=task, @@ -70,9 +121,10 @@ def nodes(self) -> list[ModelNodeConfig]: @property def losses(self) -> list[LossModuleConfig]: + """Defines the loss module for the keypoint detection task.""" return [ LossModuleConfig( - name=self.head_type.replace("Head", "Loss"), + name="EfficientKeypointBBoxLoss", attached_to="kpt_detection_head", params=self.loss_params, weight=1.0, @@ -81,6 +133,7 @@ def losses(self) -> list[LossModuleConfig]: @property def metrics(self) -> list[MetricModuleConfig]: + """Defines the metrics used for evaluation.""" return [ MetricModuleConfig( name="ObjectKeypointSimilarity", @@ -97,6 +150,8 @@ def metrics(self) -> list[MetricModuleConfig]: @property def visualizers(self) -> list[AttachedModuleConfig]: + """Defines the visualizer used for the keypoint detection + task.""" return [ AttachedModuleConfig( name="MultiVisualizer", diff --git a/luxonis_train/config/predefined_models/segmentation_model.py b/luxonis_train/config/predefined_models/segmentation_model.py index 5ec34be1..8a281131 100644 --- a/luxonis_train/config/predefined_models/segmentation_model.py +++ b/luxonis_train/config/predefined_models/segmentation_model.py @@ -1,6 +1,7 @@ -from dataclasses import dataclass, field from typing import Literal +from pydantic import BaseModel + from luxonis_train.config import ( AttachedModuleConfig, LossModuleConfig, @@ -11,52 +12,139 @@ from .base_predefined_model import BasePredefinedModel +VariantLiteral = Literal["light", "heavy"] + + +class SegmentationVariant(BaseModel): + backbone: str + backbone_params: Params + + +def get_variant(variant: VariantLiteral) -> SegmentationVariant: + """Returns the specific variant configuration for the + SegmentationModel.""" + variants = { + "light": SegmentationVariant( + backbone="DDRNet", + backbone_params={"variant": "23-slim"}, + ), + "heavy": SegmentationVariant( + backbone="DDRNet", + backbone_params={"variant": "23"}, + ), + } + + if variant not in variants: + raise ValueError( + f"Segmentation variant should be one of {list(variants.keys())}, got '{variant}'." + ) + + return variants[variant] + -@dataclass class SegmentationModel(BasePredefinedModel): - backbone: str = "MicroNet" - task: Literal["binary", "multiclass"] = "binary" - backbone_params: Params = field(default_factory=dict) - head_params: Params = field(default_factory=dict) - loss_params: Params = field(default_factory=dict) - visualizer_params: Params = field(default_factory=dict) - task_name: str | None = None + def __init__( + self, + variant: VariantLiteral = "light", + backbone: str | None = None, + backbone_params: Params | None = None, + head_params: Params | None = None, + aux_head_params: Params | None = None, + loss_params: Params | None = None, + visualizer_params: Params | None = None, + task: Literal["binary", "multiclass"] = "binary", + task_name: str | None = None, + ): + var_config = get_variant(variant) + + self.backbone = backbone or var_config.backbone + self.backbone_params = ( + backbone_params + if backbone is not None or backbone_params is not None + else var_config.backbone_params + ) or {} + self.head_params = head_params or {} + self.aux_head_params = aux_head_params or {} + self.loss_params = loss_params or {} + self.visualizer_params = visualizer_params or {} + self.task = task + self.task_name = task_name or "segmentation" @property def nodes(self) -> list[ModelNodeConfig]: - return [ + """Defines the model nodes, including backbone and head.""" + self.head_params.update({"attach_index": -1}) + self.aux_head_params.update({"attach_index": -2}) + ( + self.aux_head_params.update({"remove_on_export": True}) + if "remove_on_export" not in self.aux_head_params + else None + ) + + node_list = [ ModelNodeConfig( name=self.backbone, - alias="segmentation_backbone", + alias="ddrnet_backbone", freezing=self.backbone_params.pop("freezing", {}), params=self.backbone_params, ), ModelNodeConfig( - name="SegmentationHead", + name="DDRNetSegmentationHead", alias="segmentation_head", - inputs=["segmentation_backbone"], + inputs=["ddrnet_backbone"], freezing=self.head_params.pop("freezing", {}), params=self.head_params, task=self.task_name, ), ] + if self.backbone_params.get("use_aux_heads", True): + node_list.append( + ModelNodeConfig( + name="DDRNetSegmentationHead", + alias="aux_segmentation_head", + inputs=["ddrnet_backbone"], + freezing=self.aux_head_params.pop("freezing", {}), + params=self.aux_head_params, + task=self.task_name, + ) + ) + return node_list @property def losses(self) -> list[LossModuleConfig]: - return [ + """Defines the loss module for the segmentation task.""" + loss_list = [ LossModuleConfig( - name="BCEWithLogitsLoss" - if self.task == "binary" - else "CrossEntropyLoss", + name=( + "BCEWithLogitsLoss" + if self.task == "binary" + else "CrossEntropyLoss" + ), alias="segmentation_loss", attached_to="segmentation_head", params=self.loss_params, weight=1.0, - ) + ), ] + if self.backbone_params.get("use_aux_heads", False): + loss_list.append( + LossModuleConfig( + name=( + "BCEWithLogitsLoss" + if self.task == "binary" + else "CrossEntropyLoss" + ), + alias="aux_segmentation_loss", + attached_to="aux_segmentation_head", + params=self.loss_params, + weight=0.4, + ) + ) + return loss_list @property def metrics(self) -> list[MetricModuleConfig]: + """Defines the metrics used for evaluation.""" return [ MetricModuleConfig( name="JaccardIndex", @@ -75,6 +163,7 @@ def metrics(self) -> list[MetricModuleConfig]: @property def visualizers(self) -> list[AttachedModuleConfig]: + """Defines the visualizer used for the segmentation task.""" return [ AttachedModuleConfig( name="SegmentationVisualizer", diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py index 33df65db..a319061e 100644 --- a/luxonis_train/models/luxonis_lightning.py +++ b/luxonis_train/models/luxonis_lightning.py @@ -298,10 +298,15 @@ def _initiate_nodes( for source_name, shape in shapes.items() } - for node_name, ( - Node, - node_kwargs, - ), node_input_names, _ in traverse_graph(self.graph, nodes): + for ( + node_name, + ( + Node, + node_kwargs, + ), + node_input_names, + _, + ) in traverse_graph(self.graph, nodes): node_dummy_inputs: list[Packet[Tensor]] = [] """List of dummy input packets for the node. @@ -913,16 +918,6 @@ def _print_results( f"{stage} main metric ({self.main_metric}): {main_metric:.4f}" ) - def _is_train_eval_epoch(self) -> bool: - """Checks if train eval should be performed on current epoch - based on configured train_metrics_interval.""" - train_metrics_interval = self.cfg.trainer.train_metrics_interval - # add +1 to current_epoch because starting epoch is at 0 - return ( - train_metrics_interval != -1 - and (self.current_epoch + 1) % train_metrics_interval == 0 - ) - def _average_losses( self, step_outputs: list[Mapping[str, Tensor | float | int]] ) -> dict[str, float]: diff --git a/configs/example_multi_input.yaml b/tests/configs/multi_input.yaml similarity index 98% rename from configs/example_multi_input.yaml rename to tests/configs/multi_input.yaml index 9632ed43..7db03d90 100644 --- a/configs/example_multi_input.yaml +++ b/tests/configs/multi_input.yaml @@ -111,4 +111,4 @@ trainer: exporter: onnx: - opset_version: 11 + opset_version: 11 \ No newline at end of file diff --git a/tests/configs/parking_lot_config.yaml b/tests/configs/parking_lot_config.yaml index bb15ac37..bf0b9da3 100644 --- a/tests/configs/parking_lot_config.yaml +++ b/tests/configs/parking_lot_config.yaml @@ -145,7 +145,6 @@ trainer: accumulate_grad_batches: 1 epochs: 200 n_workers: 8 - train_metrics_interval: -1 validation_interval: 10 n_log_images: 8 skip_last_batch: True diff --git a/tests/integration/test_simple.py b/tests/integration/test_simple.py index 069e53b0..7b2f0f91 100644 --- a/tests/integration/test_simple.py +++ b/tests/integration/test_simple.py @@ -43,11 +43,14 @@ def clear_files(): @pytest.mark.parametrize( "config_file", [ - "classification_model", - "segmentation_model", - "detection_model", - "keypoint_bbox_model", - "ddrnet_segmentation_model", + "classification_heavy_model", + "classification_light_model", + "segmentation_heavy_model", + "segmentation_light_model", + "detection_heavy_model", + "detection_light_model", + "keypoint_bbox_heavy_model", + "keypoint_bbox_light_model", ], ) def test_predefined_models( @@ -58,9 +61,11 @@ def test_predefined_models( ): config_file = f"configs/{config_file}.yaml" opts |= { - "loader.params.dataset_name": cifar10_dataset.dataset_name - if "classification_model" in config_file - else coco_dataset.dataset_name, + "loader.params.dataset_name": ( + cifar10_dataset.dataset_name + if "classification" in config_file + else coco_dataset.dataset_name + ), } model = LuxonisModel(config_file, opts) model.train() @@ -68,7 +73,7 @@ def test_predefined_models( def test_multi_input(opts: dict[str, Any]): - config_file = "configs/example_multi_input.yaml" + config_file = "tests/configs/multi_input.yaml" model = LuxonisModel(config_file, opts) model.train() model.test(view="val") @@ -198,7 +203,7 @@ def test_callbacks(opts: dict[str, Any], parking_lot_dataset: LuxonisDataset): def test_freezing(opts: dict[str, Any], coco_dataset: LuxonisDataset): - config_file = "configs/segmentation_model.yaml" + config_file = "configs/segmentation_light_model.yaml" opts = deepcopy(opts) opts |= { "model.predefined_model.params": {