class_weights cannot be passed via config file as a tensor is expected

[robmarkcole]

Description

Using the Lightning CLI we can train the SemanticSegmentationTask, but cannot use class_weights without an error. Solution is to accept a list if int in addition to tensor

Steps to reproduce

In Lightning CLI Yaml:

model:
  class_path: SemanticSegmentationTask
  init_args:
    model: unet
    backbone: resnet50
    weights: null
    lr: 0.001
    in_channels: 6
    num_classes: 2
    class_weights:
      - 1
      - 50

Will result in

      Does not validate against any of the Union subtypes
      Subtypes: (<class 'torch.Tensor'>, <class 'NoneType'>)
      Errors:
        - Not a valid subclass of Tensor
          Subclass types expect one of:
          - a class path (str)
          - a dict with class_path entry
          - a dict without class_path but with init_args entry (class path given previously)
        - Expected a <class 'NoneType'>
      Given value type: <class 'list'>
      Given value: [1, 50]

Version

main

[isaaccorley]

Care to make a PR to accept a list and convert to a tensor? If not then I can take it on this weekend.

[robmarkcole]

You will get to it way before me!

[adamjstewart]

For a bit of history, I added this in #1221 and it initially only supported lists. In #1413, @ntw-au modified this to support lists, numpy arrays, and torch tensors. Then in #1541, I modified it to only accept torch tensors. I agree we need a way to support class_weights in a YAML file (and preferably also on the command line). If omegaconf supports this, we could also easily enable omegaconf as a parser: https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_advanced_2.html#enable-variable-interpolation.

[isaaccorley]

If you want to use it with hydra.utils.instantiate and omegaconf you would only need to do the following:

class_weights:
   _target_: torch.tensor
   data: [0.5, 0.5]

[isaaccorley]

I haven't looked at the Lightning CLI in awhile but I wonder if it supports recursive instantiation like

class_weights:
   class_path: torch.tensor
   init_args:
      data: [0.5, 0.5]

[robmarkcole]

@isaaccorley that results in

    - Problem with given class_path 'torchgeo.trainers.SemanticSegmentationTask':
        Parser key "class_weights":
          Does not validate against any of the Union subtypes
          Subtypes: [<class 'NoneType'>, <class 'torch.Tensor'>]
          Errors:
            - Expected a <class 'NoneType'>
            - no signature found for builtin <built-in method tensor of type object at 0x7f290dc5f1c0>
          Given value type: <class 'dict'>
          Given value: {'class_path': 'torch.tensor', 'init_args': {'data': [1.0, 10.0]}}

I spent time trying to resolve this with a wrapper class, but best approach is simply to update the class:

        # in class signature
        class_weights: list[float] | None = None,

       # later
    def configure_losses(self) -> None:
        """Initialize the loss criterion.

        Raises:
            ValueError: If *loss* is invalid.
        """
        ignore_index = self.hparams['ignore_index']
        if self.hparams['class_weights'] is not None:
            weight = torch.tensor(self.hparams['class_weights'], dtype=torch.float)
        else:
            weight = None

        match self.hparams['loss']:
            case 'ce':
                ignore_value = -1000 if ignore_index is None else ignore_index
                self.criterion: nn.Module = nn.CrossEntropyLoss(
                    ignore_index=ignore_value, weight=weight
                )

[robmarkcole]

@adamjstewart I propose to update to

        class_weights: list[float] | None = None,

Any reason to keep tensor support?

[adamjstewart]

Just backwards compatibility. Want to submit a PR to support both lists and Tensors? We could even deprecate Tensors if you want.

[ntw-au]

I would like to ask, why deprecate Tensors as a type for class_weights? I originally added support (#1413, as stated) because our code instantiates the trainer in code, not from a YAML file, and it doesn't make sense to have an original tensor, explicitly convert it to list, then back into tensor. With this change we're back to how it was originally.

[adamjstewart]

We certainly don't have to deprecate Tensor inputs. When you say you have an original tensor, what do you mean? Couldn't that also be a list? I'm assuming this is manually typed by you, not coming from somewhere else.

[ntw-au]

Thank you @adamjstewart, I'll try to explain. In our ML use cases (not just with rasters or torchgeo) we tend to calculate the weights directly from the training dataset's class histogram. The histogram is calculated on live data rather than being precomputed and typed in. We do this so that we can reuse the same training pipeline on multiple input datasets with minimal data preparation.

For example (simplified):

counts = histogram(dataset)  # Returns array or tensor
weights = counts / counts.sum()  # Weights calculated on-the-fly from histogram
sst = SemanticSegmentationTask(
    ...
    num_classes=len(counts),
    class_weights=weights,
)

If we deprecate Tensor as an input to class_weights, then we'd need to explicity convert the tensor (or array) into a list, which under the hood is converted to a tensor:

counts = histogram(dataset)  # Returns array or tensor
weights = counts / counts.sum()  # Weights calculated on-the-fly from histogram
sst = SemanticSegmentationTask(
    ...
    num_classes=len(counts),
    class_weights=weights.tolist(),  # Explicit conversion to list, just so that the constructor
                                     # can turn it back into a tensor, which it already was in
                                     # the first place!
)

To convert to a list is only a single extra call (and highly unlikely to be a performance issue), but I would find it more ergonomic and less surprising to accept a tensor/array without modification as class_weights.

[robmarkcole]

I appreciate the argument for supporting tensors, but as these are not supported via Lightning config files, it could create confusion