✨ Docker settings, BatchEnsemble utilities rework + minor fixes

.github/workflows/build-docs.yml

@@ -4,6 +4,7 @@ on:
     branches:
       - main
   pull_request:
+    types: [opened, reopened, ready_for_review, synchronize]
     branches:
       - main
   schedule:
@@ -40,7 +41,7 @@ jobs:
 
     - name: Install dependencies
       run: |
-        python3 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+        python3 -m pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
         python3 -m pip install .[all]
 
     - name: Sphinx build

.github/workflows/run-tests.yml

@@ -7,7 +7,6 @@ on:
       - dev
   pull_request:
     branches:
-      - main
       - dev
   schedule:
     - cron: "42 7 * * 0"
@@ -64,7 +63,7 @@ jobs:
     - name: Install dependencies
       if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
-        python3 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+        python3 -m pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
         python3 -m pip install .[all]
 
     - name: Check style & format

.gitignore

@@ -1,5 +1,6 @@
 # Custom
 .vscode/
+.idea/
 data/
 logs/
 lightning_logs/
@@ -41,6 +42,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.setup_done
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

Dockerfile

@@ -0,0 +1,43 @@
+FROM pytorch/pytorch:2.6.0-cuda12.6-cudnn9-runtime
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    LC_ALL=C.UTF-8 \
+    LANG=C.UTF-8
+
+# Install Git, OpenSSH Server, and OpenGL (PyTorch's base image already includes Conda and Pip)
+RUN apt-get update && apt-get install -y \
+    git \
+    openssh-server \
+    libgl1 \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /workspace
+
+# Create an empty README.md file and an empty torch_uncertainty module to satisfy flit
+RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__init__.py
+
+# Copy dependency file
+COPY pyproject.toml /workspace/
+
+# Install dependencies all dependencies
+RUN pip install --no-cache-dir -e ".[all]"
+
+# Always activate Conda when opening a new terminal
+RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc
+
+# Configure SSH server
+RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \
+    echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config && \
+    echo "AuthorizedKeysFile .ssh/authorized_keys" >> /etc/ssh/sshd_config
+
+# Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH
+EXPOSE 8888 22
+
+# Entrypoint script (runs every time the container starts)
+COPY docker/entrypoint.sh /usr/local/bin/
+RUN chmod +x /usr/local/bin/entrypoint.sh
+ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
+
+# Note that if /workspace/ is a mounted volume, any files copied to /workspace/ during the build will be overwritten by the mounted volume
+# This is why we copy the entrypoint script to /usr/local/bin/ instead of /workspace/

README.md

