added class_weights for cross entropy loss to segmentation.py

torchgeo/trainers/segmentation.py

@@ -60,6 +60,13 @@ def config_task(self) -> None:
         if self.hyperparams["loss"] == "ce":
             ignore_value = -1000 if self.ignore_index is None else self.ignore_index
             self.loss = nn.CrossEntropyLoss(ignore_index=ignore_value)
+
+            class_weights = (
+                torch.FloatTensor(self.class_weights) if self.class_weights else None
+            )
+            self.loss = nn.CrossEntropyLoss(
+                ignore_index=ignore_value, weight=class_weights
+            )
         elif self.hyperparams["loss"] == "jaccard":
             self.loss = smp.losses.JaccardLoss(
                 mode="multiclass", classes=self.hyperparams["num_classes"]
@@ -84,6 +91,8 @@ def __init__(self, **kwargs: Any) -> None:
                 the backbone
             in_channels: Number of channels in input image
             num_classes: Number of semantic classes to predict
+            class_weights: Optional rescaling weight given to each
+                class and used with 'ce' loss
             loss: Name of the loss function, currently supports
                 'ce', 'jaccard' or 'focal' loss
             ignore_index: Optional integer class index to ignore in the loss and metrics
@@ -115,6 +124,11 @@ def __init__(self, **kwargs: Any) -> None:
                 UserWarning,
             )
         self.ignore_index = kwargs["ignore_index"]
+
+        self.class_weights = kwargs.get("class_weights", None)
+        if not isinstance(self.class_weights, (list, type(None))):
+            raise ValueError("class_weights must be a List or None")
+
         self.config_task()
 
         self.train_metrics = MetricCollection(