Python PyTorch cartesian_prod用法及代碼示例

本文簡要介紹python語言中 torch.cartesian_prod 的用法。

用法:

torch.cartesian_prod(*tensors)

參數：

*tensors-任意數量的一維張量。

返回：

一個張量相當於把所有輸入的張量轉成列表，對這些列表做itertools.product，最後把得到的列表轉成張量。

返回類型：

Tensor

對給定的張量序列進行笛卡爾積。該行為類似於 python 的 itertools.product 。

例子：

>>> a = [1, 2, 3]
>>> b = [4, 5]
>>> list(itertools.product(a, b))
[(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
>>> tensor_a = torch.tensor(a)
>>> tensor_b = torch.tensor(b)
>>> torch.cartesian_prod(tensor_a, tensor_b)
tensor([[1, 4],
        [1, 5],
        [2, 4],
        [2, 5],
        [3, 4],
        [3, 5]])