dgenerate-ultralytics-headless 8.3.222__py3-none-any.whl → 8.3.225__py3-none-any.whl

ultralytics/models/utils/loss.py

@@ -15,11 +15,10 @@ from .ops import HungarianMatcher
 
 
 class DETRLoss(nn.Module):
-    """
-    DETR (DEtection TRansformer) Loss class for calculating various loss components.
+    """DETR (DEtection TRansformer) Loss class for calculating various loss components.
 
-    This class computes classification loss, bounding box loss, GIoU loss, and optionally auxiliary losses for the
-    DETR object detection model.
+    This class computes classification loss, bounding box loss, GIoU loss, and optionally auxiliary losses for the DETR
+    object detection model.
 
     Attributes:
         nc (int): Number of classes.
@@ -47,8 +46,7 @@ class DETRLoss(nn.Module):
         gamma: float = 1.5,
         alpha: float = 0.25,
     ):
-        """
-        Initialize DETR loss function with customizable components and gains.
+        """Initialize DETR loss function with customizable components and gains.
 
         Uses default loss_gain if not provided. Initializes HungarianMatcher with preset cost gains. Supports auxiliary
         losses and various loss types.
@@ -82,8 +80,7 @@ class DETRLoss(nn.Module):
     def _get_loss_class(
         self, pred_scores: torch.Tensor, targets: torch.Tensor, gt_scores: torch.Tensor, num_gts: int, postfix: str = ""
     ) -> dict[str, torch.Tensor]:
-        """
-        Compute classification loss based on predictions, target values, and ground truth scores.
+        """Compute classification loss based on predictions, target values, and ground truth scores.
 
         Args:
             pred_scores (torch.Tensor): Predicted class scores with shape (B, N, C).
@@ -124,8 +121,7 @@ class DETRLoss(nn.Module):
     def _get_loss_bbox(
         self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, postfix: str = ""
     ) -> dict[str, torch.Tensor]:
-        """
-        Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
+        """Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes with shape (N, 4).
@@ -199,8 +195,7 @@ class DETRLoss(nn.Module):
         masks: torch.Tensor | None = None,
         gt_mask: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
-        """
-        Get auxiliary losses for intermediate decoder layers.
+        """Get auxiliary losses for intermediate decoder layers.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes from auxiliary layers.
@@ -261,8 +256,7 @@ class DETRLoss(nn.Module):
 
     @staticmethod
     def _get_index(match_indices: list[tuple]) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
-        """
-        Extract batch indices, source indices, and destination indices from match indices.
+        """Extract batch indices, source indices, and destination indices from match indices.
 
         Args:
             match_indices (list[tuple]): List of tuples containing matched indices.
@@ -279,8 +273,7 @@ class DETRLoss(nn.Module):
     def _get_assigned_bboxes(
         self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, match_indices: list[tuple]
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Assign predicted bounding boxes to ground truth bounding boxes based on match indices.
+        """Assign predicted bounding boxes to ground truth bounding boxes based on match indices.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes.
@@ -317,8 +310,7 @@ class DETRLoss(nn.Module):
         postfix: str = "",
         match_indices: list[tuple] | None = None,
     ) -> dict[str, torch.Tensor]:
-        """
-        Calculate losses for a single prediction layer.
+        """Calculate losses for a single prediction layer.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes.
@@ -364,8 +356,7 @@ class DETRLoss(nn.Module):
         postfix: str = "",
         **kwargs: Any,
     ) -> dict[str, torch.Tensor]:
-        """
-        Calculate loss for predicted bounding boxes and scores.
+        """Calculate loss for predicted bounding boxes and scores.
 
         Args:
             pred_bboxes (torch.Tensor): Predicted bounding boxes, shape (L, B, N, 4).
@@ -400,8 +391,7 @@ class DETRLoss(nn.Module):
 
 
 class RTDETRDetectionLoss(DETRLoss):
-    """
-    Real-Time DeepTracker (RT-DETR) Detection Loss class that extends the DETRLoss.
+    """Real-Time DeepTracker (RT-DETR) Detection Loss class that extends the DETRLoss.
 
     This class computes the detection loss for the RT-DETR model, which includes the standard detection loss as well as
     an additional denoising training loss when provided with denoising metadata.
@@ -415,8 +405,7 @@ class RTDETRDetectionLoss(DETRLoss):
         dn_scores: torch.Tensor | None = None,
         dn_meta: dict[str, Any] | None = None,
     ) -> dict[str, torch.Tensor]:
-        """
-        Forward pass to compute detection loss with optional denoising loss.
+        """Forward pass to compute detection loss with optional denoising loss.
 
         Args:
             preds (tuple[torch.Tensor, torch.Tensor]): Tuple containing predicted bounding boxes and scores.
@@ -452,8 +441,7 @@ class RTDETRDetectionLoss(DETRLoss):
     def get_dn_match_indices(
         dn_pos_idx: list[torch.Tensor], dn_num_group: int, gt_groups: list[int]
     ) -> list[tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Get match indices for denoising.
+        """Get match indices for denoising.
 
         Args:
             dn_pos_idx (list[torch.Tensor]): List of tensors containing positive indices for denoising.

ultralytics/models/utils/ops.py

@@ -14,8 +14,7 @@ from ultralytics.utils.ops import xywh2xyxy, xyxy2xywh
 
 
 class HungarianMatcher(nn.Module):
-    """
-    A module implementing the HungarianMatcher for optimal assignment between predictions and ground truth.
+    """A module implementing the HungarianMatcher for optimal assignment between predictions and ground truth.
 
     HungarianMatcher performs optimal bipartite assignment over predicted and ground truth bounding boxes using a cost
     function that considers classification scores, bounding box coordinates, and optionally mask predictions. This is
@@ -56,8 +55,7 @@ class HungarianMatcher(nn.Module):
         alpha: float = 0.25,
         gamma: float = 2.0,
     ):
