gnn_tracking.training.tc
========================

.. py:module:: gnn_tracking.training.tc

.. autoapi-nested-parse::

   Lightning module for object condensation training.



Classes
-------

.. autoapisummary::

   gnn_tracking.training.tc.TCModule


Module Contents
---------------

.. py:class:: TCModule(*, loss_fct: gnn_tracking.metrics.losses.MultiLossFct, cluster_scanner: gnn_tracking.postprocessing.clusterscanner.ClusterScanner | None = None, **kwargs)

   Bases: :py:obj:`gnn_tracking.training.base.TrackingModule`


   Object condensation for tracks. This lightning module implements
   losses, training, and validation steps. k:w


   :param loss_fct:
   :param cluster_scanner:
   :param \*\*kwargs: Passed on to `TrackingModule`


   .. py:attribute:: loss_fct


   .. py:attribute:: cluster_scanner


   .. py:attribute:: _cluster_scan_input


   .. py:attribute:: _best_cluster_params


   .. py:method:: is_last_val_batch(batch_idx: int) -> bool

      Are we validating the last batch of the validation set?



   .. py:method:: get_losses(out: dict[str, Any], data: torch_geometric.data.Data) -> tuple[torch.Tensor, dict[str, float]]


   .. py:method:: training_step(data: torch_geometric.data.Data, batch_idx: int) -> torch.Tensor


   .. py:method:: validation_step(data: torch_geometric.data.Data, batch_idx: int) -> None


   .. py:method:: _evaluate_cluster_metrics(out: dict[str, Any], data: torch_geometric.data.Data, batch_idx: int) -> dict[str, float]

      Evaluate cluster metrics.



   .. py:method:: highlight_metric(metric: str) -> bool