Самый быстрый способ разбить тензоры на группы по две и для каждой группы случайным образом назначить один из них 0, а другой - 1

Question

У меня большой тензор с размером (128, 64, 1, 1). Для каждых двух маленьких тензоров размера (1, 1) в этом большом тензоре я хочу случайным образом присвоить одному значение 1, а другому - 0.

Мой следующий код работает, но он очень медленный.

# a is a big tensor of size (128, 64, 1, 1)
for i in range(a.size()[0]):
    j = 0
    while j < a.size()[1] - 1:
        r = int(torch.randint(0, 2, (1,)))
        a[i][j + r] = 1
        a[i][j + 1 - r] = 0
        j += 2

Есть ли способ избежать циклов for и выполнять каждое назначение параллельно? Спасибо.

Answer 1

Я думаю, что это может быть быстрее:

import torch
x = torch.randn(128, 32, 2, 1, 1)
y = x.max(dim=2, keepdim=True)[0]
z = (y-x) > 0
r = torch.flatten(z, 1, 2)

Идея состоит в том, что вы разбиваете свое измерение 64 на 2, выбираете случайные значения и выбираете одно из минимальных значений равным 1, а другое равным 0, а затем рекомбинируете 32 * 2 в размер 64.