post-processing

[NastaranVB]

Hello,
I am doing multi class 3D segmentation which contains 3 classes including background. I have used Network, Loss and Metric from Monai. Training step in my code is as below:

class Net(pytorch_lightning.LightningModule):
    def __init__(self):
        super().__init__()
        self._model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=3,
            channels=(16, 32, 64, 128, 256),
            strides=(2, 2, 2, 2),
       )
        self.loss_function = DiceLoss(include_background=False, to_onehot_y=True, softmax=True)
        self.post_pred = AsDiscrete(argmax=True, to_onehot=3)  
        self.post_label = AsDiscrete(to_onehot=3)
        self.training_step_outputs = [] 
        self.validation_step_outputs = []

     def forward(self, data):
        pred = self.model(data)
        return pred
          
    def training_step(self, batch, batch_idx):       
        img = batch["Image"]["data"]
        mask = batch["Label"]["data"]       
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        self.training_step_outputs.append(loss)
        self.logger.experiment.add_scalars('loss', {'Train': loss}, self.global_step)
        return loss

   def on_train_epoch_end(self):
        train_epoch_average = torch.stack(self.training_step_outputs).mean()
        self.logger.experiment.add_scalars('loss/epoch', {'Train': train_epoch_average}, global_step=self.current_epoch)
        self.training_step_outputs.clear()
        
    def validation_step(self, batch, batch_idx):        
        img = batch["Image"]["data"]
        mask = batch["Label"]["data"]     
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        preds = [self.post_pred(i) for i in decollate_batch(pred)]
        masks = [self.post_label(i) for i in decollate_batch(mask)]
         self.validation_step_outputs.append(loss)     
        self.logger.experiment.add_scalars('loss', {'valid': loss},self.global_step) 
        return loss

   def on_validation_epoch_end(self):
        val_epoch_average = torch.stack(self.validation_step_outputs).mean()
        self.logger.experiment.add_scalars('loss/epoch', {'Val':val_epoch_average}, global_step=self.current_epoch)
        self.validation_step_outputs.clear()
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self._model.parameters(), 1e-4)
        return optimizer

The model has been trained in a patch wise manner and getting a result for the whole volume has been done as following:

imgs = val_dataset[IDX]["Image"]["data"]
mask = val_dataset[IDX]["Label"]["data"]
grid_sampler = tio.inference.GridSampler(val_dataset[IDX], patch_size=patch_size, patch_overlap=patch_overlap)
aggregator = tio.inference.GridAggregator(grid_sampler)
patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=val_batch_size)
with torch.no_grad():
    for patches_batch in patch_loader:
        input_tensor = patches_batch['Image']["data"].to(device)  
        locations = patches_batch[tio.LOCATION]  
        pred = model(input_tensor)  
        aggregator.add_batch(pred, locations)
output_tensor = aggregator.get_output_tensor()

I wonder if you could please tell me where I have to do the post-processing?
I want to see the results even in average loss in the Tensorboard , so should I do post-processing in the training step as well?
As a last question am i right in "def on_train_epoch_end(self)" (because def training_epoch_end(self, outputs) [3D Multi-organ Segmentation with UNETR] (https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/unetr_btcv_segmentation_3d_lightning.ipynb) has been removed from the last version of pytorch lightning)
Thank you in advance.

[KumoLiu]

Hi @NastaranVB,
Post-processing typically happens after the model inference step, usually with the aim to convert model outputs into a more presentable and interpretable form. If you're looking to visualize epoch-wise average loss in Tensorboard, that's typically handled within the training step and doesn't necessarily fall under "post-processing". Loss values are calculated within the epoch and are then logged to Tensorboard.
If you're using PyTorch Lightning for training your model, you can indeed use the on_train_epoch_end(self) method to perform any operations at the end of each training epoch. This will log the value of train_loss to Tensorboard, and you can do similar for any other metrics you want to keep track of. You can indeed do this for the average loss, typically by keeping a running total of the loss values and the number of batches, and then calculating the average at the end of each epoch.
Regarding the absent training_epoch_end in recent PyTorch Lightning versions, it's been replaced by on_train_epoch_end. You're correct. You should replace it with on_train_epoch_end when using more recent versions of PyTorch Lightning.

Thanks.

[NastaranVB]

Hi @KumoLiu ,
Thank you so much for explanation. As I am new in both pytorch lightning and Monai, please let me ask another question related to my previous question. In Medical image segmentation with TorchIO, MONAI & PyTorch Lightning (https://github.com/Project-MONAI/tutorials/blob/main/modules/TorchIO_MONAI_PyTorch_Lightning.ipynb), if I want to log both the training and validation loss curves in Tensorboard, could you please tell me what changes I have to do? And how should I consider "post-processing"?
Thank you in advance.

[NastaranVB]

Hi @KumoLiu,
I wanted to follow up on an inquiry I made some days ago. Since you receive many emails daily, I thought it might have been missed. Could you please revisit my question and assist ? Any help would be appreciated.
Warm regards.

[KumoLiu]

Hi @NastaranVB,

I believe in this tutorial, both training and validation losses are logged.
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d_lightning.ipynb

Thanks.

[NastaranVB]

Hi @KumoLiu,
Thank you for coming back to me so quickly.
My question was regarding to In Medical image segmentation with TorchIO, MONAI & PyTorch Lightning (https://github.com/Project-MONAI/tutorials/blob/main/modules/TorchIO_MONAI_PyTorch_Lightning.ipynb), which does not include "post-processing" and I want to know how should I consider "post-processing"? Isn't "post-processing" needed to be considered for logging validation loss?
Thanks in advance.

[KumoLiu]

Post-processing in the context of medical image segmentation usually refers to operations performed on the model's outputs to convert them into a more interpretable or useful form. Some common post-processing steps include applying a threshold to get binary segmentation masks, removing small isolated regions, filling holes, and mapping the segmented volumes back to the original image space from any preprocessed space. The goal is to clean up and refine the raw segmentation output of the model to make it suitable for analysis or comparison with ground truth.
Logging the validation loss typically doesn't require post-processing. Validation loss is a measure computed during model training to assess how well the model's current predictions match the true labels. These predictions are usually directly outputted by the model, and don't need to be post-processed before calculating the loss. The loss function used (like Dice loss, Cross-Entropy loss etc.) should align with the raw outputs of the model (like logits or probabilities).

Thanks

[NastaranVB]

Thank you for a nice explanation.