为什么Trax会为已声明的Serial层内子层自动额外创建Serial层？

问题原因 #

你遇到的嵌套Serial层是旧版Trax的已知设计问题：低于1.3.6版本的Trax中，tl.Relu()这类激活函数的内部实现本身就是用Serial层封装的，所以你传入外层Serial的Relu()实例本身就是一个Serial对象，打印时就会出现嵌套结构，这个嵌套完全不影响模型的训练、推理效率和精度，只是结构展示不符合预期。

解决方案 #

• 升级Trax到最新正式版本即可，官方在后续版本中已经重构了激活函数的实现，所有激活层改为直接继承Layer基类，不再额外套Serial层，升级后打印模型就会和你的预期输出一致。
• 如果你暂时无法升级Trax版本，不需要做额外处理，嵌套结构不会对模型效果产生任何负面影响。

验证代码（升级后运行） #

model = tl.Serial(
        tl.Dense(n_units=512),
        tl.Relu()
        )
print(model)

预期输出 #

Serial[
  Dense_512
  Relu
]

内容的提问来源于stack exchange，提问作者Aditya Khanke