Tensor dimension order in primate reaching benchmark

[morenzoe]

Hi, I would like to ask about the order of tensor dimension for primate reaching task benchmark using LSTM model.

This is my code:

import torch #2.3.0
from neurobench.datasets import PrimateReaching #1.0.5
from neurobench.models.torch_model import TorchModel
from neurobench.benchmarks import Benchmark
from torch.utils.data import DataLoader, Subset

class LSTMCELL(torch.nn.Module):
    def __init__(self, hidden):
        super().__init__()

        self.hidden = hidden

        self.lstm_cell = torch.nn.LSTMCell(96, self.hidden)  # [channel, hidden_size]
        self.fc = torch.nn.Linear(self.hidden, out_features=2)

    def forward(self, x):
        # x [batch_size, num_steps, channel]
        hx = torch.randn(x.shape[0], self.hidden)  # [batch_size, hidden_size]
        cx = torch.randn(x.shape[0], self.hidden)

        for i in range(x.shape[1]):
            hx, cx = self.lstm_cell(x[:, i, :], (hx, cx))

        x = self.fc(hx)

        return x

filename = "indy_20160622_01"
data_dir = ".../neurobench/data"

dataset = PrimateReaching(file_path=data_dir, filename=filename,
                        num_steps=7, train_ratio=0.5, bin_width=0.004, label_series=False,
                        biological_delay=0, remove_segments_inactive=False, download=False)

dataloader = DataLoader(Subset(dataset, dataset.ind_train), batch_size=1024, shuffle=False)

model = LSTMCELL(256)

model = TorchModel(model)

static_metrics = ["footprint", "connection_sparsity"]
workload_metrics = ["r2", "activation_sparsity", "synaptic_operations"]

benchmark = Benchmark(model, dataloader, [], [], [static_metrics, workload_metrics])
results = benchmark.run()
print(results)

This is the error I got:

File "...\neurobench\neurobench\utils.py", line 296, in single_layer_MACs
    c_1 = ifgo[1, :] * inputs[1][1] + ifgo[0, :] * ifgo[2, :]
RuntimeError: The size of tensor a (262144) must match the size of tensor b (256) at non-singleton dimension 1

I am using a simple LSTMCell from PyTorch example but I change the for loop to iterate on the second dimension since Neurobench dataloader give samples in the shape of [batch_size, num_steps, channel].

Is there something wrong with how I set up the LSTM model? Does the forward function has to iterate through the batch_size to use the benchmark? Thank you in advance!

[jasonlyik]

@morenzoe Thanks for bringing this up, there was a small bug in tensor dimensions for the synops calculation of LSTMCell and GRUCell. I have fixed it in #227 and it should be pulled in soon.

[morenzoe]

Thank you for the quick fix! Looking forward to the pip version update 👍

[morenzoe]

Hi @jasonlyik, it would be very helpful if this fix can be pulled in this week since the BioCAS Grand Challenge deadline is approaching. However I totally understand if it still needs more checks done to ensure proper update. Thanks!

[jasonlyik]

@morenzoe Pulled the fix into latest 1.0.6 pip release.

Also in case you don't want to wait for a new pip release, you can use/edit a local version of the repo by cloning and using poetry, here are instructions: https://github.com/NeuroBench/neurobench#development

Good luck with the BioCAS challenge!