Fix optimizer loading + Improve configuration for imbalanced categories

dh_segment_torch/data/color_labels.py

@@ -22,11 +22,13 @@ def __init__(
         colors: List[Tuple[int, int, int]],
         one_hot_encoding: Optional[List[List[Union[int, float]]]] = None,
         labels: Optional[List[str]] = None,
+        weights: Optional[List[float]] = None,
     ):
         self.colors = colors
         self.one_hot_encoding = one_hot_encoding
         self.labels = labels
         self.log_labels = None
+        self._weights = weights
 
         if one_hot_encoding:
             if len(colors) != len(one_hot_encoding):
@@ -53,6 +55,13 @@ def __init__(
                 self.log_labels = labels
                 assert len(self.log_labels) == len(colors)
 
+        if weights:
+            if self.num_classes != len(self._weights):
+                raise ValueError(
+                    f"Cannot have a different number of classes, {self.num_classes}"
+                    f", and weights, {len(self._weights)}"
+                )
+
     @property
     def multilabel(self):
         return self.one_hot_encoding is not None
@@ -64,6 +73,10 @@ def num_classes(self):
         else:
             return len(self.colors)
 
+    @property
+    def weights(self):
+        return self._weights
+
     @classmethod
     def from_filter_by_colors(cls, color_labels, colors: Set[Tuple[int, int, int]]):
         new_colors = []

dh_segment_torch/inference/inference_model.py

@@ -1,5 +1,5 @@
 from itertools import product
-from typing import Dict, Any, Tuple, Union
+from typing import Dict, Any, Tuple, Union, Optional, List
 
 import numpy as np
 import torch
@@ -189,9 +189,10 @@ def from_partial(
         patches_batch_size: int = 4,
         model_state_dict: str = None,
         device: str = "cuda:0" if torch.cuda.is_available() else "cpu",
+        weights: Optional[List[float]] = None,
     ):
         model = model.construct(
-            num_classes=num_classes, multilabel=multilabel, margin=margin
+            num_classes=num_classes, multilabel=multilabel, margin=margin, weights=weights,
         )
 
         model_state_dict = torch.load(model_state_dict, map_location=device)
@@ -240,6 +241,7 @@ def from_color_labels_and_dataset(
             patches_batch_size,
             model_state_dict,
             device,
+            color_labels.weights,
         )
 
 

dh_segment_torch/models/model.py

@@ -131,6 +131,7 @@ def from_partial(
         classes_labels: Optional[List[str]] = None,
         ignore_padding: bool = False,
         margin: int = 0,
+        weights: Optional[List[float]] = None,
     ):
         decoder = decoder.construct(
             encoder_channels=encoder.output_dims, num_classes=num_classes
@@ -189,12 +190,12 @@ def from_partial(
         metrics = dict(zip(metric_names, metrics_built))
 
         if loss:
-            loss = loss.construct(ignore_padding=ignore_padding, margin=margin)
+            loss = loss.construct(ignore_padding=ignore_padding, margin=margin, weights=weights)
         else:
             if multilabel:
-                loss = BCEWithLogitsLoss(ignore_padding=ignore_padding, margin=margin)
+                loss = BCEWithLogitsLoss(ignore_padding=ignore_padding, margin=margin, weights=weights)
             else:
-                loss = CrossEntropyLoss(ignore_padding=ignore_padding, margin=margin)
+                loss = CrossEntropyLoss(ignore_padding=ignore_padding, margin=margin, weights=weights)
 
         return cls(encoder, decoder, loss, metrics)
 

dh_segment_torch/training/trainer.py

@@ -385,6 +385,7 @@ def from_partial(
             classes_labels=color_labels.labels,
             ignore_padding=ignore_padding,
             margin=training_margin,
+            weights=color_labels.weights,
         )
 
         parameters = [(n, p) for n, p in model.named_parameters() if p.requires_grad]

scripts/train.py

@@ -52,6 +52,10 @@
         state_dict["model"] = torch.load(args.model_checkpoint)
 
     trainer.load_state_dict(state_dict)
+    for state in trainer.optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(trainer.device)
 
     try:
         trainer.train()