Masksemble

Masksemble#

[1]:

%pip install git+https://github.com/lightning-uq-box/lightning-uq-box.git

Collecting git+https://github.com/lightning-uq-box/lightning-uq-box.git
  Cloning https://github.com/lightning-uq-box/lightning-uq-box.git to /tmp/pip-req-build-eomisgjn
  Running command git clone --filter=blob:none --quiet https://github.com/lightning-uq-box/lightning-uq-box.git /tmp/pip-req-build-eomisgjn
  Resolved https://github.com/lightning-uq-box/lightning-uq-box.git to commit bb7f5757d8a82b4dd45347280551f01996c50da7
  Installing build dependencies ... - \ done
  Getting requirements to build wheel ... - done
  Preparing metadata (pyproject.toml) ... - done
Requirement already satisfied: einops>=0.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.8.2)
Requirement already satisfied: jsonargparse>=4.28.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (4.46.0)
Requirement already satisfied: lightning>=2.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: matplotlib>=3.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.10.8)
Requirement already satisfied: numpy>=1.21.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.26.4)
Requirement already satisfied: omegaconf>=2.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.3.0)
Requirement already satisfied: pandas>=1.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.0.0)
Requirement already satisfied: torch>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.10.0)
Requirement already satisfied: torchmetrics>=1.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.8.2)
Requirement already satisfied: torchvision>=0.16.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.25.0)
Requirement already satisfied: scikit-learn>=1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.8.0)
Requirement already satisfied: gpytorch>=1.11 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.15.1)
Requirement already satisfied: laplace-torch>=0.2.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.2.2.2)
Requirement already satisfied: curvlinops-for-pytorch==2.0.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.0.1)
Requirement already satisfied: uncertainty-toolbox>=0.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.1.1)
Requirement already satisfied: kornia>=0.6.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.8.2)
Requirement already satisfied: timm>=0.9.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.0.24)
Requirement already satisfied: torchseg>=0.0.1a1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.0.1a4)
Requirement already satisfied: h5py>=3.12.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.15.1)
Requirement already satisfied: ema-pytorch>=0.7.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.7.9)
Requirement already satisfied: vbll>=0.4.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.4.9)
Requirement already satisfied: backpack-for-pytorch<2.0.0,>=1.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (1.7.1)
Requirement already satisfied: scipy<2.0.0,>=1.7.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (1.17.0)
Requirement already satisfied: tqdm<5.0.0,>=4.61.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (4.67.3)
Requirement already satisfied: einconv in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (0.1.0)
Requirement already satisfied: unfoldNd<1.0.0,>=0.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from backpack-for-pytorch<2.0.0,>=1.6.0->curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (0.2.3)
Requirement already satisfied: packaging in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from unfoldNd<1.0.0,>=0.2.0->backpack-for-pytorch<2.0.0,>=1.6.0->curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (26.0)
Requirement already satisfied: jaxtyping in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.3.7)
Requirement already satisfied: mpmath<=1.3,>=0.19 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (1.3.0)
Requirement already satisfied: linear_operator>=0.6 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.6)
Requirement already satisfied: PyYAML>=3.13 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse>=4.28.0->jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (6.0.3)
Requirement already satisfied: docstring-parser>=0.17 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (0.17.0)
Requirement already satisfied: typeshed-client>=2.8.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (2.8.2)
Requirement already satisfied: kornia_rs>=0.1.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from kornia>=0.6.9->lightning-uq-box==0.3.0.dev0) (0.1.10)
Requirement already satisfied: asdfghjkl==0.1a4 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from laplace-torch>=0.2.1->lightning-uq-box==0.3.0.dev0) (0.1a4)
Requirement already satisfied: opt_einsum in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from laplace-torch>=0.2.1->lightning-uq-box==0.3.0.dev0) (3.4.0)
Requirement already satisfied: fsspec<2028.0,>=2022.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2026.1.0)
Requirement already satisfied: lightning-utilities<2.0,>=0.10.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (0.15.2)
Requirement already satisfied: typing-extensions<6.0,>4.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (4.15.0)
Requirement already satisfied: pytorch-lightning in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (3.13.3)
Requirement already satisfied: setuptools in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-utilities<2.0,>=0.10.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (80.9.0)
Requirement already satisfied: filelock in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.20.3)
Requirement already satisfied: sympy>=1.13.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.14.0)
Requirement already satisfied: networkx>=2.5.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.6.1)
Requirement already satisfied: jinja2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.1.6)
Requirement already satisfied: cuda-bindings==12.9.4 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.9.4)
Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.93)
Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (9.10.2.21)
Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.4.1)
Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (11.3.3.83)
Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (10.3.9.90)
Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (11.7.3.90)
Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.5.8.93)
Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (0.7.1)
Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (2.27.5)
Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.4.5)
Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.93)
Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.13.1.3)
Requirement already satisfied: triton==3.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.6.0)
Requirement already satisfied: cuda-pathfinder~=1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from cuda-bindings==12.9.4->torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.3.3)
Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: aiosignal>=1.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.4.0)
Requirement already satisfied: attrs>=17.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (25.4.0)
Requirement already satisfied: frozenlist>=1.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.8.0)
Requirement already satisfied: multidict<7.0,>=4.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (6.7.1)
Requirement already satisfied: propcache>=0.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (0.4.1)
Requirement already satisfied: yarl<2.0,>=1.17.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.22.0)
Requirement already satisfied: idna>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (3.11)
Requirement already satisfied: contourpy>=1.0.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.3.3)
Requirement already satisfied: cycler>=0.10 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (4.61.1)
Requirement already satisfied: kiwisolver>=1.3.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.4.9)
Requirement already satisfied: pillow>=8 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (12.1.0)
Requirement already satisfied: pyparsing>=3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (3.3.2)
Requirement already satisfied: python-dateutil>=2.7 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (2.9.0.post0)
Requirement already satisfied: antlr4-python3-runtime==4.9.* in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from omegaconf>=2.3.0->lightning-uq-box==0.3.0.dev0) (4.9.3)
Requirement already satisfied: six>=1.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.17.0)
Requirement already satisfied: joblib>=1.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from scikit-learn>=1.3->lightning-uq-box==0.3.0.dev0) (1.5.3)
Requirement already satisfied: threadpoolctl>=3.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from scikit-learn>=1.3->lightning-uq-box==0.3.0.dev0) (3.6.0)
Requirement already satisfied: huggingface_hub in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.3.7)
Requirement already satisfied: safetensors in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.7.0)
Requirement already satisfied: importlib_resources>=1.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from typeshed-client>=2.8.2->jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (6.5.2)
Requirement already satisfied: hf-xet<2.0.0,>=1.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.2.0)
Requirement already satisfied: httpx<1,>=0.23.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.28.1)
Requirement already satisfied: shellingham in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.5.4)
Requirement already satisfied: typer-slim in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.21.1)
Requirement already satisfied: anyio in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (4.12.1)
Requirement already satisfied: certifi in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (2026.1.4)
Requirement already satisfied: httpcore==1.* in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.0.9)
Requirement already satisfied: h11>=0.16 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.16.0)
Requirement already satisfied: wadler-lindig>=0.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jaxtyping->gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.1.7)
Requirement already satisfied: MarkupSafe>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jinja2->torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.0.3)
Requirement already satisfied: click>=8.0.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from typer-slim->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (8.3.1)
Note: you may need to restart the kernel to use updated packages.

