切片保存在列表中的 Pytorch 张量

Question

我有以下代码段来生成随机样本。生成的 samples 是一个 list，其中列表的每个条目都是一个 tensor。每个张量都有两个元素。我想从列表中的所有张量中提取第一个元素；并从列表中的所有张量中提取第二个元素。如何进行这种张量切片操作

import torch
import pyro.distributions as dist
num_samples = 250
# note that both covariance matrices are diagonal
mu1 = torch.tensor([0., 5.])
sig1 = torch.tensor([[2., 0.], [0., 3.]])
dist1 = dist.MultivariateNormal(mu1, sig1)
samples1 = [pyro.sample('samples1', dist1) for _ in range(num_samples)]

samples1

Answer 1

我建议 torch.cat 具有列表理解力：

col1 = torch.cat([t[0] for t in samples1])
col2 = torch.cat([t[1] for t in samples1])

torch.cat 的文档：https://pytorch.org/docs/stable/generated/torch.cat.html

或者

您可以使用 torch.stack 将一维张量列表转换为一个大的二维张量，然后进行普通切片：

samples1_t = torch.stack(samples1)
col1 = samples1_t[:, 0]  # : means all rows
col2 = samples1_t[:, 1]

torch.stack 的文档：https://pytorch.org/docs/stable/generated/torch.stack.html

Answer 2

我应该提到 PyTorch 张量开箱即用，这意味着您可以将第一个轴解包为多个变量，而无需额外考虑。这里 torch.stack will output a tensor of shape (rows, cols), we just need to transpose it to (cols, rows) and unpack:

>>> c1, c2 = torch.stack(samples1).T

所以你得到 c1 和 c2 形 (rows,):

>>> c1
tensor([0.6433, 0.4667, 0.6811, 0.2006, 0.6623, 0.7033])

>>> c2
tensor([0.2963, 0.2335, 0.6803, 0.1575, 0.9420, 0.6963])

Answer 3

建议 .stack() 或 .cat() 的其他答案从 PyTorch 的角度来看完全没问题。

但是，由于问题的上下文涉及到pyro，请允许我添加以下内容：

因为你在做 IID 样本

[pyro.sample('samples1', dist1) for _ in range(num_samples)]

使用 pyro 的更好方法是

dist1 = dist.MultivariateNormal(mu1, sig1).expand([num_samples])

这告诉 pyro 分布是 批处理的 ，批量大小为 num_samples。从中抽样将产生

>> dist1.sample()
tensor([[-0.8712,  6.6087],
    [ 1.6076, -0.2939],
    [ 1.4526,  6.1777],
    ...
    [-0.0168,  7.5085],
    [-1.6382,  2.1878]])

现在很容易解决你原来的问题。像这样切片就可以了

samples = dist1.sample()
samples[:, 0]  # all first elements
samples[:, 1]  # all second elements

切片保存在列表中的 Pytorch 张量

slice Pytorch tensors which are saved in a list

pyro

pytorch

tensor