Skip to content

Commit

Permalink
feat: add dynamic aliases to all predefined models
Browse files Browse the repository at this point in the history
  • Loading branch information
@sokovninn
sokovninn committed Oct 29, 2024
1 parent 70adfe8 commit 163fbae
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 42 deletions.
25 changes: 13 additions & 12 deletions luxonis_train/config/predefined_models/classification_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,14 +74,14 @@ def nodes(self) -> list[ModelNodeConfig]:
return [
ModelNodeConfig(
name=self.backbone,
alias="classification_backbone",
alias=f"{self.backbone}-{self.task_name}",
freezing=self.backbone_params.pop("freezing", {}),
params=self.backbone_params,
),
ModelNodeConfig(
name="ClassificationHead",
alias="classification_head",
inputs=["classification_backbone"],
alias=f"ClassificationHead-{self.task_name}",
inputs=["f{self.backbone}-{self.task_name}"],
freezing=self.head_params.pop("freezing", {}),
params=self.head_params,
task=self.task_name,
Expand All @@ -94,8 +94,8 @@ def losses(self) -> list[LossModuleConfig]:
return [
LossModuleConfig(
name="CrossEntropyLoss",
alias="classification_loss",
attached_to="classification_head",
alias=f"CrossEntropyLoss-{self.task_name}",
attached_to=f"ClassificationHead-{self.task_name}",
params=self.loss_params,
weight=1.0,
)
Expand All @@ -107,21 +107,21 @@ def metrics(self) -> list[MetricModuleConfig]:
return [
MetricModuleConfig(
name="F1Score",
alias="classification_f1_score",
alias=f"F1Score-{self.task_name}",
is_main_metric=True,
attached_to="classification_head",
attached_to=f"ClassificationHead-{self.task_name}",
params={"task": self.task},
),
MetricModuleConfig(
name="Accuracy",
alias="classification_accuracy",
attached_to="classification_head",
alias=f"Accuracy-{self.task_name}",
attached_to=f"ClassificationHead-{self.task_name}",
params={"task": self.task},
),
MetricModuleConfig(
name="Recall",
alias="classification_recall",
attached_to="classification_head",
alias=f"Recall-{self.task_name}",
attached_to=f"ClassificationHead-{self.task_name}",
params={"task": self.task},
),
]
Expand All @@ -132,7 +132,8 @@ def visualizers(self) -> list[AttachedModuleConfig]:
return [
AttachedModuleConfig(
name="ClassificationVisualizer",
attached_to="classification_head",
alias=f"ClassificationVisualizer-{self.task_name}",
attached_to=f"ClassificationHead-{self.task_name}",
params=self.visualizer_params,
)
]
31 changes: 16 additions & 15 deletions luxonis_train/config/predefined_models/keypoint_detection_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,8 +71,8 @@ def __init__(
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
self.bbox_task_name = bbox_task_name or "boundingbox"
self.kpt_task_name = kpt_task_name or "keypoint_detection"

@property
def nodes(self) -> list[ModelNodeConfig]:
Expand All @@ -81,7 +81,7 @@ def nodes(self) -> list[ModelNodeConfig]:
nodes = [
ModelNodeConfig(
name=self.backbone,
alias="kpt_detection_backbone",
alias=f"{self.backbone}-{self.kpt_task_name}",
freezing=self.backbone_params.pop("freezing", {}),
params=self.backbone_params,
),
Expand All @@ -90,8 +90,8 @@ def nodes(self) -> list[ModelNodeConfig]:
nodes.append(
ModelNodeConfig(
name="RepPANNeck",
alias="kpt_detection_neck",
inputs=["kpt_detection_backbone"],
alias=f"RepPANNeck-{self.kpt_task_name}",
inputs=[f"{self.backbone}-{self.kpt_task_name}"],
freezing=self.neck_params.pop("freezing", {}),
params=self.neck_params,
)
Expand All @@ -106,11 +106,11 @@ def nodes(self) -> list[ModelNodeConfig]:
nodes.append(
ModelNodeConfig(
name="EfficientKeypointBBoxHead",
alias="kpt_detection_head",
alias=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
inputs=(
["kpt_detection_neck"]
[f"RepPANNeck-{self.kpt_task_name}"]
if self.use_neck
else ["kpt_detection_backbone"]
else [f"{self.backbone}-{self.kpt_task_name}"]
),
freezing=self.head_params.pop("freezing", {}),
params=self.head_params,
Expand All @@ -125,7 +125,8 @@ def losses(self) -> list[LossModuleConfig]:
return [
LossModuleConfig(
name="EfficientKeypointBBoxLoss",
attached_to="kpt_detection_head",
alias=f"EfficientKeypointBBoxLoss-{self.kpt_task_name}",
attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
params=self.loss_params,
weight=1.0,
)
Expand All @@ -137,14 +138,14 @@ def metrics(self) -> list[MetricModuleConfig]:
return [
MetricModuleConfig(
name="ObjectKeypointSimilarity",
alias="kpt_detection_oks",
attached_to="kpt_detection_head",
alias=f"ObjectKeypointSimilarity-{self.kpt_task_name}",
attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
is_main_metric=True,
),
MetricModuleConfig(
name="MeanAveragePrecisionKeypoints",
alias="kpt_detection_map",
attached_to="kpt_detection_head",
alias=f"MeanAveragePrecisionKeypoints-{self.kpt_task_name}",
attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
),
]

Expand All @@ -155,8 +156,8 @@ def visualizers(self) -> list[AttachedModuleConfig]:
return [
AttachedModuleConfig(
name="MultiVisualizer",
alias="kpt_detection_visualizer",
attached_to="kpt_detection_head",
alias=f"MultiVisualizer-{self.kpt_task_name}",
attached_to=f"EfficientKeypointBBoxHead-{self.kpt_task_name}",
params={
"visualizers": [
{
Expand Down
38 changes: 23 additions & 15 deletions luxonis_train/config/predefined_models/segmentation_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,14 +79,14 @@ def nodes(self) -> list[ModelNodeConfig]:
node_list = [
ModelNodeConfig(
name=self.backbone,
alias="ddrnet_backbone",
alias=f"{self.backbone}-{self.task_name}",
freezing=self.backbone_params.pop("freezing", {}),
params=self.backbone_params,
),
ModelNodeConfig(
name="DDRNetSegmentationHead",
alias="segmentation_head",
inputs=["ddrnet_backbone"],
alias=f"DDRNetSegmentationHead-{self.task_name}",
inputs=[f"{self.backbone}-{self.task_name}"],
freezing=self.head_params.pop("freezing", {}),
params=self.head_params,
task=self.task_name,
Expand All @@ -96,8 +96,8 @@ def nodes(self) -> list[ModelNodeConfig]:
node_list.append(
ModelNodeConfig(
name="DDRNetSegmentationHead",
alias="aux_segmentation_head",
inputs=["ddrnet_backbone"],
alias=f"DDRNetSegmentationHead_aux-{self.task_name}",
inputs=[f"{self.backbone}-{self.task_name}"],
freezing=self.aux_head_params.pop("freezing", {}),
params=self.aux_head_params,
task=self.task_name,
Expand All @@ -118,8 +118,12 @@ def losses(self) -> list[LossModuleConfig]:
if self.task == "binary"
else "OHEMCrossEntropyLoss"
),
alias="segmentation_loss",
attached_to="segmentation_head",
alias=(
f"OHEMBCEWithLogitsLoss-{self.task_name}"
if self.task == "binary"
else f"OHEMCrossEntropyLoss-{self.task_name}"
),
attached_to=f"DDRNetSegmentationHead-{self.task_name}",
params=self.loss_params,
weight=1.0,
),
Expand All @@ -132,8 +136,12 @@ def losses(self) -> list[LossModuleConfig]:
if self.task == "binary"
else "OHEMCrossEntropyLoss"
),
alias="aux_segmentation_loss",
attached_to="aux_segmentation_head",
alias=(
f"OHEMBCEWithLogitsLoss_aux-{self.task_name}"
if self.task == "binary"
else f"OHEMCrossEntropyLoss_aux-{self.task_name}"
),
attached_to=f"DDRNetSegmentationHead_aux-{self.task_name}",
params=self.loss_params,
weight=0.4,
)
Expand All @@ -146,15 +154,15 @@ def metrics(self) -> list[MetricModuleConfig]:
return [
MetricModuleConfig(
name="JaccardIndex",
alias="segmentation_jaccard_index",
attached_to="segmentation_head",
alias=f"JaccardIndex-{self.task_name}",
attached_to=f"DDRNetSegmentationHead-{self.task_name}",
is_main_metric=True,
params={"task": self.task},
),
MetricModuleConfig(
name="F1Score",
alias="segmentation_f1_score",
attached_to="segmentation_head",
alias=f"F1Score-{self.task_name}",
attached_to=f"DDRNetSegmentationHead-{self.task_name}",
params={"task": self.task},
),
]
Expand All @@ -165,8 +173,8 @@ def visualizers(self) -> list[AttachedModuleConfig]:
return [
AttachedModuleConfig(
name="SegmentationVisualizer",
alias="segmentation_visualizer",
attached_to="segmentation_head",
alias=f"SegmentationVisualizer-{self.task_name}",
attached_to=f"DDRNetSegmentationHead-{self.task_name}",
params=self.visualizer_params,
)
]

0 comments on commit 163fbae

Please sign in to comment.