PyTorch 相当于 tf.dynamic_partition答案

【问题标题】：PyTorch equivalent of tf.dynamic_partitionPyTorch 相当于 tf.dynamic_partition
【发布时间】：2019-12-21 22:05:43
【问题描述】：

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

【问题讨论】：

标签： python tensorflow pytorch

【解决方案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])]

【讨论】：