сгруппированное среднее значение в pytorch

Question

У меня есть 2D-тензор:

samples = torch.Tensor([
                 [0.1, 0.1],    #-> group / class 1
                 [0.2, 0.2],    #-> group / class 2
                 [0.4, 0.4],    #-> group / class 2
                 [0.0, 0.0]     #-> group / class 0
          ])

и метка для каждого образца, соответствующего классу:

labels = torch.LongTensor([1, 2, 2, 0])

так len(samples) == len(labels). Теперь я хочу рассчитать среднее значение для каждого класса / ярлыка. Поскольку существует 3 класса (0, 1 и 2), конечный вектор должен иметь размерность [n_classes, samples.shape[1]] Таким образом, ожидаемое решение должно быть:

result == torch.Tensor([
                 [0.1, 0.1],
                 [0.3, 0.3], # -> mean of [0.2, 0.2] and [0.4, 0.4]
                 [0.0, 0.0]
          ])

Вопрос : Как это можно сделать на чистом pytorch (то есть, нет ничего тупого, чтобы я мог автоградироваться) и в идеале без циклов for?

Answer 1

Все, что вам нужно сделать, это сформировать матрицу mxn (m = num классов, n = num выборок), которая выберет подходящие веса и соответствующим образом масштабирует среднее. Затем вы можете выполнить умножение матриц между вновь сформированной матрицей и матрицей выборок.

Учитывая ваши метки, ваша матрица должна быть (каждая строка является номером класса, каждый класс - номером пробы и его весом):

[[0.0000, 0.0000, 0.0000, 1.0000],
 [1.0000, 0.0000, 0.0000, 0.0000],
 [0.0000, 0.5000, 0.5000, 0.0000]]

Который вы можете сформировать следующим образом:

M = torch.zeros(labels.max()+1, len(samples))
M[labels, torch.arange(4)] = 1
M = torch.nn.functional.normalize(M, p=1, dim=1)
torch.mm(M, samples)

Выход:

tensor([[0.0000, 0.0000],
        [0.1000, 0.1000],
        [0.3000, 0.3000]])

Обратите внимание, что средства вывода правильно отсортированы в порядке классов.

Answer 2

Повторное размещение здесь ответа от @ptrblck_de на форумах Pytorch

labels = labels.view(labels.size(0), 1).expand(-1, samples.size(1))

unique_labels, labels_count = labels.unique(dim=0, return_counts=True)

res = torch.zeros_like(unique_labels, dtype=torch.float).scatter_add_(0, labels, samples)
res = res / labels_count.float().unsqueeze(1)