ResNet as an NNModuleSpec #
The ResNet spec model provides ResNetSpec.forward and related gradients.
This file adds a module wrapper for composition and export tooling.
def
Spec.ResNetModuleSpec
{α : Type}
[Context α]
{cfg : ResNetConfig}
{inputChannels numClasses inH inW : ℕ}
(h1 : inputChannels ≠ 0)
(h2 : numClasses ≠ 0)
(hCfg : cfg.WF)
(resnet : ResNetSpec cfg α inputChannels numClasses h1 h2 hCfg)
(h3 : inH ≠ 0)
(h4 : inW ≠ 0)
:
ModSpec.NNModuleSpec α (Shape.dim inputChannels (Shape.dim inH (Shape.dim inW Shape.scalar)))
(Shape.dim numClasses Shape.scalar)
Wrap ResNetSpec.forward as an NNModuleSpec.