équivalent PyTorch de tf.dynamic_partition

Question

Je dois utiliser un équivalent de tf.dynamic_partition dans PyTorch. Y a-t-il quelque chose avec une fonctionnalité similaire dans PyTorch ou dans une autre bibliothèque ou existe-t-il un moyen simple et astucieux de le coder pour PyTorch et de travailler rapidement ?

Answer 1

Oui, vous pouvez créer itérativement des masques booléens pour des indices uniques :

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)]

Et exemple tiré de tf.dynamic_partition documentation :

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)

Ce qui donne les mêmes résultats que l'exemple, à savoir :

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