如何在 JAX 中创建非全连接卷积层？

问题描述

我正在尝试在 JAX 中实现 LeNet5。我想我可以做到大部分，这里是架构供参考：

我想我知道如何实现前两层：

c1 = stax.Conv(out_chan=6,filter_shape=(5,5),strides=(1,1))
s2 = stax.parallel(
    stax.AvgPool(window_shape=(2,2),strides=(2,2)),stax.AvgPool(window_shape=(2,2))
)

然而，我被卡住的地方是 C3 层。问题是 C3 中的通道没有连接到 S2 中的每个通道。根据下表，C3 中的 16 个通道中的每一个仅连接到 S2 中的一个通道子集：

 0123456789ABCDEF
0x...xxx..xxxx.xx
1xx...xxx..xxxx.x
2xxx...xxx..x.xxx
3.xxx..xxxx..x.xx
4..xxx..xxxx.xx.x
5...xxx..xxxx.xxx

我无法从（相当简陋的）JAX/STAX docs 中看出如何设置像这样的稀疏连接的卷积层。

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）