Iterate through tensor dimensions

1

I have a 3D tensor and I want to iterate through the second dimension. The easiest approach is to use a for loop on the second dimension. My question is what is the most efficient/appropriate way to do it. For example:

t = torch.rand(2, 5, 3)

print(t)
tensor([[[0.1582, 0.7643, 0.1352],
         [0.7157, 0.4503, 0.7935],
         [0.9954, 0.6714, 0.0386],
         [0.7655, 0.2867, 0.0764],
         [0.2259, 0.3079, 0.8390]],

        [[0.2655, 0.2578, 0.8421],
         [0.7871, 0.7559, 0.2978],
         [0.4458, 0.6559, 0.9398],
         [0.6346, 0.4585, 0.3734],
         [0.8566, 0.1053, 0.0559]]])

for idx in range(t.shape[1]):
    print(t[:,idx,:])
tensor([[0.1582, 0.7643, 0.1352],
        [0.2655, 0.2578, 0.8421]])
tensor([[0.7157, 0.4503, 0.7935],
        [0.7871, 0.7559, 0.2978]])
tensor([[0.9954, 0.6714, 0.0386],
        [0.4458, 0.6559, 0.9398]])
tensor([[0.7655, 0.2867, 0.0764],
        [0.6346, 0.4585, 0.3734]])
tensor([[0.2259, 0.3079, 0.8390],
        [0.8566, 0.1053, 0.0559]])
2

If you want to iterate through a tensor dimension, using a for loop is a valid approach. However, it’s unclear if the loop is an optimal approach since you did not explain what your actual goal is besides iterating.

1 Like
3

@ptrblck thank you for your answer. Let my clarify what the tensor represents. The tensor is a batch of sequences, where each sequence element is features, thus the tensor is

[batch_size, sequence_length, features]

What I want, is to process each sequence element separately in batches, for example to pass them through a linear layer etc. That is why I want the for loop. I wonder if there is a better way of doing this.

4

nn.Linear layers accept inputs tensors with additional dimensions and will loop over these internally, as seen here:

import torch
import torch.nn as nn

# setup
batch_size, sequence_length, features = 2, 3, 4
x = torch.randn(batch_size, sequence_length, features)

lin = nn.Linear(features, features*2)

# loop approach
outputs = []
for i in range(sequence_length):
    x_ = x[:, i]
    out = lin(x_)
    outputs.append(out)
outputs = torch.stack(outputs, dim=1)
print(outputs.shape)
# torch.Size([2, 3, 8])

# single call
ref = lin(x)
print(ref.shape)
# torch.Size([2, 3, 8])

# compare
print((outputs - ref).abs().max())
# tensor(0., grad_fn=<MaxBackward1>)

so if you are using nn.Linear layers you would not have to iterate the input tensor manually.

1 Like
5

@ptrblck thank you once again for your detailed answer.

I am building something more complex, a version of an RNN model, specifically, a VRNN model (arxiv), where the computation of earlier sequence elements are used by the computations of the following elements, thus they cannot happen simultaneously.

Generally, I think that a for loop is the most basic way of doing things, which is the reason of the post.I guess that, in this case, this for loop is the best you can do.