micro_sam.training.semantic_sam_trainer

  1import time
  2from typing import Optional
  3
  4import torch
  5import torch.nn as nn
  6
  7from torch_em.loss import DiceLoss
  8from torch_em.trainer import DefaultTrainer
  9
 10
 11class CustomDiceLoss(nn.Module):
 12    """Loss for computing dice over one-hot labels.
 13
 14    Expects prediction and target with `num_classes` channels: the number of classes for semantic segmentation.
 15
 16    Args:
 17        num_classes: The number of classes for semantic segmentation (including background class).
 18        softmax: Whether to use softmax over the predictions.
 19    """
 20    def __init__(self, num_classes: int, softmax: bool = True) -> None:
 21        super().__init__()
 22        self.num_classes = num_classes
 23        self.dice_loss = DiceLoss()
 24        self.softmax = softmax
 25
 26    def _one_hot_encoder(self, input_tensor):
 27        tensor_list = []
 28        for i in range(self.num_classes):
 29            temp_prob = input_tensor == i  # * torch.ones_like(input_tensor)
 30            tensor_list.append(temp_prob)
 31        output_tensor = torch.cat(tensor_list, dim=1)
 32        return output_tensor.float()
 33
 34    def __call__(self, pred, target):
 35        if self.softmax:
 36            pred = torch.softmax(pred, dim=1)
 37        target = self._one_hot_encoder(target)
 38        loss = self.dice_loss(pred, target)
 39        return loss
 40
 41
 42class SemanticSamTrainer(DefaultTrainer):
 43    """Trainer class for training the Segment Anything model for semantic segmentation.
 44
 45    This class is derived from `torch_em.trainer.DefaultTrainer`.
 46    Check out https://github.com/constantinpape/torch-em/blob/main/torch_em/trainer/default_trainer.py
 47    for details on its usage and implementation.
 48
 49    Args:
 50        convert_inputs: The class that converts outputs of the dataloader to the expected input format of SAM.
 51            The class `micro_sam.training.util.ConvertToSemanticSamInputs` can be used here.
 52        num_classes: The number of classes for semantic segmentation (including the background class).
 53        dice_weight: The weighing for the dice loss in the combined dice-cross entropy loss function.
 54        kwargs: The keyword arguments of the DefaultTrainer super class.
 55    """
 56
 57    def __init__(self, convert_inputs, num_classes: int, dice_weight: Optional[float] = None, **kwargs):
 58        assert num_classes > 1
 59
 60        if "loss" not in kwargs:
 61            kwargs["loss"] = CustomDiceLoss(num_classes=num_classes)
 62
 63        if "metric" not in kwargs:
 64            kwargs["metric"] = CustomDiceLoss(num_classes=num_classes)
 65
 66        super().__init__(**kwargs)
 67
 68        self.convert_inputs = convert_inputs
 69        self.num_classes = num_classes
 70        self.compute_ce_loss = nn.CrossEntropyLoss()
 71        self.dice_weight = dice_weight
 72
 73        if self.dice_weight is not None and (self.dice_weight < 0 or self.dice_weight > 1):
 74            raise ValueError("The weight factor should lie between 0 and 1.")
 75
 76        self._kwargs = kwargs
 77
 78    def _compute_loss(self, y, masks):
 79        """Compute the combined (weighted) dice loss and cross-entropy loss between the prediction and target.
 80        """
 81        target = y.to(self.device, non_blocking=True)
 82        # Compute dice loss for the predictions
 83        dice_loss = self.loss(masks, target)
 84
 85        # Compute cross entropy loss for the predictions
 86        ce_loss = self.compute_ce_loss(masks, target.squeeze(1).long())
 87
 88        if self.dice_weight is None:
 89            net_loss = dice_loss + ce_loss
 90        else:
 91            net_loss = self.dice_weight * dice_loss + (1 - self.dice_weight) * ce_loss
 92
 93        return net_loss
 94
 95    def _get_model_outputs(self, batched_inputs):
 96        """Get the predictions from the model.
 97        """
 98        # Precompute the image embeddings if the model exposes it as functionality.
 99        if hasattr(self.model, "image_embeddings_oft"):
100            image_embeddings, batched_inputs = self.model.image_embeddings_oft(batched_inputs)
101            batched_outputs = self.model(batched_inputs, image_embeddings, multimask_output=True)
102        else:  # Otherwise we assume that the embeddings are computed internally as part of the forward pass.
103            # We need to take care of sending things to the device here.
104            batched_inputs = [
105                {"image": inp["image"].to(self.device, non_blocking=True), "original_size": inp["original_size"]}
106                for inp in batched_inputs
107            ]
108            batched_outputs = self.model(batched_inputs, multimask_output=True)
109
110        masks = torch.stack([output["masks"].squeeze(0) for output in batched_outputs])
111        return masks
112
113    def _train_epoch_impl(self, progress, forward_context, backprop):
114        self.model.train()
115
116        t_per_iter = time.time()
117        for x, y in self.train_loader:
118            self.optimizer.zero_grad()
119
120            batched_inputs = self.convert_inputs(x, y)
121
122            with forward_context():
123                masks = self._get_model_outputs(batched_inputs)
124                net_loss = self._compute_loss(y, masks)
125
126            backprop(net_loss)
127
128            self._iteration += 1
129
130            if self.logger is not None:
131                lr = [pm["lr"] for pm in self.optimizer.param_groups][0]
132                self.logger.log_train(
133                    self._iteration, net_loss, lr, x, y, torch.softmax(masks, dim=1), log_gradients=False
134                )
135
136            if self._iteration >= self.max_iteration:
137                break
138            progress.update(1)
139
140        t_per_iter = (time.time() - t_per_iter)
141        return t_per_iter
142
143    def _validate_impl(self, forward_context):
144        self.model.eval()
145
146        metric_val, loss_val = 0.0, 0.0
147
148        with torch.no_grad():
149            for x, y in self.val_loader:
150                batched_inputs = self.convert_inputs(x, y)
151
152                with forward_context():
153                    masks = self._get_model_outputs(batched_inputs)
154                    net_loss = self._compute_loss(y, masks)
155
156                loss_val += net_loss.item()
157                metric_val += net_loss.item()
158
159        loss_val /= len(self.val_loader)
160        metric_val /= len(self.val_loader)
161        dice_metric = 1 - (metric_val / self.num_classes)
162        print()
163        print(f"The Average Validation Metric Score for the Current Epoch is {dice_metric}")
164
165        if self.logger is not None:
166            self.logger.log_validation(
167                self._iteration, metric_val, loss_val, x, y, torch.softmax(masks, dim=1)
168            )
169
170        return metric_val
171
172
173class SemanticMapsSamTrainer(SemanticSamTrainer):
174    def _compute_loss(self, y, masks):
175        target = y.to(self.device, non_blocking=True)
176
177        # Compute loss for the predictions
178        net_loss = self.loss(target, masks)
179
180        return net_loss