Configurable Metrics?

[e-yi]

Currently, metrics are defined in src. models (LightningModule), along with the training procedures. Maybe these two can be decoupled? I think metircs are usually task-specific, and the training procedures are more model-specific?

I've made some modifications to enable configurable metrics here.

Any feedback on this idea will be greatly appreciated!

[ashleve]

@e-yi Thanks for interesting proposition. I took a closer look at the code.

So the way I see it:

• Task - in the template, I consider “task” to be the entry file, like train.py. It can have many phases and make use of multiple lightning modules if needed.
• LightningModule - architecture is often closely tied to the way train/val/test/predict steps are performed, and also to loss and optimiser. So personally, when I say “model” I usually mean LightningModule (architecture + procedures).

Generally, I think metrics are strongly tied to the LightningModule. As such, making them pluggable through a separate subconfig doesn’t seem to bring much value for me. I currently don’t see a scenario when it’s useful to switch metrics during development - I might add more metrics to my lightning module over time, but I don’t see the need of replacing one metric with the other when I can just log everything.

If I switch metric, then usually it means I switch the whole LightningModule since many of training procedures also need to be changed.

So I think it's easiest to treat metrics as something specific for given LightningModule. You could decouple metrics if it makes it faster for you to implement new LightningModules. But if the project isn't enourmous then I like the minimal design aesthetic. In small projects there is no need to spread your code over complex abstractions making it harder to read

[e-yi]

Hi @ashleve. Thanks for your reply! I think we have different understanding over the word 'task', I'm using 'task' as a specific definition of a problem, like giving the probability that a picture is of a cat. I'd like to provide some situations where I think pluggable metrics might be an elegant solution:

1. Situations where a problem can be defined as a classification task, a regression task, or maybe a rank task, such as finding candidate books that are likely to become a hit. In this situation, you may have to try all different tasks to know which one performs best. Some metrics, such as the Spearman correlation coefficient calculated based on the ranking of test data, can be shared across different tasks, while others, such as AUROC in classification or MAE in regression, cannot. In the meantime, the overall training procedure stays the same as well as the backbone model.

2. Situations where the goal is to evaluate the performance of a new model across different tasks. For example, if a new GNN structure is developed, one might want to evaluate its performance across different tasks, such as link prediction, node property prediction, and node classification.

3. Situations where modifications are made to the training or prediction procedures, but the evaluation remains the same. For example, when one wants to add a regularization term to the loss function or ensemble different models during prediction.

In the above siutations, pluggable metrics could be helpful in decressing the redundancy in the LightingModule. Besides, using configurable metrics eliminates the need to define a different metric for each pharse (train/val/test) and manually update and log it. I think this makes managing metrics much easier.

On the other hand, the listed situations are common in my daily development, though perhaps not for others. In that case, I agree with you that a template following the minimal design aesthetic may not need this feature.

[ashleve]

Thank you for the comprehensive explanation. I can see more clearly now pluggable metrics are very useful when using along with your metric group abstraction.

Specifically, I like how we could have exactly the same logging calls in all lightning modules. For example, one could slightly modify your implementation to allow for configs like this:

metrics_cfg = {
        "acc": {
            "metric": {
                "_target_": "torchmetrics.Accuracy",
                "task": "multiclass",
                "num_classes": 10,
            },
            "ordered_inputs": ["preds", "targets"],
            "phases": ["train", "val", "test"],
            "log_levels": ["epoch"],
            "track_max": True,
            "track_min": False,
        },
        "loss": {
            "metric": {
                "_target_": "torchmetrics.MeanMetric",
            },
            "ordered_inputs": ["loss"],
            "phases": ["train", "val", "test"],
            "log_levels": ["epoch"],
            "track_max": False,
            "track_min": True,
        },
    }

metrics = AdvancedMetricCollection(metrics_cfg)

I added track_max, track_max so that tracking maximum or minimum of given metric could happen automatically.

Now If I'm not mistaken, we could have exactly the same logging calls in all LightningModules:

def on_train_start(self):
        self.metrics.reset(phase="all")

def train_step_end(self, outputs: Dict):
        step_metrics = self.metrics.update(outputs, phase="train")
        self.log_dict(step_metrics)

def on_train_epoch_end(self):
        epoch_metrics = self.metrics.compute_epoch_metrics(phase="train")
        self.log_dict(epoch_metrics)

def validation_step_end(self, outputs: Dict):
        step_metrics = self.metrics.update(outputs, phase="val")
        self.log_dict(step_metrics)

def on_validation_epoch_end(self):
        epoch_metrics = self.metrics.compute_epoch_metrics(phase="val")
        self.log_dict(epoch_metrics)

def test_step_end(self, outputs: Dict):
        step_metrics = self.metrics.update(outputs, phase="test")
        self.log_dict(step_metrics)

def on_test_epoch_end(self):
        epoch_metrics = self.metrics.compute_epoch_metrics(phase="test")
        self.log_dict(epoch_metrics)

def on_epoch_end(self):
        self.metrics.reset(phase="all")

This indeed seems very convenient.

Since all calls are the same, maybe we could also encapsulate them in a new abstaction that inherits from LightningModule.

Some issues I have:

• I suspect there might be some edge cases for some types of metrics, when this metric collection doesn't allow to achieve something, but so far I have not been able to spot any issue.
• According to torchmetrics docs:

“Modular metrics are automatically placed on the correct device when properly defined inside a LightningModule. This means that your data will always be placed on the same device as your metrics. No need to call .to(device) anymore!”

I think this means such metric collection might not work correctly on GPU or in DDP mode but I'm not sure, would need to investigate more.

• Since template is also meant to be educational, I don’t want to make metric collection as default so that instead users can understand torchmetrics trouroughly. We could mention this approach in the readme though.

Perhaps it would be worth to implement this all as a separate pip package. Have you thought about that?

[e-yi]

I think this means such metric collection might not work correctly on GPU or in DDP mode but I'm not sure, would need to investigate more.

I agree more investagation is needed.

Perhaps it would be worth to implement this all as a separate pip package. Have you thought about that?

No, but now that you mentioned it, I will consider this. It may take some time since I am kind busy right now. I think maybe this package could provide two interfaces, a class inherits from LightningModule and a class decotrator of Lightning Module.

@log_metrics(config_path="path/to/config_file")
class MyModule(LightningModule):
    ...

or

class MyModule(LightningModuleWithMetrics):
    def __init__(self, metric_config:Path, ...):
        super(metric_config)
        ...

[ashleve]

I like these interfaces. Looking forward to your package!

[e-yi]

For anyone interested in this topic, I have developed a package named PackMetric that provides a callback for seamless integration into PyTorch Lightning and can be easily configured.

As this is my first package, I'm actively seeking feedback to improve its design and usability. I'd really appreciate any suggestions from the community, especially from those who have tackled similar challenges.

@ashleve, given this cool template you built, your insights would be particularly valuable. Please take a look at PackMetric when you have a moment and let me know what you think!