@@ -43,6 +43,11 @@ pip install torch-uncertainty
 
 The installation procedure for contributors is different: have a look at the [contribution page](https://torch-uncertainty.github.io/contributing.html).
 
+### :whale: Docker image for contributors
+
+For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. 
+This allows you to quickly launch an experiment-ready container with minimal setup. Please refer to [DOCKER.md](docker/DOCKER.md) for further details.
+
 ## :racehorse: Quickstart
 
 We make a quickstart available at [torch-uncertainty.github.io/quickstart](https://torch-uncertainty.github.io/quickstart.html).
@@ -94,4 +99,4 @@ The following projects use TorchUncertainty:
 
 - *A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors* - [ICLR 2024](https://arxiv.org/abs/2310.08287)
 
-**If you are using TorchUncertainty in your project, please let us know, we will add your project to this list!**
+**If you are using TorchUncertainty in your project, please let us know, and we will add your project to this list!**

docker/DOCKER.md

@@ -0,0 +1,71 @@
+# :whale: Docker image for contributors
+
+### Pre-built Docker image
+1. To pull the pre-built image from Docker Hub, simply run:
+    ```bash
+    docker pull docker.io/tonyzamyatin/torch-uncertainty:latest
+    ```
+
+    This image includes:
+    - PyTorch with CUDA support
+    - OpenGL (for visualization tasks)
+    - Git, OpenSSH, and all Python dependencies
+
+    Checkout the [registry on Docker Hub](https://hub.docker.com/repository/docker/tonyzamyatin/torch-uncertainty/general) for all available images.
+
+2. To start a container using this image, set up the necessary environment variables and run:
+    ```bash
+    docker run --rm -it --gpus all -p 8888:8888 -p 22:22 \
+        -e VM_SSH_PUBLIC_KEY="your-public-key" \
+        -e GITHUB_SSH_PRIVATE_KEY="your-github-key" \
+        -e GITHUB_USER="your-github-username" \
+        -e GIT_USER_EMAIL="your-git-email" \
+        -e GIT_USER_NAME="your-git-name" \
+        docker.io/tonyzamyatin/torch-uncertainty
+    ```
+
+    Optionally, you can also set `-e USER_COMPACT_SHELL_PROMPT="true"`  
+    to make the VM's shell prompts compact and colorized.
+
+    **Note:** Some cloud providers offer templates, in which you can preconfigure  
+    in advance which Docker image to pull and which environment variables to set.  
+    In this case, the provider will pull the image, set all environment variables,  
+    and start the container for you.
+
+3. Once your cloud provider has deployed the VM, it will display the host address and SSH port.  
+    You can connect to the container via SSH using:
+    ```bash
+    ssh -i /path/to/private_key root@<VM_HOST> -p <VM_PORT>
+    ```
+
+    Replace `<VM_HOST>` and `<VM_PORT>` with the values provided by your cloud provider,  
+    and `/path/to/private_key` with the private key that corresponds to `VM_SSH_PUBLIC_KEY`.
+
+4. The container exposes port `8888` in case you want to run Jupyter Notebooks or TensorBoard.
+
+    **Note:** The `/workspace` directory is mounted from your local machine or cloud storage,  
+    so changes persist across container restarts.  
+    If using a cloud provider, ensure your network volume is correctly attached to avoid losing data.
+
+### Modifying and publishing custom Docker image
+
+If you want to make changes to the Dockerfile, follow these steps:
+1. Edit the Dockerfile to fit your needs.
+
+2. Build the modified image:
+    ```
+    docker build -t my-custom-image .
+    ```
+
+3. Push to a Docker registry (if you want to use it on another VM):
+    ```
+    docker tag my-custom-image mydockerhubuser/my-custom-image:tag
+    docker push mydockerhubuser/my-custom-image:tag
+    ```
+
+4. Pull the custom image onto your VM:
+    ```
+    docker pull mydockerhubuser/my-custom-image
+    ```
+
+5. Run the container using the same docker run command with the new image name.

docker/entrypoint.sh

@@ -0,0 +1,131 @@
+#!/bin/bash
+set -e  # Exit immediately if a command fails
+
+echo "🚀 Starting container..."
+
+# Ensure SSH directory exists and has correct permissions
+if [ ! -d /root/.ssh ]; then
+    echo "📂 Creating SSH directory..."
+    mkdir -p /root/.ssh && chmod 700 /root/.ssh
+fi
+
+# Ensure the VM's public SSH key is added for authentication
+if [ -z "$VM_SSH_PUBLIC_KEY" ]; then
+    echo "❌ Error: Please set the VM_SSH_PUBLIC_KEY environment variable."
+    exit 1
+fi
+if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q "$VM_SSH_PUBLIC_KEY" /root/.ssh/authorized_keys; then
+    echo "🔑 Adding VM SSH public key..."
+    echo "$VM_SSH_PUBLIC_KEY" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys
+fi
+
+# Ensure GitHub SSH private key is set up for authentication
+if [ -z "$GITHUB_SSH_PRIVATE_KEY" ]; then
+    echo "❌ Error: Please set the GITHUB_SSH_PRIVATE_KEY environment variable."
+    exit 1
+fi
+if [ ! -f /root/.ssh/github_rsa ]; then
+    echo "🔐 Adding GitHub SSH private key..."
+    echo "$GITHUB_SSH_PRIVATE_KEY" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa
+fi
+
+# Configure SSH client for GitHub authentication
+if [ ! -f /root/.ssh/config ] || ! grep -q 'Host github.com' /root/.ssh/config; then
+    echo "⚙️  Configuring SSH client for GitHub authentication..."
+    cat <<EOF > /root/.ssh/config
+Host github.com
+  User git
+  IdentityFile /root/.ssh/github_rsa
+EOF
+    chmod 600 /root/.ssh/config
+fi
+
+# Add GitHub to known hosts (to avoid SSH verification prompts)
+echo "📌 Ensuring GitHub is a known host..."
+ssh-keygen -F github.com > /dev/null 2>&1 || ssh-keyscan github.com >> /root/.ssh/known_hosts
+
+# Start SSH agent and add GitHub private key (if not already added)
+if ! pgrep -x "ssh-agent" > /dev/null; then
+    echo "🕵️  Starting SSH agent..."
+    eval "$(ssh-agent -s)"
+fi
+if ssh-add -l | grep -q github_rsa; then
+    echo "✅ GitHub SSH key already added."
+else
+    echo "🔑 Adding GitHub SSH key to agent..."
+    ssh-add /root/.ssh/github_rsa
+fi
+
+# Set Git user name and email (if provided)
+if [ -n "$GIT_USER_NAME" ]; then
+    echo "👤 Setting Git username: $GIT_USER_NAME"
+    git config --global user.name "$GIT_USER_NAME"
+fi
+if [ -n "$GIT_USER_EMAIL" ]; then
+    echo "📧 Setting Git email: $GIT_USER_EMAIL"
+    git config --global user.email "$GIT_USER_EMAIL"
+fi
+
+# Ensure first-time setup runs only once
+if [ ! -f /workspace/.setup_done ]; then
+    echo "🛠️ Running first-time setup..."
+
+    # Ensure GitHub username is set
+    if [ -z "$GITHUB_USER" ]; then
+        echo "❌ Error: Please set the GITHUB_USER environment variable."
+        exit 1
+    fi
+
+    # Clone GitHub repo if not already cloned
+    if [ ! -d "/workspace/.git" ]; then
+        echo "📦 Cloning repository: $GITHUB_USER/torch-uncertainty..."
+        git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace
+    fi
+
+    # Mark setup as completed
+    touch /workspace/.setup_done
+    echo "✅ First-time setup complete!"
+else
+    echo "⏩ Skipping first-time setup (already done)."
+fi
+
+# Apply compact shell prompt customization (if enabled)
+if [ -n "$USE_COMPACT_SHELL_PROMPT" ]; then
+    echo "🎨 Applying compact shell prompt customization..."
+    echo 'force_color_prompt=yes' >> /root/.bashrc
+    echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc
+    echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc
+    echo '    test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc
+    echo '    alias ls="ls --color=auto"' >> /root/.bashrc
+    echo '    alias grep="grep --color=auto"' >> /root/.bashrc
+    echo '    alias fgrep="fgrep --color=auto"' >> /root/.bashrc
+    echo '    alias egrep="egrep --color=auto"' >> /root/.bashrc
+    echo 'fi' >> /root/.bashrc
+fi
+
+# Ensure /workspace is in PYTHONPATH
+if ! echo "$PYTHONPATH" | grep -q "/workspace"; then
+    echo "📌 Adding /workspace to PYTHONPATH"
+    export PYTHONPATH="/workspace:$PYTHONPATH"
+else
+    echo "✅ PYTHONPATH is already correctly set."
+fi
+
+# Check if torch_uncertainty is installed in editable mode
+if pip show torch_uncertainty | grep -q "Editable project location: /workspace"; then
+    echo "✅ torch_uncertainty is already installed in editable mode. 🎉"
+else
+    echo "🔄 Reinstalling torch_uncertainty in editable mode..."
+    pip uninstall -y torch-uncertainty
+    pip install -e /workspace
+    echo "✅ torch_uncertainty is now installed in editable mode! 🚀"
+fi
+
+# Activate pre-commit hooks (if enabled)
+echo "🔗 Activating pre-commit hooks..."
+pre-commit install
+
+# Ensure SSH server is started
+echo "🔑 Starting SSH server..."
+mkdir -p /run/sshd && chmod 755 /run/sshd
+/usr/sbin/sshd -D

docs/source/api.rst

@@ -201,6 +201,7 @@ Functions
     :toctree: generated/
     :nosignatures:
 
+    batch_ensemble
     deep_ensembles
     mc_dropout
 
@@ -212,6 +213,7 @@ Classes
     :nosignatures:
     :template: class.rst
 
+    BatchEnsemble
     CheckpointEnsemble
     EMA
     MCDropout

docs/source/conf.py

@@ -15,7 +15,7 @@
     f"{datetime.now().year!s}, Adrien Lafage and Olivier Laurent"
 )
 author = "Adrien Lafage and Olivier Laurent"
-release = "0.4.1"
+release = "0.4.2"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

experiments/classification/mnist/configs/bayesian_lenet.yaml

@@ -41,7 +41,6 @@ model:
           norm: torch.nn.Identity
           groups: 1
           dropout_rate: 0
-          last_layer_dropout: false
           layer_args: {}
       num_samples: 16
   num_classes: 10

experiments/classification/mnist/configs/lenet.yaml

@@ -38,7 +38,6 @@ model:
       norm: torch.nn.Identity
       groups: 1
       dropout_rate: 0
-      last_layer_dropout: false
       layer_args: {}
   num_classes: 10
   loss: CrossEntropyLoss