Imports#

[2]:

import os
import tempfile
from functools import partial

import matplotlib.pyplot as plt
import torch.nn as nn
from lightning import Trainer
from lightning.pytorch import seed_everything
from lightning.pytorch.loggers import CSVLogger
from torch.optim import Adam

from lightning_uq_box.datamodules import TwoMoonsDataModule
from lightning_uq_box.models import MLP
from lightning_uq_box.uq_methods import MasksemblesClassification
from lightning_uq_box.viz_utils import (
    plot_predictions_classification,
    plot_training_metrics,
    plot_two_moons_data,
)

plt.rcParams["figure.figsize"] = [14, 5]

%load_ext autoreload
%autoreload 2

[3]:

seed_everything(0)

Seed set to 0

[3]:

0

[4]:

my_temp_dir = tempfile.mkdtemp()

Datamodule#

[5]:

dm = TwoMoonsDataModule(batch_size=128)

X_train, Y_train, X_test, Y_test, test_grid_points = (
    dm.X_train,
    dm.Y_train,
    dm.X_test,
    dm.Y_test,
    dm.test_grid_points,
)

[6]:

fig = plot_two_moons_data(X_train, Y_train, X_test, Y_test)
../../_images/tutorials_classification_masksembles_8_0.png

Model#

[7]:

