Векторизация для применения функции к 3D-массиву

Question

Я пытаюсь применить функцию к тензору 3d-факела, в то время как функция применяется к 2-мерному тензору, который читается через ось 1 тензора 3d-факела.

Например, у меня есть тензор факела:форма (51, 128, 20100) (переменная с именем autoencode_logprob) и функция (rawid2sentence) выполняются на входе формы (51, 20100).

. Сейчас я написал код для запуска с использованием цикла naive for,цикл по одному с диапазоном (128).

Однако это слишком медленно.Ниже приведена важная часть кода.

autoencode_logprobs - это трехмерный тензор, и мне нужно применить функцию rawids2sentence вдоль ее второй оси.Любая помощь, чтобы векторизовать это?

for i in range(128):
    output_sent = self.dictionary.rawids2sentence(
        autoencode_logprobs[:, i].max(1)[
            1].data.cpu().numpy(),
        oov_dicts[i],
    )
    output_sent_encoding = ifst_model.encode([output_sent])

Answer 1

Поскольку я не знаю, что делает функция rawids2sentence или encode, я могу помочь вам выполнить операцию max.

В следующем утверждении

autoencode_logprobs[:, i].max(1)[1]

Вы определяете индекс максимальных значений вдоль dim=1 для каждого 51 x 20100 тензора. Итак, на выходе получается вектор размером 51.

Вы можете выполнить ту же операцию в своем полном тензоре формы 51 x 128 x 20100 и получить вывод в виде 128 x 51 тензор.

autoencode_logprobs.transpose(0, 1).max(2)[1] # 128 x 51

Итак, если ваши методы rawids2sentence или encode могут работать с пакетными входами, указанное выше изменение должно работать для вас без какой-либо петли.