Partial Monte-Carlo dropout

[pieterblok]

Monte-Carlo dropout is a simple but time-consuming procedure when determining the uncertainty score of an image. To alleviate this problem, I have implemented the so-called Partial Monte-Carlo dropout (PMCD) procedure in Detectron2 Active Learning (see line 173 in https://github.com/pieterblok/maskal/blob/maskAL/active_learning/sampling/montecarlo_dropout.py).

With PMCD an image is inferred only once through the neural network parts that don't have dropout. Then, the resulting features from these parts are copy-paste to the network parts with dropout where the actual repeated inference happens. PMCD is highly scalable and significantly faster than inferring the same image multiple times through the entire network. I was wondering if I can implement this PMCD with one of your developers in BAAL. @Dref360 @parmidaatg

[parmidaatg]

Hi @pieterblok, thank you so much for your comment and yes definitely please go ahead and make a PR. If you were interested you can then maybe generate some metrics to compare the two using BaaL and we post it in our tutorial. Otherwise I can do that part :)

[Dref360]

cc @orangetoaster

That sounds great!

[pieterblok]

@parmidaatg @Dref360 @orangetoaster

Perfect! I will have a deep dive into the BAAL source code to look how I can integrate such method. If I'm stuck somewhere or lost in integration, I will ask in this Discussion page.

[Dref360]

@pieterblok and I met on Friday to discuss this issue.
He proposed a great idea where users would implement a custom function that would have the following signature:

def montecarlo_forward(x : Tensor[BatchSize, ...], iterations:int) -> Tensor[BatchSize, ..., ITERATION]:
pass

This would make the case where we have "replicate_in_memory=False" much faster. Pieter shared a speedup of 20 to 40% on his models.

We also discussed another way of doing it with a LRU cache.

This PR is the first implementation of such proposal: #254

Let me know what you think!

[Dref360]

Hello! I'm finally back from Italy and I'm determined to figure this out.

Could you tell me what is your Pytorch/CUDA version? I suspect the torch.equal has been optimized recently. I personally have CUDA 11.7 and pytorch 2.0. I have a GTX 1060 on my laptop.

Also, I have a new test that is more "real-life like".

from torchvision.datasets import CIFAR10
from torchvision.models import vgg16
from torchvision.transforms import ToTensor

from baal.bayesian.caching_utils import MCCachingModule
from baal.bayesian.dropout import MCDropoutModule
from baal.modelwrapper import ModelWrapper

vgg = vgg16().cuda()
vgg.eval()

ds = CIFAR10('/tmp', train=False, transform=ToTensor(), download=True)

# Regular: 2:58
# Cached : 50 seconds
with MCCachingModule(vgg) as model:
    with MCDropoutModule(model) as model_2:
        wrapper = ModelWrapper(model_2, None, replicate_in_memory=False)
        wrapper.predict_on_dataset(ds, batch_size=32, iterations=20, use_cuda=True)

with MCDropoutModule(vgg) as model_2:
    wrapper = ModelWrapper(model_2, None, replicate_in_memory=False)
    wrapper.predict_on_dataset(ds, batch_size=32, iterations=20, use_cuda=True)

[pieterblok]

Hello! I'm finally back from Italy and I'm determined to figure this out.

Could you tell me what is your Pytorch/CUDA version? I suspect the torch.equal has been optimized recently. I personally have CUDA 11.7 and pytorch 2.0. I have a GTX 1060 on my laptop.

Also, I have a new test that is more "real-life like".

from torchvision.datasets import CIFAR10
from torchvision.models import vgg16
from torchvision.transforms import ToTensor

from baal.bayesian.caching_utils import MCCachingModule
from baal.bayesian.dropout import MCDropoutModule
from baal.modelwrapper import ModelWrapper

vgg = vgg16().cuda()
vgg.eval()

ds = CIFAR10('/tmp', train=False, transform=ToTensor(), download=True)

# Regular: 2:58
# Cached : 50 seconds
with MCCachingModule(vgg) as model:
    with MCDropoutModule(model) as model_2:
        wrapper = ModelWrapper(model_2, None, replicate_in_memory=False)
        wrapper.predict_on_dataset(ds, batch_size=32, iterations=20, use_cuda=True)

with MCDropoutModule(vgg) as model_2:
    wrapper = ModelWrapper(model_2, None, replicate_in_memory=False)
    wrapper.predict_on_dataset(ds, batch_size=32, iterations=20, use_cuda=True)

@Dref360 welcome back, and thanks for the update!

I'm using Pytorch version 1.8.0 with CUDA 10.1. Compared to your versions pretty outdated. I will try to get my laptop or another computer updated, so that I can replicate your experiment. I will immediately try our active learning with this cached approach, that might take some more time.

[pieterblok]

@Dref360 fortunately my colleague had already a system with Pytorch 2.0.1 and CUDA 11.7 installed, so it was easy to replicate this test. We have found the following:

"The MCCachingModule is indeed faster, approximately twice as fast. The total it/s of MCCachingModule is 20.72 while the one without caching was 10.83."

Next step is testing the prediction speed in our specific active learning application.

So far, so good... happy that we were able to reproduce your code with the newest libraries!

[Dref360]

@pieterblok I was able to run some tests on a UNET. The results are less impressive. I'm able to improve the results by 45% using the following script.

from typing import List

import segmentation_models_pytorch as smp
import torch
from segmentation_models_pytorch.base import SegmentationHead
from tqdm import tqdm

from baal.bayesian.caching_utils import MCCachingModule
from baal.bayesian.dropout import MCDropoutModule


DS_SIZE=100
ITERATIONS=20

def add_dropout(
        model: smp.Unet,
        decoder_channels: List[int] = (256, 128, 64, 32, 16),
        classes=1,
        activation=None,
):
    seg_head = SegmentationHead(
        in_channels=decoder_channels[-1],
        out_channels=classes,
        activation=activation,
        kernel_size=3,
    )
    model.add_module("segmentation_head", seg_head)
    model.initialize()


unet = smp.Unet(
    encoder_name="resnext50_32x4d",
    encoder_depth=5,
    encoder_weights="imagenet",
    decoder_use_batchnorm=False,
    classes=10,
)
# Add a Dropout layerto use MC-Dropout
add_dropout(unet, classes=10, activation=None)
unet = unet.eval().cuda()

# Regular: 1:30
# Cached: 52 seconds
with MCCachingModule(unet) as model:
    with MCDropoutModule(model) as model_2:
        for batch in tqdm(range(DS_SIZE)):
            input = torch.randn(10, 3, 224, 224).cuda()
            [model_2(input).detach().cpu() for _ in range(ITERATIONS)]

My hypothesis is that the new MCCachingModule doesn't cache concatenation hence the small speedup. This is something we could try to improve in the future. 50% is still a nice speedup.

I would propose merging MCCachingModule

[pieterblok]

@Dref360 perfect! Indeed that makes sense. Thanks for the effort.