PyTorch创建范围组合

eoigrqb6 于 2022-12-26 发布在其他

关注(0)|答案(2)|浏览(205)

有没有一种方法可以很容易地嵌套arange，从而在PyTorch中创建两个范围的所有组合？例如：

x = torch.arange(2, 4)
y = torch.arange(0, 3)

something(x, y)
# should be [[2,0], [2,1], [2,3], [3,0], [3,1], [3,2]]

也就是说，它的功能与下面的python代码相同：

l = []
for x in range(2, 4):
   for y in range(0, 3):
       l.append([x, y])

我们可以改变range(x,y)。

pytorch

来源：https://stackoverflow.com/questions/74904373/pytorch-create-combination-of-ranges

2条答案

按热度按时间

3hvapo4f1#

幸运的是，有一个内置函数可以满足您的需要，它就是torch.cartesian_prod()函数。
下面是一个例子：

x = torch.arange(2, 4)
y = torch.arange(0, 3)
l = torch.cartesian_prod(x,y)

赞(0）回复(0）举报 2022-12-26

4uqofj5v2#

请考虑以下几点：

x = torch.arange(2, 4)
y = torch.arange(0, 3)

m,n = len(x),len(y)
res = torch.stack([
    x[:,None].broadcast_to(m,n),
    y.broadcast_to(m,n)]
    ).permute(1,2,0).reshape([m*n,2])

结果：

tensor([[2, 0],
        [2, 1],
        [2, 2],
        [3, 0],
        [3, 1],
        [3, 2]])

赞(0）回复(0）举报 2022-12-26

我来回答

PyTorch创建范围组合

2条答案

相关问题

热门标签

最新问答