-        """
-        Initialize HungarianMatcher for optimal assignment of predicted and ground truth bounding boxes.
+        """Initialize HungarianMatcher for optimal assignment of predicted and ground truth bounding boxes.
 
         Args:
             cost_gain (dict[str, float], optional): Dictionary of cost coefficients for different matching cost
@@ -88,8 +86,7 @@ class HungarianMatcher(nn.Module):
         masks: torch.Tensor | None = None,
         gt_mask: list[torch.Tensor] | None = None,
     ) -> list[tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Compute optimal assignment between predictions and ground truth using Hungarian algorithm.
+        """Compute optimal assignment between predictions and ground truth using Hungarian algorithm.
 
         This method calculates matching costs based on classification scores, bounding box coordinates, and optionally
         mask predictions, then finds the optimal bipartite assignment between predictions and ground truth.
@@ -105,10 +102,10 @@ class HungarianMatcher(nn.Module):
             gt_mask (list[torch.Tensor], optional): Ground truth masks, each with shape (num_masks, Height, Width).
 
         Returns:
-            (list[tuple[torch.Tensor, torch.Tensor]]): A list of size batch_size, each element is a tuple
-                (index_i, index_j), where index_i is the tensor of indices of the selected predictions (in order)
-                and index_j is the tensor of indices of the corresponding selected ground truth targets (in order).
-                For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes).
+            (list[tuple[torch.Tensor, torch.Tensor]]): A list of size batch_size, each element is a tuple (index_i,
+                index_j), where index_i is the tensor of indices of the selected predictions (in order) and index_j is
+                the tensor of indices of the corresponding selected ground truth targets (in order).
+            For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes).
         """
         bs, nq, nc = pred_scores.shape
 
@@ -198,16 +195,15 @@ def get_cdn_group(
     box_noise_scale: float = 1.0,
     training: bool = False,
 ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, dict[str, Any] | None]:
-    """
-    Generate contrastive denoising training group with positive and negative samples from ground truths.
+    """Generate contrastive denoising training group with positive and negative samples from ground truths.
 
-    This function creates denoising queries for contrastive denoising training by adding noise to ground truth
-    bounding boxes and class labels. It generates both positive and negative samples to improve model robustness.
+    This function creates denoising queries for contrastive denoising training by adding noise to ground truth bounding
+    boxes and class labels. It generates both positive and negative samples to improve model robustness.
 
     Args:
-        batch (dict[str, Any]): Batch dictionary containing 'gt_cls' (torch.Tensor with shape (num_gts,)),
-            'gt_bboxes' (torch.Tensor with shape (num_gts, 4)), and 'gt_groups' (list[int]) indicating number of
-            ground truths per image.
+        batch (dict[str, Any]): Batch dictionary containing 'gt_cls' (torch.Tensor with shape (num_gts,)), 'gt_bboxes'
+            (torch.Tensor with shape (num_gts, 4)), and 'gt_groups' (list[int]) indicating number of ground truths
+            per image.
         num_classes (int): Total number of object classes.
         num_queries (int): Number of object queries.
         class_embed (torch.Tensor): Class embedding weights to map labels to embedding space.

ultralytics/models/yolo/classify/predict.py

@@ -11,11 +11,10 @@ from ultralytics.utils import DEFAULT_CFG, ops
 
 
 class ClassificationPredictor(BasePredictor):
-    """
-    A class extending the BasePredictor class for prediction based on a classification model.
+    """A class extending the BasePredictor class for prediction based on a classification model.
 
-    This predictor handles the specific requirements of classification models, including preprocessing images
-    and postprocessing predictions to generate classification results.
+    This predictor handles the specific requirements of classification models, including preprocessing images and
+    postprocessing predictions to generate classification results.
 
     Attributes:
         args (dict): Configuration arguments for the predictor.
@@ -24,20 +23,19 @@ class ClassificationPredictor(BasePredictor):
         preprocess: Convert input images to model-compatible format.
         postprocess: Process model predictions into Results objects.
 
-    Notes:
-        - Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'.
-
     Examples:
         >>> from ultralytics.utils import ASSETS
         >>> from ultralytics.models.yolo.classify import ClassificationPredictor
         >>> args = dict(model="yolo11n-cls.pt", source=ASSETS)
         >>> predictor = ClassificationPredictor(overrides=args)
         >>> predictor.predict_cli()
+
+    Notes:
+        - Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'.
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize the ClassificationPredictor with the specified configuration and set task to 'classify'.
+        """Initialize the ClassificationPredictor with the specified configuration and set task to 'classify'.
 
         This constructor initializes a ClassificationPredictor instance, which extends BasePredictor for classification
         tasks. It ensures the task is set to 'classify' regardless of input configuration.
@@ -72,8 +70,7 @@ class ClassificationPredictor(BasePredictor):
         return img.half() if self.model.fp16 else img.float()  # Convert uint8 to fp16/32
 
     def postprocess(self, preds, img, orig_imgs):
-        """
-        Process predictions to return Results objects with classification probabilities.
+        """Process predictions to return Results objects with classification probabilities.
 
         Args:
             preds (torch.Tensor): Raw predictions from the model.

ultralytics/models/yolo/classify/train.py

@@ -17,8 +17,7 @@ from ultralytics.utils.torch_utils import is_parallel, torch_distributed_zero_fi
 
 
 class ClassificationTrainer(BaseTrainer):
-    """
-    A trainer class extending BaseTrainer for training image classification models.
+    """A trainer class extending BaseTrainer for training image classification models.
 
     This trainer handles the training process for image classification tasks, supporting both YOLO classification models
     and torchvision models with comprehensive dataset handling and validation.
@@ -51,8 +50,7 @@ class ClassificationTrainer(BaseTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
-        """
-        Initialize a ClassificationTrainer object.
+        """Initialize a ClassificationTrainer object.
 
         Args:
             cfg (dict[str, Any], optional): Default configuration dictionary containing training parameters.
@@ -71,8 +69,7 @@ class ClassificationTrainer(BaseTrainer):
         self.model.names = self.data["names"]
 
     def get_model(self, cfg=None, weights=None, verbose: bool = True):
-        """
-        Return a modified PyTorch model configured for training YOLO classification.
+        """Return a modified PyTorch model configured for training YOLO classification.
 
         Args:
             cfg (Any, optional): Model configuration.
@@ -96,8 +93,7 @@ class ClassificationTrainer(BaseTrainer):
         return model
 
     def setup_model(self):
-        """
-        Load, create or download model for classification tasks.
+        """Load, create or download model for classification tasks.
 
         Returns:
             (Any): Model checkpoint if applicable, otherwise None.
@@ -115,8 +111,7 @@ class ClassificationTrainer(BaseTrainer):
         return ckpt
 
     def build_dataset(self, img_path: str, mode: str = "train", batch=None):
-        """
-        Create a ClassificationDataset instance given an image path and mode.
+        """Create a ClassificationDataset instance given an image path and mode.
 
         Args:
             img_path (str): Path to the dataset images.
@@ -129,8 +124,7 @@ class ClassificationTrainer(BaseTrainer):
         return ClassificationDataset(root=img_path, args=self.args, augment=mode == "train", prefix=mode)
 
     def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
-        """
-        Return PyTorch DataLoader with transforms to preprocess images.
+        """Return PyTorch DataLoader with transforms to preprocess images.
 
         Args:
             dataset_path (str): Path to the dataset.
@@ -177,8 +171,7 @@ class ClassificationTrainer(BaseTrainer):
         )
 
     def label_loss_items(self, loss_items: torch.Tensor | None = None, prefix: str = "train"):
-        """
-        Return a loss dict with labeled training loss items tensor.
+        """Return a loss dict with labeled training loss items tensor.
 
         Args:
             loss_items (torch.Tensor, optional): Loss tensor items.
@@ -195,8 +188,7 @@ class ClassificationTrainer(BaseTrainer):
         return dict(zip(keys, loss_items))
 
     def plot_training_samples(self, batch: dict[str, torch.Tensor], ni: int):
-        """
-        Plot training samples with their annotations.
+        """Plot training samples with their annotations.
 
         Args:
             batch (dict[str, torch.Tensor]): Batch containing images and class labels.

ultralytics/models/yolo/classify/val.py

@@ -16,11 +16,10 @@ from ultralytics.utils.plotting import plot_images
 
 
 class ClassificationValidator(BaseValidator):
-    """
-    A class extending the BaseValidator class for validation based on a classification model.
+    """A class extending the BaseValidator class for validation based on a classification model.
 
-    This validator handles the validation process for classification models, including metrics calculation,
-    confusion matrix generation, and visualization of results.
+    This validator handles the validation process for classification models, including metrics calculation, confusion
+    matrix generation, and visualization of results.
 
     Attributes:
         targets (list[torch.Tensor]): Ground truth class labels.
@@ -55,8 +54,7 @@ class ClassificationValidator(BaseValidator):
     """
 
     def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks=None) -> None:
