如何从Pytorch的2D张量列表中获取列

时间:2020-09-09 11:30:59

标签: python pytorch

我有一个一维2D列表,其中包含张量:

[
 [tensor([-0.0705,  1.2019]), tensor([[0]]), tensor([-1.3865], dtype=torch.float64), tensor([-0.0744,  1.1880]), tensor([False])],
 [tensor([-0.0187,  1.3574]), tensor([[2]]), tensor([0.3373], dtype=torch.float64), tensor([-0.0221,  1.3473]), tensor([False])],
 [....] ]

外部列表包含64个小列表。一个小列表包含5个不同的张量元素。

我想获取诸如tensor([-0.0705, 1.2019])tensor([-0.0187, 1.3574])之类的内部列表的第一个元素,并创建诸如64x2之类的张量来馈入神经网络。

如何以最快的方式做到这一点?

谢谢

2 个答案:

答案 0 :(得分:1)

如何使用切片?


import torch
import numpy as np
x = [
 [torch.tensor([-0.0705,  1.2019]), torch.tensor([0]), torch.tensor([-1.3865], dtype=torch.float64), torch.tensor([-0.0744,  1.1880]), torch.tensor([False])],
 [torch.tensor([-0.0187,  1.3574]), torch.tensor([2]), torch.tensor([0.3373], dtype=torch.float64), torch.tensor([-0.0221,  1.3473]), torch.tensor([False])]]
x = list(map(lambda x:list(map(lambda z:z.tolist(), x)), x))
print(x)
x = np.array(x)[:, 0]
x = list(map(lambda z:torch.tensor(z), x))
print(x)

答案 1 :(得分:0)

使用list comprehension

[item[0] for item in your_list]

示例:

li = [[tensor([-0.0705,  1.2019]), tensor([[0]]), tensor([-1.3865], dtype=torch.float64), tensor([-0.0744,  1.1880]), tensor([False])],
     [tensor([-0.0187,  1.3574]), tensor([[2]]), tensor([0.3373], dtype=torch.float64), tensor([-0.0221,  1.3473]), tensor([False])]]

[item[0] for item in li]
[tensor([-0.0705,  1.2019]), tensor([-0.0187,  1.3574])]