How to add model metrics like accuracy, recall, precision etc. of segmentation task

[Raji-Ahmad]

I'm new to PyTorch and TorchGeo was what made me inerested. I have a segmentation task and a trainer as;

task = SemanticSegmentationTask(
    segmentation_model="unet",
    encoder_name="resnet18",
    encoder_weights="imagenet",
    in_channels=4,
    num_classes=len(class_names),
    loss="jaccard",
    learning_rate=learning_rate,
    learning_rate_schedule_patience=10,
)

trainer = pl.Trainer(
    default_root_dir=experiment_dir,
    min_epochs=3,
    max_epochs=10,
    accelerator='gpu',
    devices=[gpu_id]
)

After calling trainer.fit method, The process only prints the loss of train set to the screen after every batch step. How can I print loss, accuracy or f1-score of both train set and validation set after every epoch?

[adamjstewart]

You'll want to read https://pytorch-lightning.readthedocs.io/en/stable/extensions/logging.html to understand how to combine TorchMetrics and PyTorch Lightning for logging within TorchGeo.

There are actually a few questions here:

How can I print ...

Currently, only loss is displayed on the command line. To add other metrics to the progress bar, you'll need to add prog_bar=True to all self.log(...) and self.log_dict(...) calls in torchgeo/trainers/segmentation.py. However, you'll notice that the progress bar is rather short and you probably won't have enough room to see all metrics printed unless you have an ultrawide monitor. A better solution would be to use a logger like CSVLogger or TensorBoardLogger which lets you plot and visualize all metrics across all epochs. https://torchgeo.readthedocs.io/en/stable/tutorials/trainers.html has an example of the former.

... loss, accuracy or f1-score ...

At the moment, SemanticSegmentationTask only records loss, accuracy, and jaccard index. We use TorchMetrics for all metric calculations, so any metrics that TorchMetrics supports can be easily added. To add F1-score, you would add MulticlassF1Score to self.train_metrics in torchgeo/trainers/segmentation.py.

... of both train set and validation set ...

The same metrics are recorded for train, val, and test, so you shouldn't need to change anything there.

... after every epoch?

The aforementioned self.log(...) and self.log_dict(...) methods also have on_step=True and on_epoch=True parameters to control whether they are logged only during that step or for the entire epoch or both. Try playing around with these if you don't see what you want.

Hopefully this wasn't information overload, and let me know if you have any questions that the PyTorch Lightning docs didn't answer. I'd be happy to accept PRs that make any of the above modifications for all tasks!

[Raji-Ahmad]

Your answer is spot on. It's exactly what I need. I didn't want to edit the base Segmatation.py code so I made a subclass of it and override the required methods. For example;

class MySegentation(SemanticSegmentationTask):
    def __init__(self,**kwargs):
        super().__init__(**kwargs)
        
    def training_epoch_end(self, output) -> None:
        """Logs epoch level training metrics.
        Args:
            outputs: list of items returned by training_step
        """
        self.log_dict(self.train_metrics.compute(), prog_bar=True)
        self.train_metrics.reset()

Thanks for your time.