Merge branch 'main' into add_options_mixed

Author: stes

.github/workflows/build.yml

@@ -14,24 +14,15 @@ jobs:
  fail-fast: true
  matrix:
  os: [ubuntu-latest]
- python-version: ["3.8", "3.10"]
+ python-version: ["3.9", "3.10", "3.12"]
  # We aim to support the versions on pytorch.org
  # as well as selected previous versions on
  # https://pytorch.org/get-started/previous-versions/
- torch-version: ["1.12.1", "2.0.0"]
+ torch-version: ["2.2.2", "2.4.0"]
  include:
- - os: ubuntu-latest
- python-version: 3.8
- torch-version: 1.9.0
  - os: windows-latest
- torch-version: 2.0.0
+ torch-version: 2.4.0
  python-version: "3.10"
- - os: ubuntu-latest
- torch-version: 2.1.1
- python-version: "3.11"
- #- os: macos-latest
- # torch-version: 2.0.0
- # python-version: "3.10"
 
  runs-on: ${{ matrix.os }}
 

Dockerfile

@@ -1,15 +1,13 @@
 ## EXPERIMENT BASE CONTAINER
-FROM nvidia/cuda:11.7.1-cudnn8-runtime-ubuntu20.04 AS cebra-base
+FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS cebra-base
 
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update -y \
 && apt-get install --no-install-recommends -yy git python3 python3-pip python-is-python3 \
 && rm -rf /var/lib/apt/lists/*
 
-RUN pip install --no-cache-dir torch==2.0.0+cu117 \
---index-url https://download.pytorch.org/whl/cu117
-RUN pip install --no-cache-dir --pre 'cebra[dev,datasets,integrations]' \
-&& pip uninstall -y cebra
+RUN pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124
+RUN pip install --upgrade pip
 
 
 ## GIT repository

Makefile

@@ -24,6 +24,9 @@ test: clean_test
 doctest: clean_test
 python -m pytest --ff --doctest-modules -vvv ./docs/source/usage.rst
 
+docker:
+./tools/build_docker.sh
+
 test_parallel: clean_test
 python -m pytest -n auto --ff -m "not requires_dataset" tests
 
@@ -98,4 +101,7 @@ report: check_docker format .coverage .pylint
 cat .pylint
 coverage report
 
-.PHONY: dist build archlinux clean_test test doctest test_parallel test_parallel_debug test_all test_fast test_debug test_benchmark interrogate docs docs-touch docs-strict serve_docs serve_page format codespell check_for_binary
+.PHONY: dist build docker archlinux clean_test test doctest test_parallel \
+test_parallel_debug test_all test_fast test_debug test_benchmark \
+interrogate docs docs-touch docs-strict serve_docs serve_page \
+format codespell check_for_binary

cebra/integrations/plotly.py

@@ -22,6 +22,7 @@
 """Plotly interface to CEBRA."""
 from typing import Optional, Tuple, Union
 
+import matplotlib.cm
 import matplotlib.colors
 import numpy as np
 import numpy.typing as npt

docs/source/usage.rst

@@ -465,13 +465,13 @@ Similarly, for the discrete case a discrete label can be provided and the CEBRA
  discrete_label1 = np.random.randint(0,10,(timesteps1, ))
  discrete_label2 = np.random.randint(0,10,(timesteps2, ))
 
- multi_cebra_model = cebra.CEBRA(batch_size=512,
+ multi_cebra_model_discrete = cebra.CEBRA(batch_size=512,
  output_dimension=out_dim,
  max_iterations=10,
  max_adapt_iterations=10)
 
 
- multi_cebra_model.fit([neural_session1, neural_session2], [discrete_label1, discrete_label2])
+ multi_cebra_model_discrete.fit([neural_session1, neural_session2], [discrete_label1, discrete_label2])
 
 .. admonition:: See API docs
  :class: dropdown
@@ -1348,15 +1348,15 @@ Below is the documentation on the available arguments.
  --valid-ratio 0.1 Ratio of validation set after the train data split. The remaining will be test split
  --share-model
 
-Model training using the Torch API 
+Model training using the Torch API
 ----------------------------------
 
 The scikit-learn API provides parametrization to many common use cases.
-The Torch API however allows for more flexibility and customization, for e.g. 
+The Torch API however allows for more flexibility and customization, for e.g.
 sampling, criterions, and data loaders.
 
 In this minimal example we show how to initialize a CEBRA model using the Torch API.
-Here the :py:class:`cebra.data.single_session.DiscreteDataLoader` 
+Here the :py:class:`cebra.data.single_session.DiscreteDataLoader`
 gets initialized which also allows the `prior` to be directly parametrized.
 
 👉 For an example notebook using the Torch API check out the :doc:`demo_notebooks/Demo_Allen`.
@@ -1367,45 +1367,45 @@ gets initialized which also allows the `prior` to be directly parametrized.
  import numpy as np
  import cebra.datasets
  import torch
- 
+
  if torch.cuda.is_available():
  device = "cuda"
  else:
  device = "cpu"
- 
+
  neural_data = cebra.load_data(file="neural_data.npz", key="neural")
- 
+
  discrete_label = cebra.load_data(
  file="auxiliary_behavior_data.h5", key="auxiliary_variables", columns=["discrete"],
  )
- 
+
  # 1. Define a CEBRA-ready dataset
  input_data = cebra.data.TensorDataset(
  torch.from_numpy(neural_data).type(torch.FloatTensor),
  discrete=torch.from_numpy(np.array(discrete_label[:, 0])).type(torch.LongTensor),
  ).to(device)
- 
+
  # 2. Define a CEBRA model
  neural_model = cebra.models.init(
  name="offset10-model",
  num_neurons=input_data.input_dimension,
  num_units=32,
  num_output=2,
  ).to(device)
- 
+
  input_data.configure_for(neural_model)
- 
+
  # 3. Define the Loss Function Criterion and Optimizer
  crit = cebra.models.criterions.LearnableCosineInfoNCE(
  temperature=1,
  ).to(device)
- 
+
  opt = torch.optim.Adam(
  list(neural_model.parameters()) + list(crit.parameters()),
  lr=0.001,
  weight_decay=0,
  )
- 
+
  # 4. Initialize the CEBRA model
  solver = cebra.solver.init(
  name="single-session",
@@ -1414,27 +1414,27 @@ gets initialized which also allows the `prior` to be directly parametrized.
  optimizer=opt,
  tqdm_on=True,
  ).to(device)
- 
+
  # 5. Define Data Loader
  loader = cebra.data.single_session.DiscreteDataLoader(
  dataset=input_data, num_steps=10, batch_size=200, prior="uniform"
  )
- 
+
  # 6. Fit Model
  solver.fit(loader=loader)
- 
+
  # 7. Transform Embedding
  train_batches = np.lib.stride_tricks.sliding_window_view(
  neural_data, neural_model.get_offset().__len__(), axis=0
  )
- 
+
  x_train_emb = solver.transform(
  torch.from_numpy(train_batches[:]).type(torch.FloatTensor).to(device)
  ).to(device)
- 
+
  # 8. Plot Embedding
  cebra.plot_embedding(
- x_train_emb,
+ x_train_emb.cpu(),
  discrete_label[neural_model.get_offset().__len__() - 1 :, 0],
  markersize=10,
  )

setup.cfg

@@ -28,7 +28,7 @@ packages = find:
 where =
  - .
  - tests
-python_requires = >=3.8
+python_requires = >=3.9
 install_requires =
  joblib
  literate-dataclasses
@@ -68,7 +68,8 @@ docs =
  matplotlib<=3.5.2
  pandas
  seaborn
- scikit-learn<1.3
+ scikit-learn
+ numpy<2.0.0
 demos =
  ipykernel
  jupyter
@@ -89,12 +90,12 @@ dev =
  isort
  toml
  coverage
- pytest==7.4.4
+ pytest
  pytest-benchmark
  pytest-xdist
  pytest-timeout
- pytest-sphinx==0.5.0
- tables<=3.8
+ pytest-sphinx
+ tables
  licenseheaders
  # TODO(stes) Add back once upstream issue
  # https://github.com/PyCQA/docformatter/issues/119
@@ -105,4 +106,3 @@ dev =
 
 [bdist_wheel]
 universal=1
-

tests/test_criterions.py

@@ -260,9 +260,9 @@ def _reference_infonce(pos_dist, neg_dist):
 
 def test_similiarities():
  rng = torch.Generator().manual_seed(42)
- ref = torch.randn(10, 3, generator = rng)
- pos = torch.randn(10, 3, generator = rng)
- neg = torch.randn(12, 3, generator = rng)
+ ref = torch.randn(10, 3, generator=rng)
+ pos = torch.randn(10, 3, generator=rng)
+ neg = torch.randn(12, 3, generator=rng)
 
  pos_dist, neg_dist = _reference_dot_similarity(ref, pos, neg)
  pos_dist_2, neg_dist_2 = cebra_criterions.dot_similarity(ref, pos, neg)
@@ -307,37 +307,47 @@ def test_infonce(seed):
 
 
 @pytest.mark.parametrize("seed", [42, 4242, 424242])
-def test_infonce_gradients(seed):
+@pytest.mark.parametrize("case", [0, 1, 2])
+def test_infonce_gradients(seed, case):
  pos_dist, neg_dist = _sample_dist_matrices(seed)
 
- for i in range(3):
- pos_dist_ = pos_dist.clone()
- neg_dist_ = neg_dist.clone()
- pos_dist_.requires_grad_(True)
- neg_dist_.requires_grad_(True)
- loss_ref = _reference_infonce(pos_dist_, neg_dist_)[i]
- grad_ref = _compute_grads(loss_ref, [pos_dist_, neg_dist_])
-
- pos_dist_ = pos_dist.clone()
- neg_dist_ = neg_dist.clone()
- pos_dist_.requires_grad_(True)
- neg_dist_.requires_grad_(True)
- loss = cebra_criterions.infonce(pos_dist_, neg_dist_)[i]
- grad = _compute_grads(loss, [pos_dist_, neg_dist_])
-
- # NOTE(stes) default relative tolerance is 1e-5
- assert torch.allclose(loss_ref, loss, rtol=1e-4)
-
- if i == 0:
- assert grad[0] is not None
- assert grad[1] is not None
- assert torch.allclose(grad_ref[0], grad[0])
- assert torch.allclose(grad_ref[1], grad[1])
- if i == 1:
- assert grad[0] is not None
- assert grad[1] is None
- assert torch.allclose(grad_ref[0], grad[0])
- if i == 2:
- assert grad[0] is None
- assert grad[1] is not None
- assert torch.allclose(grad_ref[1], grad[1])
+ # TODO(stes): This test seems to fail due to some recent software
+ # updates; root cause not identified. Remove this comment once
+ # fixed. (for i = 0, 1)
+ pos_dist_ = pos_dist.clone()
+ neg_dist_ = neg_dist.clone()
+ pos_dist_.requires_grad_(True)
+ neg_dist_.requires_grad_(True)
+ loss_ref = _reference_infonce(pos_dist_, neg_dist_)[case]
+ grad_ref = _compute_grads(loss_ref, [pos_dist_, neg_dist_])
+
+ pos_dist_ = pos_dist.clone()
+ neg_dist_ = neg_dist.clone()
+ pos_dist_.requires_grad_(True)
+ neg_dist_.requires_grad_(True)
+ loss = cebra_criterions.infonce(pos_dist_, neg_dist_)[case]
+ grad = _compute_grads(loss, [pos_dist_, neg_dist_])
+
+ # NOTE(stes) default relative tolerance is 1e-5
+ assert torch.allclose(loss_ref, loss, rtol=1e-4)
+
+ if case == 0:
+ assert grad[0] is not None
+ assert grad[1] is not None
+ assert torch.allclose(grad_ref[0], grad[0])
+ assert torch.allclose(grad_ref[1], grad[1])
+ if case == 1:
+ assert grad[0] is not None
+ assert torch.allclose(grad_ref[0], grad[0])
+ # TODO(stes): This is most likely not the right fix, needs more
+ # investigation. On the first run of the test, grad[1] is actually
+ # None, and then on the second run of the test it is a Tensor, but
+ # with zeros everywhere. The behavior is fine for fitting models,
+ # but there is some side-effect in our test suite we need to fix.
+ if grad[1] is not None:
+ assert torch.allclose(grad[1], torch.zeros_like(grad[1]))
+ if case == 2:
+ if grad[0] is not None:
+ assert torch.allclose(grad[0], torch.zeros_like(grad[0]))
+ assert grad[1] is not None
+ assert torch.allclose(grad_ref[1], grad[1])

tests/test_dlc.py

@@ -35,14 +35,13 @@
 # /Reaching-Mackenzie-2018-08-30/labeled-data/reachingvideo1
 # /CollectedData_Mackenzie.h5?raw=true
 # which is replaced here due to rate limitations we observed in the past.
-ANNOTATED_DLC_URL = "https://figshare.com/ndownloader/files/42303564?private_link=b917317bfab725e0b207"
+ANNOTATED_DLC_URL = "https://cebra.fra1.digitaloceanspaces.com/CollectedData_Mackenzie.h5"
 
 # NOTE(stes): The original data URL is
 # https://github.com/DeepLabCut/UnitTestData/raw/main/data.zip")
 # which is replaced here due to rate limitations we observed in the past.
 MULTISESSION_PRED_DLC_URL = (
- "https://figshare.com/ndownloader/files/42303561?private_link=b917317bfab725e0b207"
-)
+ "https://cebra.fra1.digitaloceanspaces.com/data.zip")
 
 MULTISESSION_PRED_KEYPOINTS = ["head", "tail"]
 ANNOTATED_KEYPOINTS = ["Hand", "Tongue"]

tools/build_docker.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 # Build, test and push cebra container.
 
-set -xe
+set -e
 
 if [[ -z $(git status --porcelain) ]]; then
  TAG=$(git rev-parse --short HEAD)
@@ -17,19 +17,19 @@ echo Building $DOCKERNAME
 #docker login <your registry>
 
 docker build \
---build-arg UID=$(id -u) \
---build-arg GID=$(id -g) \
---build-arg GIT_HASH=$(git rev-parse HEAD) \
- -t $DOCKERNAME .
+--build-arg UID=$(id -u) \
+--build-arg GID=$(id -g) \
+--build-arg GIT_HASH=$(git rev-parse HEAD) \
+-t $DOCKERNAME .
 docker tag $DOCKERNAME $LATEST
 
 docker run \
---gpus 2 \
--v ${CEBRA_DATADIR:-./data}:/data \
---env CEBRA_DATADIR=/data \
---network host \
--it $DOCKERNAME python -m pytest --doctest-modules tests ./docs/source/usage.rst cebra
-
+ --gpus 2 \
+ ${extra_kwargs[@]} \
+ -v ${CEBRA_DATADIR:-./data}:/data \
+ --env CEBRA_DATADIR=/data \
+ --network host \
+ -it $DOCKERNAME python -m pytest --ff -x -m "not requires_dataset" --doctest-modules ./docs/source/usage.rst tests cebra
 
 #docker push $DOCKERNAME
 #docker push $LATEST