Batched inference CEBRA & padding at the Solver level

[CeliaBenquet]

fix https://github.com/AdaptiveMotorControlLab/CEBRA-dev/pull/746
fix #199

This PR adds the following features:

• Inference (using CEBRA.transform() or solver.transform()) can be performed in batch, allowing inference on larger datasets or with larger models in a memory-efficient way (fix https://github.com/AdaptiveMotorControlLab/CEBRA-dev/issues/624).
• All sklearn-level functionalities regarding input padding and transform() are now in the PyTorch API directly, making the PyTorch API more easily usable for developers (fix https://github.com/AdaptiveMotorControlLab/CEBRA-dev/issues/637, fix https://github.com/AdaptiveMotorControlLab/CEBRA-dev/pull/594).
• Corresponding tests for batched transform with both APIs.
• Some extra tests on the xCEBRA integration (using the transform method).

Example Usage of the new PyTorch API:

    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",
        model=neural_model,
        criterion=crit,
        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
    x_train_emb = solver.transform(
        torch.from_numpy(neural_data).type(torch.FloatTensor).to(device),
        pad_before_transform=True,
        batch_size=512).to(device)

    # 8. Plot Embedding
    cebra.plot_embedding(
        x_train_emb.cpu(),
        discrete_label[:,0],
        markersize=10,
    )

all is similar to previous implementation but the inference part, which doesn't require to handle the padding of the input before passing it to the model.

[stes]

some early comments; apologies if i have asked some of these before

[stes]

why the changes here? i.e. did anything change that would cause the "old" multi session tests to break?

[CeliaBenquet]

I restablished _get_loader as it was but added a return value as I need the dataset to configure it with the model.

Else,

• I added the model_architecture as offset1-model is a special case for padding at transform.
• I added the configure_for(model) as now this is handled in the solver.
• I added some tests on the transform (was not done at all before), similar to the sklearn tests but at the pytorch level.

[stes]

can you wrap the for loops here (quite complex) in functions, and only do the assingment on the global level?

[CeliaBenquet]

I proposed something lmk if that's what you meant :)

[MMathisLab]

doc error is: /home/runner/work/CEBRA/CEBRA/cebra/data/single_session.py:docstring of cebra.data.single_session.SingleSessionDataset.configure_for:3: WARNING: py:attr reference target not found: cebra_data.Dataset.offset /home/runner/work/CEBRA/CEBRA/cebra/data/multi_session.py:docstring of cebra.data.multi_session.MultiSessionDataset.configure_for:3: WARNING: py:attr reference target not found: cebra_data.Dataset.offset

[MMathisLab]

@CeliaBenquet not sure I see your edits post review; did you push them?

[stes]

Left some initial comments; broader discussion is a bit on the api design in the solver/base class --- lets discuss offline.

[stes]

Ok, review got a bit longer again; I realized I missed a few things on the last review. High level comments:

• I made some comments in solver which could be fine; I think some arguments were moved from the sklearn class to the solver class, but the motivation for that is not entirely clear. Mostly needs one round of discussion so we can settle on a good API design for these. Specifically, what is the usecase for storing these variables now in the solver, where are they called?
• the new transform function adds a lot of duplicated code that should be unified; again, could be first discussed

[stes]

Why is this an attribute of the solver, vs. being returned directly from the model? For sklearn it makes sense to fix this, but for the solver this could also simply be a property to be returned from the model? Where is this used?

E.g. what would happen for an xCEBRA solver, where you have not a single feature dim, but multiple

[CeliaBenquet]

for the multisession case that's already the case and that's a list.

num_features cannot be a property I think, because that can be defined only based on the inputs provided to the fit(), and later if we adapt the solver, it needs to be reset. This is used to be saved with the solver as it's needed when reloading it + to be called in the sklearn + to see if the solver is fitted when calling transform().

for xcebra that's just similar to the original sklearn one but at a lower level, so yes we need to think about it but we would have had to in any case.

[MMathisLab]

At the risk of it not being perfect, lets merge this now; @CeliaBenquet can document in an issue the remaining Qs on the API design, but getting #251 merged is a priority 🦾