Merge branch 'master' into bugfix/SG-000_plot

Deci-AI · May 7, 2024 · 8a5685a · 8a5685a
2 parents 2758841 + 292e38c
commit 8a5685a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/src/super_gradients/training/sg_trainer/sg_trainer.py b/src/super_gradients/training/sg_trainer/sg_trainer.py
@@ -1421,6 +1421,14 @@ def get_finetune_lr_dict(self, lr: float) -> Dict[str, float]:
 
         self.ckpt_best_name = self.training_params.ckpt_best_name
 
+        if self.training_params.average_best_models and not self.training_params.save_model:
+            logger.warning(
+                "'training_params.average_best_models'  is enabled, but 'training_params.save_model' is disabled. \n"
+                "Model averaging requires saving snapshot checkpoints to function properly. As a result, "
+                "'training_params.average_best_models' will be disabled. "
+            )
+            self.training_params.average_best_models = False
+
         self.max_train_batches = self.training_params.max_train_batches
         self.max_valid_batches = self.training_params.max_valid_batches
 

diff --git a/src/super_gradients/training/utils/distributed_training_utils.py b/src/super_gradients/training/utils/distributed_training_utils.py
@@ -345,7 +345,6 @@ def restart_script_with_ddp(num_gpus: int = None):
         max_restarts=0,
         monitor_interval=5,
         start_method="spawn",
-        log_dir=None,
         redirects=Std.NONE,
         tee=Std.NONE,
         metrics_cfg={},

diff --git a/src/super_gradients/training/utils/predict/prediction_pose_estimation_results.py b/src/super_gradients/training/utils/predict/prediction_pose_estimation_results.py
@@ -119,7 +119,14 @@ def save(
         :param show_confidence: Whether to show confidence scores on the image.
         :param box_thickness:   (Optional) Thickness of bounding boxes. If None, will adapt to the box size.
         """
-        image = self.draw(box_thickness=box_thickness, show_confidence=show_confidence)
+        image = self.draw(
+            edge_colors=edge_colors,
+            joint_thickness=joint_thickness,
+            keypoint_colors=keypoint_colors,
+            keypoint_radius=keypoint_radius,
+            box_thickness=box_thickness,
+            show_confidence=show_confidence,
+        )
         save_image(image=image, path=output_path)