Tensorflow，使用自定义过滤器进行卷积

Question

我有两个张量A和B。

A代表张量，尺寸为N xH xW x 512。 更具体地说，A是卷积层的输出。 N是批次大小。 H和W是张量的高度和宽度A。512是张量的通道数A。

B表示张量，其尺寸为N x 1 x 1 x 512 x 256。 更具体地说，它是卷积层的N个内核，每个内核的尺寸为1 x 1 x 512 x256。这里，512和256分别表示输入通道和输出通道。

我想做的是用张量A对每个内核进行卷积，例如

all_out = []
For i in range(N):
     kernel = B[i]
     a_batch = A[i]
     a_out = conv2d(a_batch, kernel, trainable=False) #pseudo convolution function as tf.nn.conv2d does not accept trainable argument.
     all_out.append(a_out)
all_out = some_func(all_out) # convert to N x H x W x 256 tensor

在重塑两个张量A和B之后，我尝试使用tf.nn.depthwise_conv2d。但是，它很慢并且占用太多内存空间。

是否有一种有效的方法来做到这一点？有什么建议吗？