Pooling Modules #
The wrappers in this file preserve a leading channel dimension and pool over an arbitrary vector of spatial dimensions. Padding, stride, and window extents are independent on every axis.
def
Spec.Module.maxPool
{α : Type}
[TorchLean.Storage α]
[Context α]
{d C : ℕ}
{inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]}
{hKernel : ∀ (i : Fin d), kernel.getScalar i ≠ 0}
{hStride : ∀ (i : Fin d), stride.getScalar i ≠ 0}
(m : MaxPoolSpec d kernel stride padding hKernel hStride)
:
Module α (Shape.ofList (C :: inSpatial.to (List ℕ)))
(Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ)))
Wrap arbitrary-rank channels-first max pooling as a Spec.Module.
Instances For
def
Spec.Module.avgPool
{α : Type}
[TorchLean.Storage α]
[Context α]
{d C : ℕ}
{inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]}
{hKernel : ∀ (i : Fin d), kernel.getScalar i ≠ 0}
{hStride : ∀ (i : Fin d), stride.getScalar i ≠ 0}
(m : AvgPoolSpec d kernel stride padding hKernel hStride)
:
Module α (Shape.ofList (C :: inSpatial.to (List ℕ)))
(Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ)))
Wrap arbitrary-rank channels-first average pooling as a Spec.Module.