pytorch - 从生成器创建 torch 张量

Question

我尝试从生成器构建一个张量，如下所示:

>>> torch.tensor(i**2 for i in range(10))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Could not infer dtype of generator

目前我只做:

>>> torch.tensor([i**2 for i in range(10)])
tensor([ 0,  1,  4,  9, 16, 25, 36, 49, 64, 81])

有没有办法避免需要这个中间列表？

Answer 1

正如@blue-phoenox 已经指出的那样，最好使用内置的 PyTorch 函数直接创建张量。但是如果您必须处理生成器，建议使用 numpy 作为中间阶段。由于 PyTorch 避免复制 numpy 数组，所以它应该是相当性能的(与简单的列表推导相比)

>>> import torch
>>> import numpy as np
>>> torch.from_numpy(np.fromiter((i**2 for i in range(10)), int))
tensor([ 0,  1,  4,  9, 16, 25, 36, 49, 64, 81])