SemanticSegmentationTask: add class-wise metrics

[robmarkcole]

Addresses #2121 for segmentation. Mostly copied from @isaaccorley as here - he is additionally passing on_epoch=True which is also adopted here

[DimitrisMantas]

Given that most metrics of interest are broken (e.g., all of them when average="macro" and ignore_index is specified (Lightning-AI/torchmetrics#2443) andJaccardIndex which outputs NaN when average==macro instead when you try to take absent and ignored classes into account with zero_division (Lightning-AI/torchmetrics#2535)), should we make an effort to see if and how we could add our own?

I'm saying this because these are only the issues I've found so far, but I've also noticed other suspicious things like the fact that my classwise recall values are not the same as those in the confusion matrix when you normalize it with respect to ground truth (I haven't checked if this is also the case with precision, so when the matrix is normalized column-wise). I'm also pretty confident that if all of this is wrong then micro averaging is also probably wrong.

I should be pretty easy to compute all these metrics straight from the confusion matrix (assuming it at least is correct) and I've actually tried to reimplent them this way but it hasn't really been a priority because I’ve found that all these wrong (?) values are basically a lower bound of the actual ones. If you look at the official implementations, this is actually what they are doing, and my guess is that they have a bug in their logic later on. But indeed all these metrics inherit from StatScores, basically the confusion matrix.

I’m actually pretty dumbfounded these issues are not a top priority for the TorchMetrics team and instead they focus on adding to their docs but to each their own…

[robmarkcole]

@DimitrisMantas good call on my ignoring the ignore_index.! In fairness they do address issues, but have a long backlog. When I made some noise they addressed Lightning-AI/torchmetrics#2198
My opinion is it is better to work with torchmetrics to address the issues, rather than implement from scratch here. I see your comment at Lightning-AI/torchmetrics#2535 (comment) so perhaps a pragmatic approach is not to add new metrics that we have concerns about, but also to create specific issues which track these concerns

[DimitrisMantas]

Sure, that makes sense; please excuse the rant haha.

[robmarkcole]

Applied on_epoch=True, to all steps for consistency - this results in both per epoch and per step being reported for train only - perhaps this is why @isaaccorley did not apply to train?

train_loss_epoch | 0.028535427525639534
train_loss_step | 0.00008003244874998927
train_AverageAccuracy_epoch | 0.9101453423500061
train_AverageAccuracy_step | 0.9124529361724854

Note that Val is unaffected:

val_AverageAccuracy | 0.8227439522743225

For a task with 2 classes there are a grand total of Metrics (52) being reported between train & val

[isaaccorley]

I just set to be explicit but I think that pytorch lightning or torchmetrics auto sets on_epoch to be False for training and True for all else.

[DimitrisMantas]

You need to set both on_step and on_epoch to get logs only per step or per epoch.

[robmarkcole]

@DimitrisMantas now just performing on_step for train loss, so a more manageable 36 metrics now

[robmarkcole]

Not sure about this failing test ValueError: Problem with given class_path 'torchgeo.trainers.SemanticSegmentationTask'

[isaaccorley]

Must be an issue with on of the minimum versions of the package since it's passing for the other tests.

[adamjstewart]

We can definitely increase the min version of torchmetrics if we need to.

[DimitrisMantas]

@robmarkcole I can confirm the recommended approach yields consistent results.

[adamjstewart]

Sorry it's taken me so long to review. I was originally hung up on the hack required to support ClasswiseWrapper in log_dict, but I've gotten over that. If torchmetrics wants to make that easier in the future, great. But I also really want this feature, so let's not wait on that.

Only remaining concern is that the code required to loop over all metrics and averages actually makes the code more complicated and difficult to read than avoiding loops entirely. If we want to add new metrics in the future, it looks non-straightforward. I wonder if we can loop over averages only and still keep things simple.

I would also really like to see this done for ClassificationTask too so SemanticSegmentationTask doesn't have a different set of features or metrics.

[robmarkcole]

@adamjstewart please see above comment on using _epoch_end

[isaaccorley]

I think it's better practice to do on_epoch=True for val and test metrics so that you will get a single value per metric per epoch for the validation and test set. Otherwise you will get a metric per step which is not great for comparing performance across models.

[robmarkcole]

To clarify if we want on_step=False, or True for train. I had a look what others do, rslearn do not but clay do

[isaaccorley]

That's a good question, I normally set the train to only record steps and not the full epoch average because depending on how many metrics and the size of your train set this can use up a lot of memory. I think for now we can set both as true.

[robmarkcole]

From memory I think if using both the loss is named loss_step and loss_epoch, trying to get a model trained to demonstrate. I assume most people will be fine with this, but it could require some explanation

[isaaccorley]

This LGTM. Any concerns with merging this @adamjstewart?

[adamjstewart]

I don't think these concerns have been addressed: #2130 (comment)

[isaaccorley]

I agree, I think just making a list of the metrics instead of using loops is probably more readable and easier to add/remove individual metrics based on a users need

[robmarkcole]

merging main really messed this branch up, will create a new branch/PR

[calebrob6]

Coming back to this to say I just reimplemented in a private project, would love to have this upstream!

[robmarkcole]

Honestly dont know when I will get around to it - the last occasion I had a few hours I ended up debugging dataset issues

[adamjstewart]

According to Lightning-AI/torchmetrics@63c7bbe the argument didn't exist until 1.2