如何从更大的 pytorch 张量中提取张量到 numpy 数组或列表

Question

我有一个 pytorch 张量列表，如下所示：

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

现在这只是一个示例数据，实际的数据很大，但结构相似。

问题： 我想提取 tensor([1, 2, 3])、tensor([4, 5, 6]) 即索引 1 张量从 data 到 numpy 数组或扁平形式的列表。

预期输出：

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

或

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

我尝试了几种方法，其中一种包括 map 功能，例如：

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

这给出：

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

而且我无法使用我正在使用的任何其他方法获得所需的结果。

Answer 1

好的，你可以这样做。

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

Answer 2

您可以将嵌套的张量列表转换为具有嵌套堆栈的 tensor/numpy 数组：

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

然后您可以轻松地对其进行索引，并连接输出：

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

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