Convolution module wrappers #
This file exposes conv specs as NNModuleSpecs.
The wrappers for Conv2D and ConvTranspose2D are consolidated here with their public names.
Conv2D #
def
Spec.Conv2DModuleSpec
{α : Type}
[Context α]
{inC outC kH kW stride padding inH inW : ℕ}
{h1 : inC ≠ 0}
{h2 : kH ≠ 0}
{h3 : kW ≠ 0}
(m : Conv2DSpec inC outC kH kW stride padding α h1 h2 h3)
:
ModSpec.NNModuleSpec α (Shape.dim inC (Shape.dim inH (Shape.dim inW Shape.scalar)))
(Shape.dim outC
(Shape.dim (Shape.slidingWindowOutDim inH kH stride padding)
(Shape.dim (Shape.slidingWindowOutDim inW kW stride padding) Shape.scalar)))
Wrap conv2d_spec as an NNModuleSpec, with the output shape computed in the type.
Instances For
ConvTranspose2D #
def
Spec.ConvTranspose2DModuleSpec
{α : Type}
[Context α]
{inC outC kH kW stride padding inH inW : ℕ}
{h1 : inC > 0}
{h2 : kH ≠ 0}
{h3 : kW ≠ 0}
(m : ConvTranspose2DSpec inC outC kH kW stride padding α h1 h2 h3)
:
ModSpec.NNModuleSpec α (Shape.dim inC (Shape.dim inH (Shape.dim inW Shape.scalar)))
(Shape.dim outC
(Shape.dim (convTransposeOutDim inH kH stride padding)
(Shape.dim (convTransposeOutDim inW kW stride padding) Shape.scalar)))
ConvTranspose2D wrapper as an NNModuleSpec (output shape encoded at the type level).