Masksembles for Uncertainty Estimation

Project Page | Paper | Video Explanation

Official implementation of Masksembles approach from the paper "Masksembles for Uncertainty Estimation" by Nikita Durasov, Timur Bagautdinov, Pierre Baque, Pascal Fua. In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)

Installation

To install this package, use:

pip install git+http://github.com/nikitadurasov/masksembles

In addition, Masksembles requires installing at least one of the backends: torch or tensorflow2 / keras. Please follow official installation instructions for torch or tensorflow accordingly.

Usage

This package provides implementations for Masksembles{1|2|3}D layers in masksembles.{torch|keras} where {1|2|3} refers to dimensionality of input tensors (1-, 2- and 3-dimensional accordingly).

Masksembles1D: works with 1-dim inputs,[B, C] shaped tensors
Masksembles2D: works with 2-dim inputs,[B, H, W, C] (keras) or [B, C, H, W] (torch) shaped tensors
Masksembles3D : TBD

In a Nutshell, Masksembles applies binary masks to inputs via multiplying them both channel-wise. For more efficient implementation we've followed approach similar to this one. Therefore, after inference outputs[:B // N] - stores results for the first submodel, outputs[B // N : 2 * B // N] - for the second and etc.

Torch

import torch
from masksembles.torch import Masksembles1D

layer = Masksembles1D(10, 4, 2.)
layer(torch.ones([4, 10]))

tensor([[0., 1., 0., 0., 1., 0., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1., 0., 0., 1., 1.],
        [1., 0., 1., 1., 0., 0., 1., 0., 1., 1.],
        [1., 0., 0., 1., 1., 1., 0., 1., 1., 0.]], dtype=torch.float64)

Tensorflow / Keras

import tensorflow as tf 
from masksembles.keras import Masksembles1D

layer = Masksembles1D(4, 2.)
layer(tf.ones([4, 10]))

<tf.Tensor: shape=(4, 10), dtype=float32, numpy=
array([[0., 1., 1., 0., 1., 1., 1., 0., 1., 0.],
       [0., 1., 0., 1., 1., 0., 1., 1., 0., 1.],
       [1., 1., 1., 1., 0., 0., 1., 0., 0., 1.],
       [1., 0., 0., 1., 0., 1., 1., 0., 1., 1.]], dtype=float32)>

Model example

import tensorflow as tf 
from masksembles.keras import Masksembles1D, Masksembles2D

model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="elu"),
        Masksembles2D(4, 2.0),
        layers.MaxPooling2D(pool_size=(2, 2)),
     
        layers.Conv2D(64, kernel_size=(3, 3), activation="elu"),
        Masksembles2D(4, 2.0),
        layers.MaxPooling2D(pool_size=(2, 2)),
     
        layers.Flatten(),
        Masksembles1D(4, 2.),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

dkoguciuk / masksembles