.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_tutorials/Bayesian_Methods/tutorial_muad_mc_drop.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_tutorials_Bayesian_Methods_tutorial_muad_mc_drop.py: Monte Carlo Dropout for Semantic Segmentation on MUAD ===================================================== This tutorial demonstrates how to train a segmentation model on the MUAD dataset using TorchUncertainty. MUAD is a synthetic dataset designed for evaluating autonomous driving under diverse uncertainties. It includes **10,413 images** across training, validation, and test sets, featuring adverse weather, lighting conditions, and out-of-distribution (OOD) objects. The dataset supports tasks like semantic segmentation, depth estimation, and object detection. For details and access, visit the `MUAD Website `_. 1. Loading the utilities ~~~~~~~~~~~~~~~~~~~~~~~~ First, we load the following utilities from TorchUncertainty: - the TUTrainer which mostly handles the link with the hardware (accelerators, precision, etc) - the segmentation training & evaluation routine from torch_uncertainty.routines - the datamodule handling dataloaders: MUADDataModule from torch_uncertainty.datamodules - the model: small_unet from torch_uncertainty.models.segmentation.unet - the MC Dropout wrapper from torch_uncertainty.models .. GENERATED FROM PYTHON SOURCE LINES 27-43 .. code-block:: Python import matplotlib.pyplot as plt import torch import torchvision.transforms.v2.functional as F from huggingface_hub import hf_hub_download from torch import optim from torch.optim import lr_scheduler from torchvision import tv_tensors from torchvision.transforms import v2 from torchvision.utils import draw_segmentation_masks from torch_uncertainty import TUTrainer from torch_uncertainty.datamodules.segmentation import MUADDataModule from torch_uncertainty.models import mc_dropout from torch_uncertainty.models.segmentation.unet import small_unet from torch_uncertainty.routines import SegmentationRoutine .. GENERATED FROM PYTHON SOURCE LINES 44-46 2. Initializing the DataModule ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 46-93 .. code-block:: Python muad_mean = MUADDataModule.mean muad_std = MUADDataModule.std train_transform = v2.Compose( [ v2.Resize(size=(256, 512), antialias=True), v2.RandomHorizontalFlip(), v2.ToDtype( dtype={ tv_tensors.Image: torch.float32, tv_tensors.Mask: torch.int64, "others": None, }, scale=True, ), v2.Normalize(mean=muad_mean, std=muad_std), ] ) test_transform = v2.Compose( [ v2.Resize(size=(256, 512), antialias=True), v2.ToDtype( dtype={ tv_tensors.Image: torch.float32, tv_tensors.Mask: torch.int64, "others": None, }, scale=True, ), v2.Normalize(mean=muad_mean, std=muad_std), ] ) # datamodule providing the dataloaders to the trainer datamodule = MUADDataModule( root="./data", batch_size=10, version="small", train_transform=train_transform, test_transform=test_transform, num_workers=4, ) datamodule.prepare_data() datamodule.setup("fit") .. GENERATED FROM PYTHON SOURCE LINES 94-95 Visualize a validation input sample (and RGB image) .. GENERATED FROM PYTHON SOURCE LINES 95-109 .. code-block:: Python # Undo normalization on the image and convert to uint8. img, tgt = datamodule.train[0] t_muad_mean = torch.tensor(muad_mean, device=img.device) t_muad_std = torch.tensor(muad_std, device=img.device) img = img * t_muad_std[:, None, None] + t_muad_mean[:, None, None] img = F.to_dtype(img, torch.uint8, scale=True) img_pil = F.to_pil_image(img) plt.figure(figsize=(6, 6)) plt.imshow(img_pil) plt.axis("off") plt.show() .. image-sg:: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_001.png :alt: tutorial muad mc drop :srcset: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 110-111 Visualize the same image above but segmented. .. GENERATED FROM PYTHON SOURCE LINES 111-124 .. code-block:: Python tmp_tgt = tgt.masked_fill(tgt == 255, 21) tgt_masks = tmp_tgt == torch.arange(22, device=tgt.device)[:, None, None] img_segmented = draw_segmentation_masks( img, tgt_masks, alpha=1, colors=datamodule.train.color_palette ) img_pil = F.to_pil_image(img_segmented) plt.figure(figsize=(6, 6)) plt.imshow(img_pil) plt.axis("off") plt.show() .. image-sg:: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_002.png :alt: tutorial muad mc drop :srcset: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 125-128 3. Instantiating the Model ~~~~~~~~~~~~~~~~~~~~~~~~~~ We create the model easily using the blueprint from torch_uncertainty.models. .. GENERATED FROM PYTHON SOURCE LINES 128-136 .. code-block:: Python model = small_unet( in_channels=datamodule.num_channels, num_classes=datamodule.num_classes, bilinear=True, dropout_rate=0.1, # We use dropout to enable MC Dropout later ) .. GENERATED FROM PYTHON SOURCE LINES 137-139 4. Compute class weights to mitigate class inbalance ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 139-187 .. code-block:: Python def enet_weighing(dataloader, num_classes, c=1.02): """Computes class weights as described in the ENet paper. w_class = 1 / (ln(c + p_class)), where c is usually 1.02 and p_class is the propensity score of that class: propensity_score = freq_class / total_pixels. References: https://arxiv.org/abs/1606.02147 Args: dataloader (``data.Dataloader``): A data loader to iterate over the dataset. num_classes (``int``): The number of classes. c (``int``, optional): AN additional hyper-parameter which restricts the interval of values for the weights. Default: 1.02. ignore_indexes (``list``, optional): A list of indexes to ignore when computing the weights. Default to `None`. """ class_count = 0 total = 0 for _, label in dataloader: label = label.cpu() # Flatten label flat_label = label.flatten() flat_label = flat_label[flat_label != 255] flat_label = flat_label[flat_label < num_classes] # Sum up the number of pixels of each class and the total pixel # counts for each label class_count += torch.bincount(flat_label, minlength=num_classes) total += flat_label.size(0) # Compute propensity score and then the weights for each class propensity_score = class_count / total return 1 / (torch.log(c + propensity_score)) class_weights = enet_weighing(datamodule.val_dataloader(), datamodule.num_classes) print(class_weights) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([ 4.3817, 19.7927, 3.3011, 48.8031, 36.2141, 33.0049, 47.5130, 48.8560, 12.4401, 48.0600, 14.4807, 30.8762, 4.7467, 19.3913, 50.4984]) .. GENERATED FROM PYTHON SOURCE LINES 188-189 Let's define the training parameters. .. GENERATED FROM PYTHON SOURCE LINES 189-197 .. code-block:: Python BATCH_SIZE = 10 LEARNING_RATE = 1e-3 WEIGHT_DECAY = 2e-4 LR_DECAY_EPOCHS = 20 LR_DECAY = 0.1 NB_EPOCHS = 1 .. GENERATED FROM PYTHON SOURCE LINES 198-200 5. The Loss, the Routine, and the Trainer ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 200-216 .. code-block:: Python # We build the optimizer optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) # Learning rate decay scheduler lr_updater = lr_scheduler.StepLR(optimizer, step_size=LR_DECAY_EPOCHS, gamma=LR_DECAY) # Segmentation Routine seg_routine = SegmentationRoutine( model=model, num_classes=datamodule.num_classes, loss=torch.nn.CrossEntropyLoss(weight=class_weights), optim_recipe={"optimizer": optimizer, "lr_scheduler": lr_updater}, ) trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=NB_EPOCHS, enable_progress_bar=False) .. GENERATED FROM PYTHON SOURCE LINES 217-219 6. Training the model ~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 219-220 .. code-block:: Python trainer.fit(model=seg_routine, datamodule=datamodule) .. GENERATED FROM PYTHON SOURCE LINES 221-223 7. Testing the model ~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 223-224 .. code-block:: Python results = trainer.test(datamodule=datamodule, ckpt_path="best") .. rst-class:: sphx-glr-script-out .. code-block:: none ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Segmentation ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ Brier │ 0.63934 │ │ NLL │ 1.57190 │ │ mAcc │ 26.578% │ │ mIoU │ 19.289% │ │ pixAcc │ 63.958% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Calibration ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ ECE │ 32.290% │ │ aECE │ 32.290% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Selective Classification ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ AUGRC │ 13.225% │ │ AURC │ 20.583% │ │ Cov@5Risk │ nan% │ │ Risk@80Cov │ 30.103% │ └──────────────┴───────────────────────────┘ .. GENERATED FROM PYTHON SOURCE LINES 225-228 8. Loading a pre-trained model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Let's now load a fully trained model to continue this tutorial .. GENERATED FROM PYTHON SOURCE LINES 228-234 .. code-block:: Python model_path = hf_hub_download(repo_id="torch-uncertainty/muad_tutorials", filename="small_unet.pth") model.load_state_dict(torch.load(model_path)) # Replace the model in the routine seg_routine.model = model # Test the model results = trainer.test(model=seg_routine, datamodule=datamodule) .. rst-class:: sphx-glr-script-out .. code-block:: none ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Segmentation ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ Brier │ 0.20635 │ │ NLL │ 0.40168 │ │ mAcc │ 68.773% │ │ mIoU │ 50.739% │ │ pixAcc │ 85.701% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Calibration ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ ECE │ 1.731% │ │ aECE │ 1.874% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Selective Classification ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ AUGRC │ 2.733% │ │ AURC │ 3.367% │ │ Cov@5Risk │ 72.133% │ │ Risk@80Cov │ 7.174% │ └──────────────┴───────────────────────────┘ .. GENERATED FROM PYTHON SOURCE LINES 235-238 9. Uncertainty evaluations with MCP ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Here we will just use as confidence score the Maximum class probability (MCP) .. GENERATED FROM PYTHON SOURCE LINES 238-286 .. code-block:: Python sample_idx = 0 img, target = datamodule.test[sample_idx] batch_img = img.unsqueeze(0) batch_target = target.unsqueeze(0) model.eval() with torch.no_grad(): # Forward propagation output_probs = model(batch_img).softmax(dim=1) # remove the batch dimension output_probs = output_probs.squeeze(0) confidence, pred = output_probs.max(0) # Undo normalization on the image and convert to uint8. img = img * t_muad_std[:, None, None] + t_muad_mean[:, None, None] img = F.to_dtype(img, torch.uint8, scale=True) tmp_target = target.masked_fill(target == 255, 21) target_masks = tmp_target == torch.arange(22, device=target.device)[:, None, None] img_segmented = draw_segmentation_masks( img, target_masks, alpha=1, colors=datamodule.test.color_palette ) pred_masks = pred == torch.arange(22, device=pred.device)[:, None, None] pred_img = draw_segmentation_masks(img, pred_masks, alpha=1, colors=datamodule.test.color_palette) if confidence.ndim == 2: confidence = confidence.unsqueeze(0) img = F.to_pil_image(F.resize(img, 1024)) img_segmented = F.to_pil_image(F.resize(img_segmented, 1024)) pred_img = F.to_pil_image(F.resize(pred_img, 1024)) confidence_img = F.to_pil_image(F.resize(confidence, 1024)) fig, axs = plt.subplots(1, 4, figsize=(25, 7)) images = [img, img_segmented, pred_img, confidence_img] for ax, im in zip(axs, images, strict=False): ax.imshow(im) ax.axis("off") plt.subplots_adjust(left=0.01, right=0.99, top=0.99, bottom=0.01, wspace=0.05) plt.show() .. image-sg:: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_003.png :alt: tutorial muad mc drop :srcset: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 287-291 10. Apply the MC Dropout wrapper ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ This technique decribed in this `paper `_ allow us to have a better confidence score by using the dropout during test time. .. GENERATED FROM PYTHON SOURCE LINES 291-307 .. code-block:: Python # We wrap the model with the MC Dropout wrapper from torch_uncertainty.models.mc_dropout num_estimators = 10 mc_model = mc_dropout( model, num_estimators=num_estimators, last_layer=False, # We do not want to apply dropout on the last layer on_batch=False, # To reduce memory usage, we execute the forward passes sequentially ) seg_routine = SegmentationRoutine( model=mc_model, num_classes=datamodule.num_classes, loss=None, # No loss needed for testing ) .. GENERATED FROM PYTHON SOURCE LINES 308-310 11. Testing the MC Dropout model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 310-312 .. code-block:: Python results = trainer.test(model=seg_routine, datamodule=datamodule) .. rst-class:: sphx-glr-script-out .. code-block:: none ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Segmentation ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ Brier │ 0.20520 │ │ NLL │ 0.40327 │ │ mAcc │ 68.400% │ │ mIoU │ 50.871% │ │ pixAcc │ 85.739% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Calibration ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ ECE │ 2.792% │ │ aECE │ 2.792% │ └──────────────┴───────────────────────────┘ ┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ Selective Classification ┃ ┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ AUGRC │ 2.680% │ │ AURC │ 3.288% │ │ Cov@5Risk │ 72.726% │ │ Risk@80Cov │ 6.948% │ └──────────────┴───────────────────────────┘ .. GENERATED FROM PYTHON SOURCE LINES 313-317 9. Uncertainty evaluations with MCP ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Here the confidence score is the Maximum class probability on the averaged probabilities of the different estimators. .. GENERATED FROM PYTHON SOURCE LINES 317-365 .. code-block:: Python sample_idx = 0 img, target = datamodule.test[sample_idx] batch_img = img.unsqueeze(0) batch_target = target.unsqueeze(0) mc_model.eval() with torch.no_grad(): # Forward propagation output_probs_per_est = mc_model(batch_img).softmax(dim=1) output_probs = output_probs_per_est.mean(0) # Average over the estimators # remove the batch dimension confidence, pred = output_probs.max(0) # Undo normalization on the image and convert to uint8. img = img * t_muad_std[:, None, None] + t_muad_mean[:, None, None] img = F.to_dtype(img, torch.uint8, scale=True) tmp_target = target.masked_fill(target == 255, 21) target_masks = tmp_target == torch.arange(22, device=target.device)[:, None, None] img_segmented = draw_segmentation_masks( img, target_masks, alpha=1, colors=datamodule.test.color_palette ) pred_masks = pred == torch.arange(22, device=pred.device)[:, None, None] pred_img = draw_segmentation_masks(img, pred_masks, alpha=1, colors=datamodule.test.color_palette) if confidence.ndim == 2: confidence = confidence.unsqueeze(0) img = F.to_pil_image(F.resize(img, 1024)) img_segmented = F.to_pil_image(F.resize(img_segmented, 1024)) pred_img = F.to_pil_image(F.resize(pred_img, 1024)) confidence_img = F.to_pil_image(F.resize(confidence, 1024)) fig, axs = plt.subplots(1, 4, figsize=(25, 7)) images = [img, img_segmented, pred_img, confidence_img] for ax, im in zip(axs, images, strict=False): ax.imshow(im) ax.axis("off") plt.subplots_adjust(left=0.01, right=0.99, top=0.99, bottom=0.01, wspace=0.05) plt.show() .. image-sg:: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_004.png :alt: tutorial muad mc drop :srcset: /auto_tutorials/Bayesian_Methods/images/sphx_glr_tutorial_muad_mc_drop_004.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 54.962 seconds) .. _sphx_glr_download_auto_tutorials_Bayesian_Methods_tutorial_muad_mc_drop.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: tutorial_muad_mc_drop.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: tutorial_muad_mc_drop.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: tutorial_muad_mc_drop.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_