I am trying to index a 3-dimensional torch tensor based on timesteps acquired with torch.nonzeroes (for a latency decoder in a Neuromorphic Computing project), but am unsure on how to proceed.
My input tensor x has dimensions [Timesteps, Batchsize, Classes] = [48, 256, 10].
I am using torch.nonzeroes to acquire an array of 256 unique timesteps (1 for each sample in the batch), where each timestep is the first occuring nonzero value for that sample, along the time-axis of 48 timestamps ( I realize this is probably quite inefficient but havn't found a better way thus far).
nonzeroes = torch.nonzero(x,as_tuple=True)
FirstSpike = []
for i in range(0,np.size(x.cpu().detach().numpy(),1)):
nonzeroes = torch.nonzero(x[:,i,:],as_tuple=True)
FirstSpike.append(nonzeroes[0][0].cpu().detach().numpy())
This returns FirstSpike as an array of 256 integer values for indexing/slicing each timestep ( [7, 9, 13, 43,...] ).
I want to use this to end up with an array x2 with dimensions [256, 10], where each sample of the 256 batch corresponds to its appropriate slice in time. (for example, sample 17 could have timestep 7, whereas sample 57 has timestep 38 ). I know I could probably obtain this in a for loop like this:
x2 = []
for i in range(0,np.size(x.cpu().detach().numpy(),1)):
val = x[FirstSpike[i],i,:] #output dimension [1,10]
x2.append(val) #Final x2 dimension [256,10]
However, since this is part of a neural network decoder, that would be very inefficient, memory-wise. Is there perhaps a more clever operation that can do this in one go?
Kind Regards Jonathan