60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
"""PINA Callbacks Implementations"""
|
|
|
|
from pytorch_lightning.callbacks import Callback
|
|
import torch
|
|
import copy
|
|
|
|
|
|
class MetricTracker(Callback):
|
|
|
|
def __init__(self):
|
|
"""
|
|
PINA Implementation of a Lightning Callback for Metric Tracking.
|
|
|
|
This class provides functionality to track relevant metrics during the training process.
|
|
|
|
:ivar _collection: A list to store collected metrics after each training epoch.
|
|
|
|
:param trainer: The trainer object managing the training process.
|
|
:type trainer: pytorch_lightning.Trainer
|
|
|
|
:return: A dictionary containing aggregated metric values.
|
|
:rtype: dict
|
|
|
|
Example:
|
|
>>> tracker = MetricTracker()
|
|
>>> # ... Perform training ...
|
|
>>> metrics = tracker.metrics
|
|
"""
|
|
self._collection = []
|
|
|
|
def on_train_epoch_end(self, trainer, __):
|
|
"""
|
|
Collect and track metrics at the end of each training epoch.
|
|
|
|
:param trainer: The trainer object managing the training process.
|
|
:type trainer: pytorch_lightning.Trainer
|
|
:param _: Placeholder argument.
|
|
|
|
:return: None
|
|
:rtype: None
|
|
"""
|
|
self._collection.append(
|
|
copy.deepcopy(trainer.logged_metrics)
|
|
) # track them
|
|
|
|
@property
|
|
def metrics(self):
|
|
"""
|
|
Aggregate collected metrics during training.
|
|
|
|
:return: A dictionary containing aggregated metric values.
|
|
:rtype: dict
|
|
"""
|
|
common_keys = set.intersection(*map(set, self._collection))
|
|
v = {
|
|
k: torch.stack([dic[k] for dic in self._collection])
|
|
for k in common_keys
|
|
}
|
|
return v
|