How to prepare data for SupervisedNN10

[icarosadero]

I have neural data of a single individual of a single experiment over time. I also have a binary label that goes along with it. My goal is to use the SupervisedNN10 to see if it is able to predict the label at a given time.

I loaded the data using cebra.data.TensorDataset for both train and test like this:

with device:
    train = cebra.data.TensorDataset(
        neural = torch.from_numpy(X_train).type(torch.FloatTensor),
        discrete = torch.from_numpy(y_train).type(torch.LongTensor),
        offset = cebra.data.Offset(1) #It doesn't matter what you put here, it gets updated later
    )
    test = cebra.data.TensorDataset(
        neural = torch.from_numpy(X_test).type(torch.FloatTensor),
        discrete = torch.from_numpy(y_test).type(torch.LongTensor),
        offset = cebra.data.Offset(1)
    )

Then, I instantiated the model like so:

with device:
    model10 = cebra.models.SupervisedNN10(
        num_neurons = X_train.shape[1],
        num_units = 10,
        num_output = 1
    )
    train.configure_for(model10)
    test.configure_for(model10)

By the way, I don't know what "num_units" stands for in here.

Now, I expect train.offset and test.offset to have size 10.

The criterion and optimizer were instantiated like this:

with device:
    criterion = cebra.models.criterions.LearnableCosineInfoNCE()
    optimizer = torch.optim.Adam(
        list(model10.parameters()) + list(criterion.parameters()),
        lr = 1e-3,
        weight_decay= 0
    )

And finally, the solver and Data Loaders:

with device:
    solver = cebra.solver.SupervisedNNSolver(
        model = model10,
        criterion = criterion,
        optimizer = optimizer,
        tqdm_on = True
    )

    train_loader = cebra.data.single_session.DiscreteDataLoader(
        dataset = train,
        num_steps = 500,
        batch_size = 256,
        prior = "uniform"
    )
    test_loader = cebra.data.single_session.DiscreteDataLoader(
        dataset = test,
        num_steps = 500,
        batch_size = 256,
        prior = "uniform"
    )

Now, one thing that I noticed getting an example of the Data Loader output is that it can't be indexed. And indeed, when I try to run the solver, I get this error:

{
	"name": "TypeError",
	"message": "'Batch' object is not subscriptable",
	"stack": "---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[113], line 1
----> 1 solver.fit(loader=train_loader, valid_loader=test_loader, valid_frequency=10, num_steps=500)

File /path/to/supervised.py:73, in SupervisedNNSolver.fit(self, loader, num_steps, valid_loader, save_frequency, valid_frequency, decode, logdir)
     71 while True:
     72     for _, batch in enumerate(loader):
---> 73         stats = self.step(batch)
     74         self._log_checkpoint(num_steps, loader, valid_loader)
     75         step_idx += 1

File /path/to/supervised.py:89, in SupervisedNNSolver.step(self, batch)
     80 \"\"\"Perform a single gradient update.
     81
     82 Args:
   (...)
     86     Dictionary containing training metrics TODO
     87 \"\"\"
     88 self.optimizer.zero_grad()
---> 89 prediction = self._inference(batch)
     90 loss = self.criterion(prediction, batch[\"label\"].squeeze())
     91 loss.backward()

File /path/to/supervised.py:98, in SupervisedNNSolver._inference(self, batch)
     96 def _inference(self, batch):
     97     \"\"\"Compute predictions (discrete/continuous) for the batch.\"\"\"
---> 98     feature, prediction = self.model(batch[\"neural\"])
     99     return prediction

TypeError: 'Batch' object is not subscriptable"
}

Am I supposed to be using some other Data Loader?

[stes]

Hi @icarosadero , happy to check. Could you upload your whole code to repro your error?