How do i iterate over pytorch 2d tensors?

Question

X = np.array([
    [-2,4,-1],
    [4,1,-1],
    [1, 6, -1],
    [2, 4, -1],
    [6, 2, -1],
    ])

for epoch in range(1,epochs):
    for i, x in enumerate(X):

X = torch.tensor([
    [-2,4,-1],
    [4,1,-1],
    [1, 6, -1],
    [2, 4, -1],
    [6, 2, -1],
    ])

The looping was fine when it was an numpy array. But i want to work with pytorch tensors so what's alternative to enumerate or How can i loop through the above tensor in the 2nd line?

Solution 1

enumerate expects an iterable, so it works just fine with pytorch tensors too:

X = torch.tensor([
    [-2,4,-1],
    [4,1,-1],
    [1, 6, -1],
    [2, 4, -1],
    [6, 2, -1],
    ])

for i, x in enumerate(X):
    print(x)
    tensor([-2,  4, -1])
    tensor([ 4,  1, -1])
    tensor([ 1,  6, -1])
    tensor([ 2,  4, -1])
    tensor([ 6,  2, -1])

If you want to iterate over the underlying arrays:

for i, x in enumerate(X.numpy()):
    print(x)
    [-2  4 -1]
    [ 4  1 -1]
    [ 1  6 -1]
    [ 2  4 -1]
    [ 6  2 -1]

Do note however that pytorch's underlying data structure are numpy arrays, so you probably want to avoid looping over the tensor as you're doing, and should be looking at vectorising any operations either via pytorch or numpy.

Solution 2

Iterating pytorch tensor or a numpy array is significantly slower than iterating a list.

Convert your tensor to a list and iterate over it:

l = tens.tolist()

detach() is needed if you need to detach your tensor from a computation graph:

l = tens.detach().tolist()

Alternatively, you might use numpy array and use some of its fast functions on each row in your 2d array in order to get the values that you need from that row.

Solution 3

x = torch.tensor([ 
                 [-2,4,-1],
                 [4,1,-1],
                 [1, 6, -1],
                 [2, 4, -1],
                 [6, 2, -1],
                          ])

for i in x:
    print(i)

output :

tensor([-2,  4, -1])
tensor([ 4,  1, -1])
tensor([ 1,  6, -1])
tensor([ 2,  4, -1])
tensor([ 6,  2, -1])

Comments

[yatu]

That question is about aggregating multiple iterables (in this case tensors) at once @AragonS - hence zip was suggested