-        """
-        Initialize ClassificationValidator with dataloader, save directory, and other parameters.
+        """Initialize ClassificationValidator with dataloader, save directory, and other parameters.
 
         Args:
             dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
@@ -96,8 +94,7 @@ class ClassificationValidator(BaseValidator):
         return batch
 
     def update_metrics(self, preds: torch.Tensor, batch: dict[str, Any]) -> None:
-        """
-        Update running metrics with model predictions and batch targets.
+        """Update running metrics with model predictions and batch targets.
 
         Args:
             preds (torch.Tensor): Model predictions, typically logits or probabilities for each class.
@@ -112,12 +109,7 @@ class ClassificationValidator(BaseValidator):
         self.targets.append(batch["cls"].type(torch.int32).cpu())
 
     def finalize_metrics(self) -> None:
-        """
-        Finalize metrics including confusion matrix and processing speed.
-
-        Notes:
-            This method processes the accumulated predictions and targets to generate the confusion matrix,
-            optionally plots it, and updates the metrics object with speed information.
+        """Finalize metrics including confusion matrix and processing speed.
 
         Examples:
             >>> validator = ClassificationValidator()
@@ -125,6 +117,10 @@ class ClassificationValidator(BaseValidator):
             >>> validator.targets = [torch.tensor([0])]  # Ground truth class
             >>> validator.finalize_metrics()
             >>> print(validator.metrics.confusion_matrix)  # Access the confusion matrix
