向量化将函数应用于三维数组

Question

我试图将一个函数应用于3d torch张量，同时将该函数应用于通过3d torch张量的轴1读取的2d张量

例如，我有一个(51, 128, 20100)形状的torch张量（一个名为autoencode_logprob的变量），函数（rawid2sentence）在(51, 20100)形状的输入上运行

现在，我编写了使用naivefor循环运行的代码，使用range（128）逐个循环

但是，它太慢了。下面是重要的代码部分

autoencodelogprobs是3d张量，我需要沿着它的第二个轴应用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]

为每个51 x 20100张量标识沿dim=1的最大值的索引。因此，输出是一个大小为51的向量

您可以在形状51 x 128 x 20100的全张量中执行相同的操作，并将输出作为128 x 51张量

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

因此，如果您的rawids2sentence或encode方法可以处理批输入，那么上面的更改应该可以在没有任何循环的情况下为您工作