mc-dropout-pytorch 0.1.0__tar.gz

mc_dropout_pytorch-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mc_dropout_pytorch-0.1.0/PKG-INFO

@@ -0,0 +1,189 @@
+Metadata-Version: 2.4
+Name: mc-dropout-pytorch
+Version: 0.1.0
+Summary: MC Dropout (Gal & Ghahramani, 2016) - Pytorch
+Home-page: https://github.com/lucidrains/mc-dropout-pytorch
+Author: lucidrains
+Author-email: lucidrains@gmail.com
+License: MIT
+Keywords: artificial intelligence,deep learning,bayesian deep learning,uncertainty estimation,monte carlo dropout
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: accelerate
+Requires-Dist: einops>=0.7
+Requires-Dist: ema-pytorch>=0.4.2
+Requires-Dist: torch>=2.0
+Requires-Dist: tqdm
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
+
+
+
+## MC Dropout, in Pytorch
+
+[![PyPI version](https://badge.fury.io/py/mc-dropout-pytorch.svg)](https://badge.fury.io/py/mc-dropout-pytorch)
+
+Implementation of [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://arxiv.org/abs/1506.02142) (Gal & Ghahramani, ICML 2016) in Pytorch.
+
+Standard dropout NNs cast as approximate Bayesian inference over deep Gaussian processes — giving free, calibrated uncertainty estimates with no architectural changes and zero inference overhead beyond T forward passes.
+
+## Install
+
+```bash
+$ pip install mc-dropout-pytorch
+```
+
+## Usage
+
+### Regression with uncertainty
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, MCDropoutInference, Trainer
+
+# build model
+model = BayesianMLP(
+    input_dim    = 1,
+    output_dim   = 1,
+    hidden_dims  = (256, 256),
+    dropout_rate = 0.1,
+    activation   = 'relu',
+)
+
+# wrap for MC inference (T=50 stochastic passes)
+mc = MCDropoutInference(model, num_samples = 50, task = 'regression', tau = 1.0)
+
+x = torch.linspace(-3, 3, 100).unsqueeze(-1)
+out = mc(x)
+
+out.mean      # predictive mean     — (100, 1)
+out.variance  # predictive variance — (100, 1)  includes τ⁻¹ noise term
+out.samples   # raw samples         — (50, 100, 1)
+```
+
+### Classification with predictive entropy
+
+```python
+import torch
+from mc_dropout_pytorch import BayesianCNN, MCDropoutInference
+
+model = BayesianCNN(
+    in_channels      = 1,
+    num_classes      = 10,
+    base_channels    = 32,
+    dropout_rate     = 0.25,
+    fc_dropout_rate  = 0.5,
+    img_size         = 28,
+)
+
+mc = MCDropoutInference(model, num_samples = 50, task = 'classification')
+
+x = torch.randn(8, 1, 28, 28)
+out = mc(x)
+
+out.mean      # class probabilities  — (8, 10)
+out.variance  # per-class variance   — (8, 10)
+
+# active learning signals (§6)
+H  = mc.predictive_entropy(x)   # (8,)  — total uncertainty
+MI = mc.mutual_information(x)   # (8,)  — epistemic uncertainty only
+```
+
+### Full training loop with the `Trainer`
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, Trainer
+
+# synthetic regression dataset
+X = torch.randn(1000, 4)
+y = X[:, 0] * 2 + X[:, 1] - X[:, 2] + torch.randn(1000) * 0.1
+dataset = TensorDataset(X, y)
+
+model = BayesianMLP(
+    input_dim    = 4,
+    output_dim   = 1,
+    hidden_dims  = (128, 128),
+    dropout_rate = 0.1,
+)
+
+trainer = Trainer(
+    model,
+    dataset,
+    task             = 'regression',
+    train_lr         = 1e-3,
+    train_num_steps  = 5_000,
+    train_batch_size = 64,
+    ema_decay        = 0.995,
+    amp              = False,
+    weight_decay     = 1e-4,   # ≡ prior precision in §3
+    tau              = 1.0,    # noise precision
+    num_mc_samples   = 50,
+)
+
+trainer.train()
+
+# inference via EMA model
+mc = trainer.inference
+out = mc(X[:10])
+print(out.mean, out.variance)
+```
+
+### Multi-GPU
+
+```bash
+$ accelerate config
+$ accelerate launch train.py
+```
+
+## Key ideas from the paper
+
+**The insight (§3)**: Training a NN with dropout and L2 regularisation minimises a KL divergence to the posterior of a deep Gaussian process — no variational EM, no weight sampling required.
+
+**Test-time dropout (MC Dropout)**:
+
+```
+for t = 1 … T:
+    ŷ_t = f^ω_t(x)    # ω_t ~ q(ω)  via Bernoulli dropout
+
+E[y*]   ≈ (1/T) Σ ŷ_t                              # predictive mean
+Var[y*] ≈ τ⁻¹ I + (1/T) Σ ŷ_t ŷ_tᵀ − E[y*]²     # predictive variance  (Eq. 9)
+```
+
+**Active learning** (§6): Use `mc.mutual_information(x)` to identify the most informative unlabelled points — pure epistemic uncertainty, disentangled from aleatoric noise.
+
+**Weight correspondence** (§3.2):
+
+| Dropout training          | Bayesian GP posterior    |
+|---------------------------|--------------------------|
+| dropout probability `p`   | variational parameter    |
+| L2 weight decay `λ`       | prior precision          |
+| noise precision `τ`       | `τ = (2N λ) / (1 − p)`  |
+
+## Citations
+
+```bibtex
+@article{Gal2016Dropout,
+    title   = {Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning},
+    author  = {Yarin Gal and Zoubin Ghahramani},
+    journal = {Proceedings of the 33rd International Conference on Machine Learning (ICML)},
+    year    = {2016},
+    url     = {https://arxiv.org/abs/1506.02142}
+}
+```

mc_dropout_pytorch-0.1.0/README.md

@@ -0,0 +1,156 @@
+
+
+## MC Dropout, in Pytorch
+
+[![PyPI version](https://badge.fury.io/py/mc-dropout-pytorch.svg)](https://badge.fury.io/py/mc-dropout-pytorch)
+
+Implementation of [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://arxiv.org/abs/1506.02142) (Gal & Ghahramani, ICML 2016) in Pytorch.
+
+Standard dropout NNs cast as approximate Bayesian inference over deep Gaussian processes — giving free, calibrated uncertainty estimates with no architectural changes and zero inference overhead beyond T forward passes.
+
+## Install
+
+```bash
+$ pip install mc-dropout-pytorch
+```
+
+## Usage
+
+### Regression with uncertainty
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, MCDropoutInference, Trainer
+
+# build model
+model = BayesianMLP(
+    input_dim    = 1,
+    output_dim   = 1,
+    hidden_dims  = (256, 256),
+    dropout_rate = 0.1,
+    activation   = 'relu',
+)
+
+# wrap for MC inference (T=50 stochastic passes)
+mc = MCDropoutInference(model, num_samples = 50, task = 'regression', tau = 1.0)
+
+x = torch.linspace(-3, 3, 100).unsqueeze(-1)
+out = mc(x)
+
+out.mean      # predictive mean     — (100, 1)
+out.variance  # predictive variance — (100, 1)  includes τ⁻¹ noise term
+out.samples   # raw samples         — (50, 100, 1)
+```
+
+### Classification with predictive entropy
+
+```python
+import torch
+from mc_dropout_pytorch import BayesianCNN, MCDropoutInference
+
+model = BayesianCNN(
+    in_channels      = 1,
+    num_classes      = 10,
+    base_channels    = 32,
+    dropout_rate     = 0.25,
+    fc_dropout_rate  = 0.5,
+    img_size         = 28,
+)
+
+mc = MCDropoutInference(model, num_samples = 50, task = 'classification')
+
+x = torch.randn(8, 1, 28, 28)
+out = mc(x)
+
+out.mean      # class probabilities  — (8, 10)
+out.variance  # per-class variance   — (8, 10)
+
+# active learning signals (§6)
+H  = mc.predictive_entropy(x)   # (8,)  — total uncertainty
+MI = mc.mutual_information(x)   # (8,)  — epistemic uncertainty only
+```
+
+### Full training loop with the `Trainer`
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, Trainer
+
+# synthetic regression dataset
+X = torch.randn(1000, 4)
+y = X[:, 0] * 2 + X[:, 1] - X[:, 2] + torch.randn(1000) * 0.1
+dataset = TensorDataset(X, y)
+
+model = BayesianMLP(
+    input_dim    = 4,
+    output_dim   = 1,
+    hidden_dims  = (128, 128),
+    dropout_rate = 0.1,
+)
+
+trainer = Trainer(
+    model,
+    dataset,
+    task             = 'regression',
+    train_lr         = 1e-3,
+    train_num_steps  = 5_000,
+    train_batch_size = 64,
+    ema_decay        = 0.995,
+    amp              = False,
+    weight_decay     = 1e-4,   # ≡ prior precision in §3
+    tau              = 1.0,    # noise precision
+    num_mc_samples   = 50,
+)
+
+trainer.train()
+
+# inference via EMA model
+mc = trainer.inference
+out = mc(X[:10])
+print(out.mean, out.variance)
+```
+
+### Multi-GPU
+
+```bash
+$ accelerate config
+$ accelerate launch train.py
+```
+
+## Key ideas from the paper
+
+**The insight (§3)**: Training a NN with dropout and L2 regularisation minimises a KL divergence to the posterior of a deep Gaussian process — no variational EM, no weight sampling required.
+
+**Test-time dropout (MC Dropout)**:
+
+```
+for t = 1 … T:
+    ŷ_t = f^ω_t(x)    # ω_t ~ q(ω)  via Bernoulli dropout
+
+E[y*]   ≈ (1/T) Σ ŷ_t                              # predictive mean
+Var[y*] ≈ τ⁻¹ I + (1/T) Σ ŷ_t ŷ_tᵀ − E[y*]²     # predictive variance  (Eq. 9)
+```
+
+**Active learning** (§6): Use `mc.mutual_information(x)` to identify the most informative unlabelled points — pure epistemic uncertainty, disentangled from aleatoric noise.
+
+**Weight correspondence** (§3.2):
+
+| Dropout training          | Bayesian GP posterior    |
+|---------------------------|--------------------------|
+| dropout probability `p`   | variational parameter    |
+| L2 weight decay `λ`       | prior precision          |
+| noise precision `τ`       | `τ = (2N λ) / (1 − p)`  |
+
+## Citations
+
+```bibtex
+@article{Gal2016Dropout,
+    title   = {Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning},
+    author  = {Yarin Gal and Zoubin Ghahramani},
+    journal = {Proceedings of the 33rd International Conference on Machine Learning (ICML)},
+    year    = {2016},
+    url     = {https://arxiv.org/abs/1506.02142}
+}
+```

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch/__init__.py

@@ -0,0 +1,8 @@
+from mc_dropout_pytorch.mc_dropout_pytorch import (
+    MCDropout,
+    MCDropout2d,
+    BayesianMLP,
+    BayesianCNN,
+    MCDropoutInference,
+    Trainer,
+)

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch/mc_dropout_pytorch.py

@@ -0,0 +1,461 @@
+import math
+from pathlib import Path
+from functools import partial
+from collections import namedtuple
+from multiprocessing import cpu_count
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, TensorDataset
+
+from torch.optim import Adam
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+from einops import rearrange, reduce, repeat
+from einops.layers.torch import Rearrange
+
+from tqdm.auto import tqdm
+from ema_pytorch import EMA
+from accelerate import Accelerator
+
+# ──────────────────────────────────────────────
+# constants
+# ──────────────────────────────────────────────
+
+ModelOutput = namedtuple('ModelOutput', ['mean', 'variance', 'samples'])
+
+# ──────────────────────────────────────────────
+# helpers
+# ──────────────────────────────────────────────
+
+def exists(x):
+    return x is not None
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+def identity(t, *args, **kwargs):
+    return t
+
+def cycle(dl):
+    while True:
+        for data in dl:
+            yield data
+
+def cast_tuple(t, length = 1):
+    if isinstance(t, tuple):
+        return t
+    return ((t,) * length)
+
+def divisible_by(numer, denom):
+    return (numer % denom) == 0
+
+def num_to_groups(num, divisor):
+    groups = num // divisor
+    remainder = num % divisor
+    arr = [divisor] * groups
+    if remainder > 0:
+        arr.append(remainder)
+    return arr
+
+# ──────────────────────────────────────────────
+# MC Dropout core: enable dropout at test time
+# ──────────────────────────────────────────────
+
+class MCDropout(nn.Dropout):
+    """
+    MC Dropout layer (Gal & Ghahramani, 2016).
+
+    Standard nn.Dropout is disabled at eval time (model.eval()).
+    This subclass keeps dropout active regardless of training mode,
+    so T stochastic forward passes give a Monte Carlo estimate of
+    the posterior predictive distribution.
+
+    Eq. (8) in the paper: T forward passes → {ŷ_t}  (t=1..T)
+        predictive mean   : E[y*] ≈ (1/T) Σ ŷ_t
+        predictive var    : Var[y*] ≈ τ⁻¹ I + (1/T) Σ ŷ_t ŷ_tᵀ − E[y*]²
+    """
+
+    def forward(self, x):
+        # keep p active even in eval mode  ← key contribution of the paper
+        return F.dropout(x, self.p, training = True, inplace = self.inplace)
+
+
+class MCDropout2d(nn.Dropout2d):
+    """Spatial MC Dropout for convolutional feature maps."""
+
+    def forward(self, x):
+        return F.dropout2d(x, self.p, training = True, inplace = self.inplace)
+
+# ──────────────────────────────────────────────
+# small helper modules
+# ──────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim = -1) * self.g * self.scale
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult = 4, dropout = 0.0):
+        super().__init__()
+        inner = int(dim * mult)
+        self.net = nn.Sequential(
+            RMSNorm(dim),
+            nn.Linear(dim, inner),
+            nn.GELU(),
+            MCDropout(dropout),
+            nn.Linear(inner, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+# ──────────────────────────────────────────────
+# Bayesian MLP — regression & classification
+# ──────────────────────────────────────────────
+
+class BayesianMLP(nn.Module):
+    """
+    Dropout-regularised MLP whose test-time stochastic forward passes
+    approximate a deep Gaussian process posterior (Gal & Ghahramani, §3).
+
+    Parameters
+    ----------
+    input_dim    : int   — feature dimension
+    output_dim   : int   — number of targets / classes
+    hidden_dims  : tuple — widths of hidden layers, default (256, 256)
+    dropout_rate : float — Bernoulli dropout probability p (typically 0.1-0.5)
+    activation   : str   — 'relu' | 'tanh' | 'gelu'  (§5 ablation)
+    """
+
+    activation_map = {
+        'relu' : nn.ReLU,
+        'tanh' : nn.Tanh,
+        'gelu' : nn.GELU,
+    }
+
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        *,
+        hidden_dims = (256, 256),
+        dropout_rate = 0.1,
+        activation = 'relu',
+    ):
+        super().__init__()
+
+        act_cls = self.activation_map.get(activation, nn.ReLU)
+        dims = (input_dim, *hidden_dims)
+
+        layers = []
+        for d_in, d_out in zip(dims[:-1], dims[1:]):
+            layers += [
+                nn.Linear(d_in, d_out),
+                act_cls(),
+                MCDropout(dropout_rate),
+            ]
+        layers.append(nn.Linear(hidden_dims[-1], output_dim))
+
+        self.net = nn.Sequential(*layers)
+        self.dropout_rate = dropout_rate
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# ──────────────────────────────────────────────
+# Bayesian CNN — for image classification (§5)
+# ──────────────────────────────────────────────
+
+class BayesianCNN(nn.Module):
+    """
+    Convolutional network with MC Dropout after every conv block,
+    matching the MNIST architecture described in §5 of the paper.
+
+    Architecture: Conv → ReLU → MCDrop2d → Conv → ReLU → MCDrop2d
+                  → Flatten → Linear → ReLU → MCDrop → Linear
+    """
+
+    def __init__(
+        self,
+        in_channels = 1,
+        num_classes = 10,
+        *,
+        base_channels = 32,
+        dropout_rate = 0.25,
+        fc_dropout_rate = 0.5,
+        img_size = 28,
+    ):
+        super().__init__()
+
+        c = base_channels
+        conv_out_size = (img_size // 4) ** 2 * (c * 2)
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, c, 3, padding = 1),
+            nn.ReLU(),
+            MCDropout2d(dropout_rate),
+            nn.Conv2d(c, c * 2, 3, padding = 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            MCDropout2d(dropout_rate),
+            nn.Conv2d(c * 2, c * 2, 3, padding = 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            MCDropout2d(dropout_rate),
+        )
+
+        self.head = nn.Sequential(
+            Rearrange('b c h w -> b (c h w)'),
+            nn.Linear(conv_out_size, 256),
+            nn.ReLU(),
+            MCDropout(fc_dropout_rate),
+            nn.Linear(256, num_classes),
+        )
+
+        self.dropout_rate = dropout_rate
+
+    def forward(self, x):
+        return self.head(self.conv(x))
+
+
+# ──────────────────────────────────────────────
+# MC Inference — the inference wrapper
+# ──────────────────────────────────────────────
+
+class MCDropoutInference(nn.Module):
+    """
+    Wraps any BayesianMLP / BayesianCNN to produce predictive
+    mean, variance and full sample tensor via T stochastic passes.
+
+    Predictive uncertainty decomposition (§3, Eq. 9):
+        τ⁻¹  — noise precision (regression length-scale term)
+        Var  — model (epistemic) uncertainty from T samples
+
+    Parameters
+    ----------
+    model        : nn.Module  — a BayesianMLP or BayesianCNN
+    num_samples  : int        — T in the paper (default 50)
+    task         : str        — 'regression' | 'classification'
+    tau          : float      — noise precision τ for regression uncertainty
+    """
+
+    def __init__(
+        self,
+        model,
+        *,
+        num_samples = 50,
+        task = 'regression',
+        tau = 1.0,
+    ):
+        super().__init__()
+        self.model = model
+        self.num_samples = num_samples
+        self.task = task
+        self.tau = tau
+
+    @torch.no_grad()
+    def forward(self, x):
+        # T stochastic forward passes — shape (T, B, output_dim)
+        samples = torch.stack(
+            [self.model(x) for _ in range(self.num_samples)],
+            dim = 0,
+        )
+
+        if self.task == 'classification':
+            # softmax each sample, then average → predictive probabilities
+            probs = samples.softmax(dim = -1)          # (T, B, C)
+            mean  = reduce(probs, 't b c -> b c', 'mean')
+            var   = reduce(probs ** 2, 't b c -> b c', 'mean') - mean ** 2
+        else:
+            # regression: Eq. (9) — add noise precision term
+            mean  = reduce(samples, 't b o -> b o', 'mean')
+            var   = reduce(samples ** 2, 't b o -> b o', 'mean') \
+                  - mean ** 2 \
+                  + (1.0 / self.tau)
+
+        return ModelOutput(mean = mean, variance = var, samples = samples)
+
+    def predictive_entropy(self, x):
+        """
+        H[y | x, X, Y] — used for active learning (§6 of paper).
+        High entropy → model is uncertain → good candidate to label.
+        """
+        out = self.forward(x)
+        # clip for numerical stability
+        p = out.mean.clamp(min = 1e-8)
+        return -(p * p.log()).sum(dim = -1)
+
+    def mutual_information(self, x):
+        """
+        I[y, ω | x, X, Y] — epistemic (model) uncertainty (§6).
+        MI = H[y|x] − E_ω[H[y|x,ω]]
+        """
+        out = self.forward(x)
+        # H of predictive mean
+        p_mean = out.mean.clamp(min = 1e-8)
+        h_mean = -(p_mean * p_mean.log()).sum(dim = -1)
+
+        # expected H over samples
+        probs = out.samples.softmax(dim = -1).clamp(min = 1e-8)   # (T, B, C)
+        h_samples = -(probs * probs.log()).sum(dim = -1)           # (T, B)
+        exp_h = reduce(h_samples, 't b -> b', 'mean')
+
+        return h_mean - exp_h
+
+
+# ──────────────────────────────────────────────
+# Trainer
+# ──────────────────────────────────────────────
+
+class Trainer:
+    """
+    Training wrapper with accelerate + EMA for MC Dropout models.
+
+    Supports both regression (MSE) and classification (cross-entropy).
+
+    Parameters
+    ----------
+    model           : BayesianMLP | BayesianCNN
+    dataset         : Dataset
+    task            : 'regression' | 'classification'
+    train_lr        : float  — learning rate (default 1e-3)
+    train_num_steps : int    — total gradient steps
+    train_batch_size: int    — batch size
+    ema_decay       : float  — EMA decay for weight averaging
+    amp             : bool   — mixed precision
+    results_folder  : str    — where to save checkpoints
+    num_mc_samples  : int    — T for inference object
+    weight_decay    : float  — L2 regularisation (≡ prior precision in §3)
+    """
+
+    def __init__(
+        self,
+        model,
+        dataset,
+        *,
+        task = 'regression',
+        train_lr = 1e-3,
+        train_num_steps = 10_000,
+        train_batch_size = 128,
+        ema_decay = 0.995,
+        amp = False,
+        results_folder = './results',
+        num_mc_samples = 50,
+        weight_decay = 1e-4,
+        tau = 1.0,
+        save_every = 1000,
+    ):
+        self.accelerator = Accelerator(mixed_precision = 'fp16' if amp else 'no')
+
+        self.model    = model
+        self.task     = task
+        self.tau      = tau
+
+        self.save_every = save_every
+        self.train_num_steps = train_num_steps
+
+        self.dl = cycle(DataLoader(
+            dataset,
+            batch_size = train_batch_size,
+            shuffle = True,
+            num_workers = min(4, cpu_count()),
+            pin_memory = True,
+        ))
+
+        self.opt = Adam(model.parameters(), lr = train_lr, weight_decay = weight_decay)
+
+        self.ema = EMA(model, beta = ema_decay, update_every = 10)
+
+        self.results_folder = Path(results_folder)
+        self.results_folder.mkdir(exist_ok = True)
+
+        self.model, self.opt, self.ema = self.accelerator.prepare(
+            self.model, self.opt, self.ema
+        )
+
+        # expose inference wrapper around EMA model
+        self.inference = MCDropoutInference(
+            self.ema.ema_model,
+            num_samples = num_mc_samples,
+            task = task,
+            tau = tau,
+        )
+
+        self.step = 0
+
+    def save(self, milestone):
+        data = {
+            'step'  : self.step,
+            'model' : self.accelerator.get_state_dict(self.model),
+            'opt'   : self.opt.state_dict(),
+            'ema'   : self.ema.state_dict(),
+        }
+        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
+
+    def load(self, milestone):
+        data = torch.load(
+            str(self.results_folder / f'model-{milestone}.pt'),
+            map_location = self.accelerator.device,
+        )
+        model = self.accelerator.unwrap_model(self.model)
+        model.load_state_dict(data['model'])
+        self.step = data['step']
+        self.opt.load_state_dict(data['opt'])
+        self.ema.load_state_dict(data['ema'])
+
+    def train(self):
+        accelerator = self.accelerator
+
+        with tqdm(
+            initial = self.step,
+            total   = self.train_num_steps,
+            disable = not accelerator.is_main_process,
+        ) as pbar:
+
+            while self.step < self.train_num_steps:
+                batch = next(self.dl)
+
+                # support (x,) or (x, y) datasets
+                if isinstance(batch, (list, tuple)) and len(batch) == 2:
+                    x, y = batch
+                else:
+                    x = batch[0] if isinstance(batch, (list, tuple)) else batch
+                    y = None
+
+                with self.accelerator.autocast():
+                    logits = self.model(x)
+
+                    if self.task == 'classification' and exists(y):
+                        loss = F.cross_entropy(logits, y.long())
+                    elif exists(y):
+                        # heteroscedastic MSE — Eq. (4) negative log-likelihood
+                        loss = F.mse_loss(logits.squeeze(-1), y.float())
+                    else:
+                        raise ValueError("Dataset must return (x, y) pairs")
+
+                self.accelerator.backward(loss)
+                self.opt.step()
+                self.opt.zero_grad()
+                self.ema.update()
+
+                pbar.set_description(f'loss: {loss.item():.4f}')
+                self.step += 1
+
+                if divisible_by(self.step, self.save_every):
+                    milestone = self.step // self.save_every
+                    self.save(milestone)
+
+                pbar.update(1)
+
+        accelerator.print('training complete')

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch/version.py

@@ -0,0 +1 @@
+__version__ = '0.1.0'

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch.egg-info/PKG-INFO

@@ -0,0 +1,189 @@
+Metadata-Version: 2.4
+Name: mc-dropout-pytorch
+Version: 0.1.0
+Summary: MC Dropout (Gal & Ghahramani, 2016) - Pytorch
+Home-page: https://github.com/lucidrains/mc-dropout-pytorch
+Author: lucidrains
+Author-email: lucidrains@gmail.com
+License: MIT
+Keywords: artificial intelligence,deep learning,bayesian deep learning,uncertainty estimation,monte carlo dropout
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: accelerate
+Requires-Dist: einops>=0.7
+Requires-Dist: ema-pytorch>=0.4.2
+Requires-Dist: torch>=2.0
+Requires-Dist: tqdm
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary
+
+
+
+## MC Dropout, in Pytorch
+
+[![PyPI version](https://badge.fury.io/py/mc-dropout-pytorch.svg)](https://badge.fury.io/py/mc-dropout-pytorch)
+
+Implementation of [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://arxiv.org/abs/1506.02142) (Gal & Ghahramani, ICML 2016) in Pytorch.
+
+Standard dropout NNs cast as approximate Bayesian inference over deep Gaussian processes — giving free, calibrated uncertainty estimates with no architectural changes and zero inference overhead beyond T forward passes.
+
+## Install
+
+```bash
+$ pip install mc-dropout-pytorch
+```
+
+## Usage
+
+### Regression with uncertainty
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, MCDropoutInference, Trainer
+
+# build model
+model = BayesianMLP(
+    input_dim    = 1,
+    output_dim   = 1,
+    hidden_dims  = (256, 256),
+    dropout_rate = 0.1,
+    activation   = 'relu',
+)
+
+# wrap for MC inference (T=50 stochastic passes)
+mc = MCDropoutInference(model, num_samples = 50, task = 'regression', tau = 1.0)
+
+x = torch.linspace(-3, 3, 100).unsqueeze(-1)
+out = mc(x)
+
+out.mean      # predictive mean     — (100, 1)
+out.variance  # predictive variance — (100, 1)  includes τ⁻¹ noise term
+out.samples   # raw samples         — (50, 100, 1)
+```
+
+### Classification with predictive entropy
+
+```python
+import torch
+from mc_dropout_pytorch import BayesianCNN, MCDropoutInference
+
+model = BayesianCNN(
+    in_channels      = 1,
+    num_classes      = 10,
+    base_channels    = 32,
+    dropout_rate     = 0.25,
+    fc_dropout_rate  = 0.5,
+    img_size         = 28,
+)
+
+mc = MCDropoutInference(model, num_samples = 50, task = 'classification')
+
+x = torch.randn(8, 1, 28, 28)
+out = mc(x)
+
+out.mean      # class probabilities  — (8, 10)
+out.variance  # per-class variance   — (8, 10)
+
+# active learning signals (§6)
+H  = mc.predictive_entropy(x)   # (8,)  — total uncertainty
+MI = mc.mutual_information(x)   # (8,)  — epistemic uncertainty only
+```
+
+### Full training loop with the `Trainer`
+
+```python
+import torch
+from torch.utils.data import TensorDataset
+from mc_dropout_pytorch import BayesianMLP, Trainer
+
+# synthetic regression dataset
+X = torch.randn(1000, 4)
+y = X[:, 0] * 2 + X[:, 1] - X[:, 2] + torch.randn(1000) * 0.1
+dataset = TensorDataset(X, y)
+
+model = BayesianMLP(
+    input_dim    = 4,
+    output_dim   = 1,
+    hidden_dims  = (128, 128),
+    dropout_rate = 0.1,
+)
+
+trainer = Trainer(
+    model,
+    dataset,
+    task             = 'regression',
+    train_lr         = 1e-3,
+    train_num_steps  = 5_000,
+    train_batch_size = 64,
+    ema_decay        = 0.995,
+    amp              = False,
+    weight_decay     = 1e-4,   # ≡ prior precision in §3
+    tau              = 1.0,    # noise precision
+    num_mc_samples   = 50,
+)
+
+trainer.train()
+
+# inference via EMA model
+mc = trainer.inference
+out = mc(X[:10])
+print(out.mean, out.variance)
+```
+
+### Multi-GPU
+
+```bash
+$ accelerate config
+$ accelerate launch train.py
+```
+
+## Key ideas from the paper
+
+**The insight (§3)**: Training a NN with dropout and L2 regularisation minimises a KL divergence to the posterior of a deep Gaussian process — no variational EM, no weight sampling required.
+
+**Test-time dropout (MC Dropout)**:
+
+```
+for t = 1 … T:
+    ŷ_t = f^ω_t(x)    # ω_t ~ q(ω)  via Bernoulli dropout
+
+E[y*]   ≈ (1/T) Σ ŷ_t                              # predictive mean
+Var[y*] ≈ τ⁻¹ I + (1/T) Σ ŷ_t ŷ_tᵀ − E[y*]²     # predictive variance  (Eq. 9)
+```
+
+**Active learning** (§6): Use `mc.mutual_information(x)` to identify the most informative unlabelled points — pure epistemic uncertainty, disentangled from aleatoric noise.
+
+**Weight correspondence** (§3.2):
+
+| Dropout training          | Bayesian GP posterior    |
+|---------------------------|--------------------------|
+| dropout probability `p`   | variational parameter    |
+| L2 weight decay `λ`       | prior precision          |
+| noise precision `τ`       | `τ = (2N λ) / (1 − p)`  |
+
+## Citations
+
+```bibtex
+@article{Gal2016Dropout,
+    title   = {Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning},
+    author  = {Yarin Gal and Zoubin Ghahramani},
+    journal = {Proceedings of the 33rd International Conference on Machine Learning (ICML)},
+    year    = {2016},
+    url     = {https://arxiv.org/abs/1506.02142}
+}
+```

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+setup.py
+mc_dropout_pytorch/__init__.py
+mc_dropout_pytorch/mc_dropout_pytorch.py
+mc_dropout_pytorch/version.py
+mc_dropout_pytorch.egg-info/PKG-INFO
+mc_dropout_pytorch.egg-info/SOURCES.txt
+mc_dropout_pytorch.egg-info/dependency_links.txt
+mc_dropout_pytorch.egg-info/requires.txt
+mc_dropout_pytorch.egg-info/top_level.txt

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch.egg-info/requires.txt

@@ -0,0 +1,5 @@
+accelerate
+einops>=0.7
+ema-pytorch>=0.4.2
+torch>=2.0
+tqdm

mc_dropout_pytorch-0.1.0/mc_dropout_pytorch.egg-info/top_level.txt

@@ -0,0 +1 @@
+mc_dropout_pytorch

mc_dropout_pytorch-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mc_dropout_pytorch-0.1.0/setup.py

@@ -0,0 +1,37 @@
+from setuptools import setup, find_packages
+
+exec(open('mc_dropout_pytorch/version.py').read())
+
+setup(
+    name = 'mc-dropout-pytorch',
+    packages = find_packages(),
+    version = __version__,
+    license = 'MIT',
+    description = 'MC Dropout (Gal & Ghahramani, 2016) - Pytorch',
+    long_description = open('README.md').read(),
+    long_description_content_type = 'text/markdown',
+    author = 'lucidrains',
+    author_email = 'lucidrains@gmail.com',
+    url = 'https://github.com/lucidrains/mc-dropout-pytorch',
+    keywords = [
+        'artificial intelligence',
+        'deep learning',
+        'bayesian deep learning',
+        'uncertainty estimation',
+        'monte carlo dropout',
+    ],
+    install_requires = [
+        'accelerate',
+        'einops>=0.7',
+        'ema-pytorch>=0.4.2',
+        'torch>=2.0',
+        'tqdm',
+    ],
+    classifiers = [
+        'Development Status :: 4 - Beta',
+        'Intended Audience :: Developers',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'License :: OSI Approved :: MIT License',
+        'Programming Language :: Python :: 3.6',
+    ],
+)