+
+        Notes:
+            This method processes the accumulated predictions and targets to generate the confusion matrix,
+            optionally plots it, and updates the metrics object with speed information.
         """
         self.confusion_matrix.process_cls_preds(self.pred, self.targets)
         if self.args.plots:
@@ -161,8 +157,7 @@ class ClassificationValidator(BaseValidator):
         return ClassificationDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
 
     def get_dataloader(self, dataset_path: Path | str, batch_size: int) -> torch.utils.data.DataLoader:
-        """
-        Build and return a data loader for classification validation.
+        """Build and return a data loader for classification validation.
 
         Args:
             dataset_path (str | Path): Path to the dataset directory.
@@ -180,8 +175,7 @@ class ClassificationValidator(BaseValidator):
         LOGGER.info(pf % ("all", self.metrics.top1, self.metrics.top5))
 
     def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
-        """
-        Plot validation image samples with their ground truth labels.
+        """Plot validation image samples with their ground truth labels.
 
         Args:
             batch (dict[str, Any]): Dictionary containing batch data with 'img' (images) and 'cls' (class labels).
@@ -201,8 +195,7 @@ class ClassificationValidator(BaseValidator):
         )
 
     def plot_predictions(self, batch: dict[str, Any], preds: torch.Tensor, ni: int) -> None:
-        """
-        Plot images with their predicted class labels and save the visualization.
+        """Plot images with their predicted class labels and save the visualization.
 
         Args:
             batch (dict[str, Any]): Batch data containing images and other information.

ultralytics/models/yolo/detect/predict.py

@@ -6,8 +6,7 @@ from ultralytics.utils import nms, ops
 
 
 class DetectionPredictor(BasePredictor):
-    """
-    A class extending the BasePredictor class for prediction based on a detection model.
+    """A class extending the BasePredictor class for prediction based on a detection model.
 
     This predictor specializes in object detection tasks, processing model outputs into meaningful detection results
     with bounding boxes and class predictions.
@@ -32,8 +31,7 @@ class DetectionPredictor(BasePredictor):
     """
 
     def postprocess(self, preds, img, orig_imgs, **kwargs):
