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

Question

我有一个包含张量的 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 这样的张量来馈送我的神经网络。

我怎样才能以最快的方式做到这一点？

谢谢

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

Answer 2

使用列表理解

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