CDL/NLCD/SSL4EO: allow selection of classes

[adamjstewart]

There are too many CDL classes. #1389 revealed that the top 3 classes cover more area than the bottom 130, and only 17 classes occupy more than 1% of land. We need a way to specify a smaller set of classes in order for benchmarking to be computationally tractable.

This PR introduces a classes parameter to all 3 datasets that lets users specify which classes they actually care about. In combination with the class_weights and ignore_index parameters of SemanticSegmentationTask, this will allow us to weight these remaining classes to fight class imbalance.

@nilsleh apologies, but this undoes a lot of your hard work and changes the cmaps to their previous values (albeit without unused classes). I only did NLCD to start, but see what you think of this implementation and we can update CDL and SSL4EO-L Benchmark to match. I also haven't tested any of this so it's very likely there are bugs. We should test both the getitem and plot methods. I also don't know what will happen if the user doesn't include 0 in classes.

All of the following plots have been validated in QGIS.

NLCD

All classes:

Few classes:

CDL

All classes:

Few classes:

[calebrob6]

There are too many CDL classes. #1389 revealed that the top 3 classes cover more area than the bottom 130, and only 17 classes occupy more than 1% of land. We need a way to specify a smaller set of classes in order for benchmarking to be computationally tractable.

I don't understand this motivation:

• the number of CDL classes is fixed, there aren't "too many". Long-tailed class distributions are something that needs to be dealt with in modeling.
• <256 classes is definitely not computationally intractable (e.g. imagenet has 1,000)
• Regardless, if a user doesn't want some classes, then they can remap them to 0 in a transform (see below). I don't see this as a feature that torchgeo needs for one particular dataset.

def preprocess(sample):
  mask = sample["mask"]
  mask[mask == 42] = 0
  sample["mask"] = mask
  return sample

Do you mind expanding on your thinking here?

[adamjstewart]

<256 classes is definitely not computationally intractable (e.g. imagenet has 1,000)

It isn't for classification, but it is for semantic segmentation. Our CDL benchmarks take 60x longer than our NLCD benchmarks. At the moment, it's basically impossible for us to finish all CDL benchmarks before the deadline. This is what we decided in the last meeting.

[calebrob6]

I can't reproduce this behavior

[calebrob6]

It is identical with 20 and 130 output classes, the last classification layer should not be a large fraction of the computation done in a UNet.

[nilsleh]

I think the story changes when you also include the backward pass, that is also what we found when using the pytorch Profiler.

[calebrob6]

There's some difference but not double

[adamjstewart]

@yichiac tried training on our cluster and each epoch took 60x longer for CDL than NLCD. @nilsleh tried on his cluster and it was 10x. @calebrob6 if you want to train on your GPUs that's fine, but it's not feasible on our systems. This is what we spent a large portion of the last meeting discussing and this is the decision that we made. We've been trying to figure this out for days and no one has solved it and we're rapidly running out of time to start benchmarking. We have 65 CDL benchmarks to fill in. If each CDL benchmark takes > 24 hrs to complete, we're going to need quite a lot of GPUs...

[calebrob6]

That's fine, but I don't think it needs to be baked into the library is what I'm saying.

[adamjstewart]

Writing transforms is a little painful at the moment but we can think about moving this into a datamodule transform before the next release.

[adamjstewart]

https://openaccess.thecvf.com/content/ICCV2021/papers/Jain_Scaling_Semantic_Segmentation_Beyond_1K_Classes_on_a_Single_GPU_ICCV_2021_paper.pdf

[calebrob6]

Yes, it will take more memory, but the computation cost is linear in the number of classes.

See how the linear function fit on "time for inference + backprop" from 20 to 500 classes predicts the "time for inference + backprop" for 1000 classes: