python - Pytorch 张量平均在轴上吐出的块

Question

我正在寻找一种通过另一个掩码张量预定义的分块来执行张量平均的有效方法：

我们得到一个大小的张量[b_size,main_ax,sub_ax]和一个大小为掩码的二进制张量[b_size,main_ax]（仅包含 1 和 0 值），其中最后一个具有number_of_ones沿第二轴的恒定总和，这意味着每列的总和始终为number_of_ones。输出张量的大小为[b_size,number_of_ones,sub_ax]. 我们的目标是根据第二个定义的块对第一个张量进行平均，这意味着每个块都由一个序列定义，从 1 开始，包含 0，直到遇到下一个 1。

标签： pythonmatrixpytorchvectorizationtensor