-        """
-        Post-process predictions and return a list of Results objects.
+        """Post-process predictions and return a list of Results objects.
 
         This method applies non-maximum suppression to raw model predictions and prepares them for visualization and
         further analysis.
@@ -92,8 +90,7 @@ class DetectionPredictor(BasePredictor):
         return [feats[idx] if idx.shape[0] else [] for feats, idx in zip(obj_feats, idxs)]  # for each img in batch
 
     def construct_results(self, preds, img, orig_imgs):
-        """
-        Construct a list of Results objects from model predictions.
+        """Construct a list of Results objects from model predictions.
 
         Args:
             preds (list[torch.Tensor]): List of predicted bounding boxes and scores for each image.
@@ -109,8 +106,7 @@ class DetectionPredictor(BasePredictor):
         ]
 
     def construct_result(self, pred, img, orig_img, img_path):
-        """
-        Construct a single Results object from one image prediction.
+        """Construct a single Results object from one image prediction.
 
         Args:
             pred (torch.Tensor): Predicted boxes and scores with shape (N, 6) where N is the number of detections.

ultralytics/models/yolo/detect/train.py

@@ -22,11 +22,10 @@ from ultralytics.utils.torch_utils import torch_distributed_zero_first, unwrap_m
 
 
 class DetectionTrainer(BaseTrainer):
-    """
-    A class extending the BaseTrainer class for training based on a detection model.
+    """A class extending the BaseTrainer class for training based on a detection model.
 
-    This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models
-    for object detection including dataset building, data loading, preprocessing, and model configuration.
+    This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models for
+    object detection including dataset building, data loading, preprocessing, and model configuration.
 
     Attributes:
         model (DetectionModel): The YOLO detection model being trained.
@@ -54,8 +53,7 @@ class DetectionTrainer(BaseTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
-        """
-        Initialize a DetectionTrainer object for training YOLO object detection model training.
+        """Initialize a DetectionTrainer object for training YOLO object detection model training.
 
         Args:
             cfg (dict, optional): Default configuration dictionary containing training parameters.
@@ -65,8 +63,7 @@ class DetectionTrainer(BaseTrainer):
         super().__init__(cfg, overrides, _callbacks)
 
     def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
-        """
-        Build YOLO Dataset for training or validation.
+        """Build YOLO Dataset for training or validation.
 
         Args:
             img_path (str): Path to the folder containing images.
@@ -80,8 +77,7 @@ class DetectionTrainer(BaseTrainer):
         return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
 
     def get_dataloader(self, dataset_path: str, batch_size: int = 16, rank: int = 0, mode: str = "train"):
-        """
-        Construct and return dataloader for the specified mode.
+        """Construct and return dataloader for the specified mode.
 
         Args:
             dataset_path (str): Path to the dataset.
@@ -109,8 +105,7 @@ class DetectionTrainer(BaseTrainer):
         )
 
     def preprocess_batch(self, batch: dict) -> dict:
-        """
-        Preprocess a batch of images by scaling and converting to float.
+        """Preprocess a batch of images by scaling and converting to float.
 
         Args:
             batch (dict): Dictionary containing batch data with 'img' tensor.
@@ -150,8 +145,7 @@ class DetectionTrainer(BaseTrainer):
         # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
 
     def get_model(self, cfg: str | None = None, weights: str | None = None, verbose: bool = True):
-        """
-        Return a YOLO detection model.
+        """Return a YOLO detection model.
 
         Args:
             cfg (str, optional): Path to model configuration file.
@@ -174,8 +168,7 @@ class DetectionTrainer(BaseTrainer):
         )
 
     def label_loss_items(self, loss_items: list[float] | None = None, prefix: str = "train"):
-        """
-        Return a loss dict with labeled training loss items tensor.
+        """Return a loss dict with labeled training loss items tensor.
 
         Args:
             loss_items (list[float], optional): List of loss values.
@@ -202,8 +195,7 @@ class DetectionTrainer(BaseTrainer):
         )
 
     def plot_training_samples(self, batch: dict[str, Any], ni: int) -> None:
-        """
-        Plot training samples with their annotations.
+        """Plot training samples with their annotations.
 
         Args:
             batch (dict[str, Any]): Dictionary containing batch data.
@@ -223,8 +215,7 @@ class DetectionTrainer(BaseTrainer):
         plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)
 
     def auto_batch(self):
-        """
-        Get optimal batch size by calculating memory occupation of model.
+        """Get optimal batch size by calculating memory occupation of model.
 
         Returns:
             (int): Optimal batch size.