Avinash Toppo
Avinash Toppo

Reputation: 366

How do i iterate over pytorch 2d tensors?

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?

Upvotes: 5

Views: 54712

Answers (3)

Montoya
Montoya

Reputation: 3049

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.

Upvotes: 4

kiranr
kiranr

Reputation: 2465

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])

Upvotes: 3

yatu
yatu

Reputation: 88305

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.

Upvotes: 12

Related Questions