[onert/python] Add typing annotations across training API modules (#15183)

* [onert/python] Add typing annotations across training API modules

This commit adds typing annotations across training API modules.
- Applied `from typing` imports and added type hints to DataLoader methods and attributes
- Annotated LossFunction, MeanSquaredError, CategoricalCrossentropy constructors with `Literal`, `Dict` and return types
- Typed LossRegistry and MetricsRegistry: `create_*` and `map_*` methods use `str`, `Type[...]`, and enum return types
- Annotated Optimizer base class, Adam, SGD and OptimizerRegistry with `float`, `Optional`, `Dict`, and `Type[...]`
- Added precise types to TrainSession: `train_info: traininfo`, `optimizer: Optional[Optimizer]`, `loss: Optional[LossFunction]`, `metrics: List[Metric]`, `total_time: Dict[str, Union[float,List[float]]]`, and method signatures

ONE-DCO-1.0-Signed-off-by: ragmani <ragmani0216@gmail.com>

* Remove unnecessary commented code

* Restore removed comments

Author: ragmani

runtime/onert/api/python/package/experimental/train/losses/cce.py

@@ -1,3 +1,4 @@
+from typing import Literal
 import numpy as np
 from .loss import LossFunction
 
@@ -6,7 +7,7 @@ class CategoricalCrossentropy(LossFunction):
     """
     Categorical Cross-Entropy Loss Function with reduction type.
     """
-    def __init__(self, reduction="mean"):
+    def __init__(self, reduction: Literal["mean", "sum"] = "mean") -> None:
         """
         Initialize the Categorical Cross-Entropy loss function.
         Args:

runtime/onert/api/python/package/experimental/train/losses/loss.py

@@ -1,17 +1,18 @@
+from typing import Literal, Dict
 from onert.native.libnnfw_api_pybind import loss_reduction
 
 
 class LossFunction:
     """
     Base class for loss functions with reduction type.
     """
-    def __init__(self, reduction="mean"):
+    def __init__(self, reduction: Literal["mean", "sum"] = "mean") -> None:
         """
         Initialize the Categorical Cross-Entropy loss function.
         Args:
             reduction (str): Reduction type ('mean', 'sum').
         """
-        reduction_mapping = {
+        reduction_mapping: Dict[Literal["mean", "sum"], loss_reduction] = {
             "mean": loss_reduction.SUM_OVER_BATCH_SIZE,
             "sum": loss_reduction.SUM
         }
@@ -21,4 +22,4 @@ def __init__(self, reduction="mean"):
             raise ValueError(
                 f"Invalid reduction type. Choose from {list(reduction_mapping.keys())}.")
 
-        self.reduction = reduction_mapping[reduction]
+        self.reduction: loss_reduction = reduction_mapping[reduction]

runtime/onert/api/python/package/experimental/train/losses/mse.py

@@ -1,3 +1,4 @@
+from typing import Literal
 import numpy as np
 from .loss import LossFunction
 
@@ -6,7 +7,7 @@ class MeanSquaredError(LossFunction):
     """
     Mean Squared Error (MSE) Loss Function with reduction type.
     """
-    def __init__(self, reduction="mean"):
+    def __init__(self, reduction: Literal["mean", "sum"] = "mean") -> None:
         """
         Initialize the MSE loss function.
         Args:

runtime/onert/api/python/package/experimental/train/losses/registry.py

@@ -1,19 +1,21 @@
-from onert.native.libnnfw_api_pybind import loss as loss_type
+from typing import Type, Dict
+from .loss import LossFunction
 from .cce import CategoricalCrossentropy
 from .mse import MeanSquaredError
+from onert.native.libnnfw_api_pybind import loss as loss_type
 
 
 class LossRegistry:
     """
     Registry for creating and mapping losses by name or instance.
     """
-    _losses = {
+    _losses: Dict[str, Type[LossFunction]] = {
         "categorical_crossentropy": CategoricalCrossentropy,
         "mean_squared_error": MeanSquaredError
     }
 
     @staticmethod
-    def create_loss(name):
+    def create_loss(name: str) -> LossFunction:
         """
         Create a loss instance by name.
         Args:
@@ -26,7 +28,7 @@ def create_loss(name):
         return LossRegistry._losses[name]()
 
     @staticmethod
-    def map_loss_function_to_enum(loss_instance):
+    def map_loss_function_to_enum(loss_instance: LossFunction) -> loss_type:
         """
         Maps a LossFunction instance to the appropriate enum value.
         Args:
@@ -36,8 +38,7 @@ def map_loss_function_to_enum(loss_instance):
         Raises:
             TypeError: If the loss_instance is not a recognized LossFunction type.
         """
-        # Loss to Enum mapping
-        loss_to_enum = {
+        loss_to_enum: Dict[Type[LossFunction], loss_type] = {
             CategoricalCrossentropy: loss_type.CATEGORICAL_CROSSENTROPY,
             MeanSquaredError: loss_type.MEAN_SQUARED_ERROR
         }

runtime/onert/api/python/package/experimental/train/metrics/categorical_accuracy.py

@@ -1,26 +1,32 @@
 import numpy as np
+from typing import List
 from .metric import Metric
 
 
 class CategoricalAccuracy(Metric):
     """
     Metric for computing categorical accuracy.
     """
-    def __init__(self):
-        self.correct = 0
-        self.total = 0
-        self.axis = 0
+    def __init__(self) -> None:
+        """
+        Initialize internal counters and axis.
+        """
+        self.correct: int = 0
+        self.total: int = 0
+        self.axis: int = 0
 
-    def reset_state(self):
+    def reset_state(self) -> None:
         """
         Reset the metric's state.
         """
         self.correct = 0
         self.total = 0
 
-    def update_state(self, outputs, expecteds):
+    def update_state(self, outputs: List[np.ndarray],
+                     expecteds: List[np.ndarray]) -> None:
         """
         Update the metric's state based on the outputs and expecteds.
+
         Args:
             outputs (list of np.ndarray): List of model outputs for each output layer.
             expecteds (list of np.ndarray): List of expected ground truth values for each output layer.
@@ -45,9 +51,10 @@ def update_state(self, outputs, expecteds):
                     self.correct += 1
             self.total += batch_size
 
-    def result(self):
+    def result(self) -> float:
         """
         Compute and return the final metric value.
+
         Returns:
             float: Metric value.
         """

runtime/onert/api/python/package/experimental/train/metrics/metric.py

@@ -1,25 +1,30 @@
+from typing import Any
+
+
 class Metric:
     """
     Abstract base class for all metrics.
     """
-    def reset_state(self):
+    def reset_state(self) -> None:
         """
         Reset the metric's state.
         """
         raise NotImplementedError
 
-    def update_state(self, outputs, expecteds):
+    def update_state(self, outputs: Any, expecteds: Any) -> None:
         """
         Update the metric's state based on the outputs and expecteds.
+
         Args:
-            outputs (np.ndarray): Model outputs.
-            expecteds (np.ndarray): Expected ground truth values.
+            outputs (Any): Model outputs.
+            expecteds (Any): Expected ground truth values.
         """
         raise NotImplementedError
 
-    def result(self):
+    def result(self) -> float:
         """
         Compute and return the final metric value.
+
         Returns:
             float: Metric value.
         """

runtime/onert/api/python/package/experimental/train/metrics/registry.py

@@ -1,22 +1,26 @@
+from typing import Type, Dict
+from .metric import Metric
 from .categorical_accuracy import CategoricalAccuracy
 
 
 class MetricsRegistry:
     """
     Registry for creating metrics by name.
     """
-    _metrics = {
+    _metrics: Dict[str, Type[Metric]] = {
         "categorical_accuracy": CategoricalAccuracy,
     }
 
     @staticmethod
-    def create_metric(name):
+    def create_metric(name: str) -> Metric:
         """
         Create a metric instance by name.
+
         Args:
             name (str): Name of the metric.
+
         Returns:
-            BaseMetric: Metric instance.
+            Metric: Metric instance.
         """
         if name not in MetricsRegistry._metrics:
             raise ValueError(

runtime/onert/api/python/package/experimental/train/optimizer/adam.py

@@ -1,20 +1,26 @@
+from typing import Literal
 from .optimizer import Optimizer
 
 
 class Adam(Optimizer):
     """
     Adam optimizer.
     """
-    def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-7):
+    def __init__(self,
+                 learning_rate: float = 0.001,
+                 beta1: float = 0.9,
+                 beta2: float = 0.999,
+                 epsilon: float = 1e-7) -> None:
         """
         Initialize the Adam optimizer.
+
         Args:
             learning_rate (float): The learning rate for optimization.
             beta1 (float): Exponential decay rate for the first moment estimates.
             beta2 (float): Exponential decay rate for the second moment estimates.
             epsilon (float): Small constant to prevent division by zero.
         """
         super().__init__(learning_rate)
-        self.beta1 = beta1
-        self.beta2 = beta2
-        self.epsilon = epsilon
+        self.beta1: float = beta1
+        self.beta2: float = beta2
+        self.epsilon: float = epsilon

runtime/onert/api/python/package/experimental/train/optimizer/optimizer.py

@@ -5,11 +5,15 @@ class Optimizer:
     """
     Base class for optimizers.
     """
-    def __init__(self, learning_rate=0.001, nums_trainable_ops=trainable_ops.ALL):
+    def __init__(self,
+                 learning_rate: float = 0.001,
+                 nums_trainable_ops: int = trainable_ops.ALL) -> None:
         """
         Initialize the optimizer.
+
         Args:
             learning_rate (float): The learning rate for optimization.
+            nums_trainable_ops (int or enum): Number of trainable ops or enum mask.
         """
-        self.learning_rate = learning_rate
-        self.nums_trainable_ops = nums_trainable_ops
+        self.learning_rate: float = learning_rate
+        self.nums_trainable_ops: int = nums_trainable_ops

runtime/onert/api/python/package/experimental/train/optimizer/registry.py

@@ -1,44 +1,53 @@
-from onert.native.libnnfw_api_pybind import optimizer as optimizer_type
+from typing import Type, Dict
+from .optimizer import Optimizer
 from .adam import Adam
 from .sgd import SGD
+from onert.native.libnnfw_api_pybind import optimizer as optimizer_type
 
 
 class OptimizerRegistry:
     """
     Registry for creating optimizers by name.
     """
-    _optimizers = {"adam": Adam, "sgd": SGD}
+    _optimizers: Dict[str, Type[Optimizer]] = {"adam": Adam, "sgd": SGD}
 
     @staticmethod
-    def create_optimizer(name):
+    def create_optimizer(name: str) -> Optimizer:
         """
         Create an optimizer instance by name.
+
         Args:
             name (str): Name of the optimizer.
+
         Returns:
-            BaseOptimizer: Optimizer instance.
+            Optimizer: Optimizer instance.
         """
         if name not in OptimizerRegistry._optimizers:
             raise ValueError(
                 f"Unknown Optimizer: {name}. Custom optimizer is not supported yet")
         return OptimizerRegistry._optimizers[name]()
 
     @staticmethod
-    def map_optimizer_to_enum(optimizer_instance):
+    def map_optimizer_to_enum(optimizer_instance: Optimizer) -> optimizer_type:
         """
         Maps an optimizer instance to the appropriate enum value.
+
         Args:
             optimizer_instance (Optimizer): An instance of an optimizer.
+
         Returns:
             optimizer_type: Corresponding enum value for the optimizer.
+
         Raises:
             TypeError: If the optimizer_instance is not a recognized optimizer type.
         """
-        # Optimizer to Enum mapping
-        optimizer_to_enum = {SGD: optimizer_type.SGD, Adam: optimizer_type.ADAM}
-        for optimizer_class, enum_value in optimizer_to_enum.items():
-            if isinstance(optimizer_instance, optimizer_class):
-                return enum_value
+        optimizer_to_enum: Dict[Type[Optimizer], optimizer_type] = {
+            SGD: optimizer_type.SGD,
+            Adam: optimizer_type.ADAM
+        }
+        for cls, enum_val in optimizer_to_enum.items():
+            if isinstance(optimizer_instance, cls):
+                return enum_val
         raise TypeError(
             f"Unsupported optimizer type: {type(optimizer_instance).__name__}. "
             f"Supported types are: {list(optimizer_to_enum.keys())}.")