Pooling module wrappers #
These wrappers expose pooling specs as NNModuleSpecs.
Conventions:
- Channel-first images use shape
(C, H, W)at the spec level (.dim C (.dim H (.dim W .scalar))). MaxPool2DModuleSpecapplies the spatial max-pool independently per channel.AvgPool2DModuleSpecis provided for a single-channel 2D tensor; multi-channel usage typically maps it per channel in the same way as max-pool.
If you want a PyTorch mapping: nn.MaxPool2d / nn.AvgPool2d on a single (C,H,W) image (no
batch).
def
Spec.MaxPool2DModuleSpec
{α : Type}
[Context α]
{kH kW stride inH inW inC : ℕ}
{h1 : kH ≠ 0}
{h2 : kW ≠ 0}
{hStride : stride ≠ 0}
(m : MaxPool2DSpec kH kW stride h1 h2 hStride)
:
ModSpec.NNModuleSpec α (Shape.dim inC (Shape.dim inH (Shape.dim inW Shape.scalar)))
(Shape.dim inC (Shape.dim (poolOutDim inH kH stride 0) (Shape.dim (poolOutDim inW kW stride 0) Shape.scalar)))
MaxPool2D wrapper (channel-first, pool applied per channel).
Instances For
def
Spec.AvgPool2DModuleSpec
{α : Type}
[Context α]
{kH kW stride inH inW : ℕ}
{h1 : kH ≠ 0}
{h2 : kW ≠ 0}
{hStride : stride ≠ 0}
(m : AvgPool2DSpec kH kW stride h1 h2 hStride)
:
ModSpec.NNModuleSpec α (Shape.dim inH (Shape.dim inW Shape.scalar))
(Shape.dim (poolOutDim inH kH stride 0) (Shape.dim (poolOutDim inW kW stride 0) Shape.scalar))
AvgPool2D wrapper (2D tensor).