Question

我需要在PyTorch中使用等效的tf.dynamic_partition。 PyTorch或其他库中是否有任何具有类似功能的东西，还是有一种简单聪明的方法为PyTorch编写代码并快速工作？

Answer 1

是的，您可以迭代地为唯一索引创建布尔掩码：

import torch


# partitions should be integer-like types
def dynamic_partition(
    data: torch.Tensor, partitions: torch.Tensor, num_partitions=None
):
    assert len(partitions.shape) == 1, "Only one dimensional partitions supported"
    assert (
        data.shape[0] == partitions.shape[0]
    ), "Partitions requires the same size as data"

    if num_partitions is None:
        num_partitions = max(torch.unique(partitions))

    return [data[partitions == i] for i in range(num_partitions)]

摘自tf.dynamic_partition文档的示例：

partitions = torch.tensor([0, 0, 1, 1, 0])
num_partitions = 2
data = torch.tensor([10, 20, 30, 40, 50])

dynamic_partition(data, partitions, num_partitions)

给出的结果与示例相同，即：

[torch.tensor([10, 20, 50]), torch.tensor([30, 40])]

PyTorch等效于tf.dynamic_partition

1 个答案: