CODEWITHSUNDEEP
pythonlistnumpypytorchtensor
popeye
popeye

Reputation: 886

How to extract tensors to numpy arrays or lists from a larger pytorch tensor

I have a list of pytorch tensors as shown below:

data = [[tensor([0, 0, 0]), tensor([1, 2, 3])],
        [tensor([0, 0, 0]), tensor([4, 5, 6])]]

Now this is just a sample data, the actual one is quite large but the structure is similar.

Question: I want to extract the tensor([1, 2, 3]), tensor([4, 5, 6]) i.e., the index 1 tensors from data to either a numpy array or a list in flattened form.

Expected Output:

out = array([1, 2, 3, 4, 5, 6])

OR

out = [1, 2, 3, 4, 5, 6]

I have tried several ways one including map function like:

map(lambda x: x[1].numpy(), data)

This gives:

[array([1, 2, 3]),
 array([4, 5, 6])]

And I'm unable to get the desired result with any other method I'm using.

Upvotes: 1

Views: 2001

Answers (2)

iacob
iacob

Reputation: 24181

You can convert a nested list of tensors to a tensor/numpy array with a nested stack:

data = np.stack([np.stack([d for d in d_]) for d_ in data])

You can then easily index this, and concatenate the output:

>>> np.concatenate(data[:,1])
array([[1, 2, 3],
       [4, 5, 6]])

Upvotes: 1

Natthaphon Hongcharoen
Natthaphon Hongcharoen

Reputation: 2430

OK, you can just do this.

out = np.concatenate(list(map(lambda x: x[1].numpy(), data)))

Upvotes: 2

Related Questions