带有Conv2d的批量维度中的Tensorflow switch case

Question

假设我有一个形状为 [n, h, w, c] 的批次以及一个范围为 0-9 和 10 Conv2D convs 的 n 个索引列表，我想将其应用于数据取决于列表中的索引。索引列表随每个批次而变化。

例如。输入 x，批量大小 4 和索引 l=[1,5,1,9] 我想计算 [convs[l[0]](x[0]), convs[l[1]](x[1]), convs[l[2]](x[2]), convs[l[3]](x[3])]

一个简单的解决方案是计算每个组合并根据l 收集。但是，这需要 10 倍的内存量。这个问题有更好的解决方案吗？

Answer 1

一个“hacky”解决方案是将输入的维度从[n, h, w, c] 扩展到[1, n, h, w, c]，然后使用Conv3D 代替内核形状[1, x, y]

如果您有单独定义的权重（也可以使用layer.weights 获得权重），您可以类似地将它们堆叠在第 0 维中并通过tf.nn.conv3d 使用它们。