network = MLP(
    n_inputs=2,
    n_hidden=[50, 50, 50],
    n_outputs=2,
    dropout_p=0.2,
    activation_fn=nn.ReLU(),
)
network

[7]:

MLP(
  (model): Sequential(
    (0): Linear(in_features=2, out_features=50, bias=True)
    (1): ReLU()
    (2): Dropout(p=0.2, inplace=False)
    (3): Linear(in_features=50, out_features=50, bias=True)
    (4): ReLU()
    (5): Dropout(p=0.2, inplace=False)
    (6): Linear(in_features=50, out_features=50, bias=True)
    (7): ReLU()
    (8): Dropout(p=0.2, inplace=False)
    (9): Linear(in_features=50, out_features=2, bias=True)
  )
)

[8]:

masksemble = MasksemblesClassification(
    model=network,
    optimizer=partial(Adam, lr=1e-2),
    loss_fn=nn.CrossEntropyLoss(),
    num_estimators=5,
    scale=3.5,
)

Trainer#

[9]:

logger = CSVLogger(my_temp_dir)
trainer = Trainer(
    accelerator="cpu",
    max_epochs=100,  # number of epochs we want to train
    logger=logger,  # log training metrics for later evaluation
    log_every_n_steps=1,
    enable_checkpointing=False,
    enable_progress_bar=False,
    default_root_dir=my_temp_dir,
)

GPU available: False, used: False
TPU available: False, using: 0 TPU cores
💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

[10]:

trainer.fit(masksemble, dm)


  | Name          | Type             | Params | Mode  | FLOPs
-------------------------------------------------------------------
0 | loss_fn       | CrossEntropyLoss | 0      | train | 0
1 | model         | MLP              | 6.1 K  | train | 0
2 | train_metrics | MetricCollection | 0      | train | 0
3 | val_metrics   | MetricCollection | 0      | train | 0
4 | test_metrics  | MetricCollection | 0      | train | 0
-------------------------------------------------------------------
5.4 K     Trainable params
750       Non-trainable params
6.1 K     Total params
0.024     Total estimated model params size (MB)
29        Modules in train mode
0         Modules in eval mode
0         Total Flops
/home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages/lightning/pytorch/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
`Trainer.fit` stopped: `max_epochs=100` reached.

Training Metrics#

[11]:

fig = plot_training_metrics(
    os.path.join(my_temp_dir, "lightning_logs"), ["train_loss", "trainAcc"]
)
../../_images/tutorials_classification_masksembles_16_0.png

Prediction#

[12]:

# save predictions
trainer.test(masksemble, dm.test_dataloader())

────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
       Test metric             DataLoader 0
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
         testAcc                    1.0
     testCalibration       0.0013964951504021883
 testEmpirical Coverage             1.0
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────

[12]:

[{'testAcc': 1.0,
  'testCalibration': 0.0013964951504021883,
  'testEmpirical Coverage': 1.0}]

[13]:

preds = masksemble.predict_step(test_grid_points.to(masksemble.device))

Evaluate Predictions#

[14]:

fig = plot_predictions_classification(
    X_test, Y_test, preds["pred"].argmax(-1), test_grid_points, preds["pred_uct"]
)
../../_images/tutorials_classification_masksembles_21_0.png
Contents