supervisely 6.73.250__py3-none-any.whl → 6.73.252__py3-none-any.whl

supervisely/nn/training/gui/training_process.py

@@ -22,34 +22,36 @@ class TrainingProcess:
     lock_message = "Select hyperparametrs to unlock"
 
     def __init__(self, app_options: Dict[str, Any]):
-        self.app_options = app_options
-        self.experiment_name_input = Input("Enter experiment name")
-        self.experiment_name_field = Field(
-            title="Experiment name",
-            description="Experiment name will be saved to experiment_info.json",
-            content=self.experiment_name_input,
-        )
-
+        self.display_widgets = []
         self.success_message_text = (
             "Training completed. Training artifacts were uploaded to Team Files. "
             "You can find and open tensorboard logs in the artifacts folder via the "
             "<a href='https://ecosystem.supervisely.com/apps/tensorboard-logs-viewer' target='_blank'>Tensorboard</a> app."
         )
-        self.success_message = DoneLabel(text=self.success_message_text)
-        self.success_message.hide()
+        self.app_options = app_options
 
-        self.artifacts_thumbnail = FolderThumbnail()
-        self.artifacts_thumbnail.hide()
+        # GUI Components
+        # Optional Select CUDA device
+        if self.app_options.get("device_selector", False):
+            self.select_device = SelectCudaDevice()
+            self.select_cuda_device_field = Field(
+                title="Select CUDA device",
+                description="The device on which the model will be trained",
+                content=self.select_device,
+            )
+            self.display_widgets.extend([self.select_cuda_device_field])
+        # -------------------------------- #
 
-        self.model_benchmark_report_thumbnail = ReportThumbnail()
-        self.model_benchmark_report_thumbnail.hide()
+        self.experiment_name_input = Input("Enter experiment name")
+        self.experiment_name_field = Field(
+            title="Experiment name",
+            description="Experiment name will be saved to experiment_info.json",
+            content=self.experiment_name_input,
+        )
 
-        self.validator_text = Text("")
-        self.validator_text.hide()
         self.start_button = Button("Start")
         self.stop_button = Button("Stop", button_type="danger")
         self.stop_button.hide()  # @TODO: implement stop and hide stop button until training starts
-
         button_container = Container(
             [self.start_button, self.stop_button, Empty()],
             "horizontal",
@@ -58,29 +60,34 @@ class TrainingProcess:
             gap=1,
         )
 
-        container_widgets = [
-            self.experiment_name_field,
-            button_container,
-            self.validator_text,
-            self.artifacts_thumbnail,
-            self.model_benchmark_report_thumbnail,
-        ]
+        self.validator_text = Text("")
+        self.validator_text.hide()
 
-        if self.app_options.get("device_selector", False):
-            self.select_device = SelectCudaDevice()
-            self.select_cuda_device_field = Field(
-                title="Select CUDA device",
-                description="The device on which the model will be trained",
-                content=self.select_device,
-            )
-            container_widgets.insert(1, self.select_cuda_device_field)
+        self.artifacts_thumbnail = FolderThumbnail()
+        self.artifacts_thumbnail.hide()
 
-        container = Container(container_widgets)
+        self.display_widgets.extend(
+            [
+                self.experiment_name_field,
+                button_container,
+                self.validator_text,
+                self.artifacts_thumbnail,
+            ]
+        )
+        # -------------------------------- #
 
+        # Optional Model Benchmark
+        if app_options.get("model_benchmark", False):
+            self.model_benchmark_report_thumbnail = ReportThumbnail()
+            self.model_benchmark_report_thumbnail.hide()
+            self.display_widgets.extend([self.model_benchmark_report_thumbnail])
+        # -------------------------------- #
+
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
         )
         self.card.lock()
@@ -92,8 +99,6 @@ class TrainingProcess:
             widgets.append(self.experiment_name_input)
         return widgets
 
-        return []
-
     def validate_step(self) -> bool:
         return True
 

supervisely/nn/training/loggers/__init__.py

@@ -0,0 +1,22 @@
+from supervisely.nn.training.loggers.base_train_logger import BaseTrainLogger
+from supervisely.nn.training.loggers.tensorboard_logger import (
+    TensorboardLogger,
+    tensorboard_installed,
+)
+
+train_logger = BaseTrainLogger()
+
+
+def setup_train_logger(name="tensorboard_logger"):
+    global train_logger
+    if name == "tensorboard_logger":
+        if tensorboard_installed:
+            if not isinstance(train_logger, TensorboardLogger):
+                train_logger = TensorboardLogger()
+        else:
+            raise ImportError("TensorboardX is not installed")
+    elif name == "default_logger":
+        if not isinstance(train_logger, BaseTrainLogger):
+            train_logger = BaseTrainLogger()
+    else:
+        raise ValueError(f"Logger {name} is not supported")

supervisely/nn/training/loggers/base_train_logger.py

@@ -7,7 +7,7 @@ class BaseTrainLogger:
         self._on_train_finished_callbacks = []
         self._on_epoch_started_callbacks = []
         self._on_epoch_finished_callbacks = []
-        self._on_step_callbacks = []
+        self._on_step_finished_callbacks = []
         self.epoch_idx = 0
         self.step_idx = 0
 
@@ -28,9 +28,9 @@ class BaseTrainLogger:
         for callback in self._on_epoch_finished_callbacks:
             callback()
 
-    def on_step_end(self):
+    def step_finished(self):
         self.step_idx += 1
-        for callback in self._on_step_callbacks:
+        for callback in self._on_step_finished_callbacks:
             callback()
 
     def add_on_train_started_callback(self, callback: Callable):
@@ -45,8 +45,8 @@ class BaseTrainLogger:
     def add_on_epoch_finish_callback(self, callback: Callable):
         self._on_epoch_finished_callbacks.append(callback)
 
-    def add_on_step_callback(self, callback: Callable):
-        self._on_step_callbacks.append(callback)
+    def add_on_step_finished_callback(self, callback: Callable):
+        self._on_step_finished_callbacks.append(callback)
 
     def log(self, logs: dict, idx: int):
         raise NotImplementedError
@@ -56,3 +56,6 @@ class BaseTrainLogger:
 
     def log_epoch(self, logs: dict):
         raise NotImplementedError
+    
+    def close(self):
+        pass

supervisely/nn/training/loggers/tensorboard_logger.py

@@ -1,13 +1,10 @@
-from supervisely import logger
-
+from supervisely.nn.training.loggers.base_train_logger import BaseTrainLogger
+tensorboard_installed = False
 try:
     from tensorboardX import SummaryWriter
-    _tensorboard_supported = True
+    tensorboard_installed = True
 except ImportError:
-    logger.warning("TensorboardX is not installed. TensorboardLogger will not work.")
-    _tensorboard_supported = False
-
-from supervisely.nn.training.loggers.base_train_logger import BaseTrainLogger
+    pass
 
 
 class TensorboardLogger(BaseTrainLogger):
@@ -40,7 +37,3 @@ class TensorboardLogger(BaseTrainLogger):
 
     def log_epoch(self, logs: dict):
         self.log(logs, self.epoch_idx)
-
-
-if _tensorboard_supported:
-    tb_logger = TensorboardLogger()