816 lines
30 KiB
Go
816 lines
30 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package signal
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
)
|
||
|
||
// Pooling operations for 2-D feature maps. All functions take
|
||
// NCHW input (N, C, H, W) and return (N, C, H_out, W_out). Padding is
|
||
// a single int applied to both spatial dims; stride is a single int.
|
||
// Padding must stay below the kernel: a window entirely inside the
|
||
// padding holds no data, and a max would answer −Inf there, an
|
||
// average 0.
|
||
//
|
||
// The kernels hold two constraints fixed:
|
||
//
|
||
// - Element order. Average pooling adds its window elements in the
|
||
// (kernel height, kernel width) nesting and divides once, with the
|
||
// countIncludePad divisor chosen exactly as before; any other
|
||
// order or divisor changes bits. Max pooling is order free but
|
||
// keeps the same walk for a single shared code path, and a NaN in
|
||
// a window wins: v > best is false for NaN, so the comparison
|
||
// would silently drop it and answer the largest finite neighbour
|
||
// instead.
|
||
// - Parallel split. Work is distributed over output rows, which own
|
||
// disjoint output elements, so no split can influence a result.
|
||
//
|
||
// Non-float inputs widen once into a scratch payload (see
|
||
// widenFloats): FloatAt widens exactly, so the values
|
||
// compared and summed are bit-identical to the per-element accessor
|
||
// reads.
|
||
|
||
// parallelWorkBudget is the number of element updates a worker should
|
||
// carry before splitting a kernel pays for the worker's spawn. Every
|
||
// scheduling floor in the package is derived from it: an item body
|
||
// costing a few hundred nanoseconds still leaves a worker below the
|
||
// budget on a tiny input, where the split would spawn goroutines to do
|
||
// less work than the spawn itself costs.
|
||
const parallelWorkBudget = 1 << 12
|
||
|
||
// workFloorFor returns the smallest number of items a worker should
|
||
// carry when one item costs about workPerItem element updates. It is
|
||
// the floor argument engine.ParallelMin wants: below it the kernel runs
|
||
// whole on the calling goroutine, and above it the split is the one
|
||
// engine.Parallel computes, so the worker choice and the chunk
|
||
// boundaries of a split kernel are unchanged.
|
||
func workFloorFor(workPerItem int) int {
|
||
if workPerItem < 1 {
|
||
return parallelWorkBudget
|
||
}
|
||
return max(1, parallelWorkBudget/workPerItem)
|
||
}
|
||
|
||
// MaxPool2D returns the maximum over each kernel-sized window. A NaN
|
||
// in a window propagates: the maximum answers NaN.
|
||
func MaxPool2D(input *core.Array, kernel, stride, padding int) (*core.Array, error) {
|
||
return pool2D(input, kernel, stride, padding, true, false)
|
||
}
|
||
|
||
// AvgPool2D returns the average over each kernel-sized window. When
|
||
// countIncludePad is true the divisor is kH*kW; otherwise it's the
|
||
// number of non-padded elements in the window.
|
||
func AvgPool2D(input *core.Array, kernel, stride, padding int, countIncludePad bool) (*core.Array, error) {
|
||
return pool2D(input, kernel, stride, padding, false, countIncludePad)
|
||
}
|
||
|
||
// MaxPool1D returns the maximum over each kernel-sized window of a
|
||
// 3-D tensor (N, C, L). A NaN in a window propagates: the maximum
|
||
// answers NaN.
|
||
func MaxPool1D(input *core.Array, kernel, stride, padding int) (*core.Array, error) {
|
||
return pool1D(input, kernel, stride, padding, true, false)
|
||
}
|
||
|
||
// AvgPool1D returns the average over each kernel-sized window of a
|
||
// 3-D tensor (N, C, L). When countIncludePad is true the divisor is
|
||
// the full kernel size; otherwise it is the number of non-padded
|
||
// elements in the window.
|
||
func AvgPool1D(input *core.Array, kernel, stride, padding int, countIncludePad bool) (*core.Array, error) {
|
||
return pool1D(input, kernel, stride, padding, false, countIncludePad)
|
||
}
|
||
|
||
// MaxPool3D returns the maximum over each kernel-sized window of a
|
||
// 5-D tensor (N, C, D, H, W). A NaN in a window propagates: the
|
||
// maximum answers NaN.
|
||
func MaxPool3D(input *core.Array, kernel [3]int, stride [3]int, padding [3]int) (*core.Array, error) {
|
||
return pool3D(input, kernel, stride, padding, true, false)
|
||
}
|
||
|
||
// AvgPool3D returns the average over each kernel-sized window of a
|
||
// 5-D tensor (N, C, D, H, W). When countIncludePad is true the
|
||
// divisor is the full kernel size; otherwise it is the number of
|
||
// non-padded elements in the window.
|
||
func AvgPool3D(input *core.Array, kernel [3]int, stride [3]int, padding [3]int, countIncludePad bool) (*core.Array, error) {
|
||
return pool3D(input, kernel, stride, padding, false, countIncludePad)
|
||
}
|
||
|
||
// AdaptiveMaxPool2D pools to the requested output size by taking the
|
||
// maximum over each output window. Window o covers the input rows
|
||
// from floor(o·H_in/H_out) to ceil((o+1)·H_in/H_out), so every input
|
||
// sample lands in some window. A NaN in a window propagates: the
|
||
// maximum answers NaN.
|
||
func AdaptiveMaxPool2D(input *core.Array, outputH, outputW int) (*core.Array, error) {
|
||
if input.NDim() != 4 {
|
||
return nil, base.Errf("AdaptiveMaxPool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveMaxPool2D: complex arrays are not supported")
|
||
}
|
||
if outputH < 1 || outputW < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool2D: output size must be at least 1")
|
||
}
|
||
n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
|
||
// A zero spatial dimension leaves every window empty: max would
|
||
// answer -Inf, so it is an error like every other shape refusal.
|
||
if hIn < 1 || wIn < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool2D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputH, outputW}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output row; rows are disjoint. A row walks
|
||
// about one window height of input rows per window column.
|
||
items := n * c * outputH
|
||
floor := workFloorFor(max(1, (hIn+outputH-1)/outputH) * wIn)
|
||
engine.ParallelMin(items, floor, func(bs, be int) {
|
||
for item := bs; item < be; item++ {
|
||
oy := item % outputH
|
||
row := item / outputH
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*hIn*wIn + ch*hIn*wIn
|
||
yStart := oy * hIn / outputH
|
||
yEnd := ((oy+1)*hIn + outputH - 1) / outputH
|
||
for ox := range outputW {
|
||
xStart := ox * wIn / outputW
|
||
xEnd := ((ox+1)*wIn + outputW - 1) / outputW
|
||
best := math.Inf(-1)
|
||
for y := yStart; y < yEnd; y++ {
|
||
inRow := chanBase + y*wIn
|
||
for x := xStart; x < xEnd; x++ {
|
||
if v := inF[inRow+x]; v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
}
|
||
}
|
||
outF[batch*c*outputH*outputW+ch*outputH*outputW+oy*outputW+ox] = best
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// GlobalAvgPool2D returns AdaptiveAvgPool2D reduced to (1, 1), the
|
||
// global average pooling operation common in modern CNNs.
|
||
func GlobalAvgPool2D(input *core.Array) (*core.Array, error) {
|
||
return AdaptiveAvgPool2D(input, 1, 1)
|
||
}
|
||
|
||
// GlobalMaxPool2D returns the global maximum of a 4-D feature map.
|
||
func GlobalMaxPool2D(input *core.Array) (*core.Array, error) {
|
||
return AdaptiveMaxPool2D(input, 1, 1)
|
||
}
|
||
|
||
// GlobalAvgPool1D reduces a 3-D input (N, C, L) to shape (N, C, 1).
|
||
func GlobalAvgPool1D(input *core.Array) (*core.Array, error) {
|
||
if input.NDim() != 3 {
|
||
return nil, base.Errf("GlobalAvgPool1D: input must be 3-D (N, C, L), got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("GlobalAvgPool1D: complex arrays are not supported")
|
||
}
|
||
if input.Shape()[2] == 0 {
|
||
return nil, base.Errf("GlobalAvgPool1D: the spatial dimension must not be empty")
|
||
}
|
||
return pool1D(input, input.Shape()[2], 1, 0, false, false)
|
||
}
|
||
|
||
// GlobalMaxPool1D reduces a 3-D input (N, C, L) to shape (N, C, 1).
|
||
func GlobalMaxPool1D(input *core.Array) (*core.Array, error) {
|
||
if input.NDim() != 3 {
|
||
return nil, base.Errf("GlobalMaxPool1D: input must be 3-D (N, C, L), got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("GlobalMaxPool1D: complex arrays are not supported")
|
||
}
|
||
if input.Shape()[2] == 0 {
|
||
return nil, base.Errf("GlobalMaxPool1D: the spatial dimension must not be empty")
|
||
}
|
||
return pool1D(input, input.Shape()[2], 1, 0, true, false)
|
||
}
|
||
|
||
// GlobalAvgPool3D reduces a 5-D input (N, C, D, H, W) to (N, C, 1, 1, 1).
|
||
func GlobalAvgPool3D(input *core.Array) (*core.Array, error) {
|
||
if input.NDim() != 5 {
|
||
return nil, base.Errf("GlobalAvgPool3D: input must be 5-D (N, C, D, H, W), got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("GlobalAvgPool3D: complex arrays are not supported")
|
||
}
|
||
if input.Shape()[2] == 0 || input.Shape()[3] == 0 || input.Shape()[4] == 0 {
|
||
return nil, base.Errf("GlobalAvgPool3D: the spatial dimensions must not be empty")
|
||
}
|
||
return pool3D(input, [3]int{input.Shape()[2], input.Shape()[3], input.Shape()[4]}, [3]int{1, 1, 1}, [3]int{0, 0, 0}, false, false)
|
||
}
|
||
|
||
// GlobalMaxPool3D reduces a 5-D input (N, C, D, H, W) to (N, C, 1, 1, 1).
|
||
func GlobalMaxPool3D(input *core.Array) (*core.Array, error) {
|
||
if input.NDim() != 5 {
|
||
return nil, base.Errf("GlobalMaxPool3D: input must be 5-D (N, C, D, H, W), got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("GlobalMaxPool3D: complex arrays are not supported")
|
||
}
|
||
if input.Shape()[2] == 0 || input.Shape()[3] == 0 || input.Shape()[4] == 0 {
|
||
return nil, base.Errf("GlobalMaxPool3D: the spatial dimensions must not be empty")
|
||
}
|
||
return pool3D(input, [3]int{input.Shape()[2], input.Shape()[3], input.Shape()[4]}, [3]int{1, 1, 1}, [3]int{0, 0, 0}, true, false)
|
||
}
|
||
|
||
// AdaptiveMaxPool1D pools the input to the requested output size by
|
||
// taking the maximum over each output window. Window o covers the
|
||
// input positions from floor(o·L_in/L_out) to ceil((o+1)·L_in/L_out),
|
||
// so every input sample lands in some window. Input is (N, C, L). A
|
||
// NaN in a window propagates: the maximum answers NaN.
|
||
func AdaptiveMaxPool1D(input *core.Array, outputL int) (*core.Array, error) {
|
||
if input.NDim() != 3 {
|
||
return nil, base.Errf("AdaptiveMaxPool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveMaxPool1D: complex arrays are not supported")
|
||
}
|
||
if outputL < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool1D: output size must be at least 1")
|
||
}
|
||
n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
|
||
if lIn < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool1D: the spatial dimension must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputL}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per (batch, channel) row; rows are disjoint.
|
||
engine.ParallelMin(n*c, workFloorFor(lIn), func(bs, be int) {
|
||
for row := bs; row < be; row++ {
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*lIn + ch*lIn
|
||
for ol := range outputL {
|
||
start := ol * lIn / outputL
|
||
end := ((ol+1)*lIn + outputL - 1) / outputL
|
||
best := math.Inf(-1)
|
||
for l := start; l < end; l++ {
|
||
if v := inF[chanBase+l]; v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
}
|
||
outF[batch*c*outputL+ch*outputL+ol] = best
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// AdaptiveMaxPool3D pools the input to the requested output size by
|
||
// taking the maximum over each output window, with the same
|
||
// floor-start, ceil-end convention as AdaptiveMaxPool1D. Input is
|
||
// (N, C, D, H, W). A NaN in a window propagates: the maximum answers
|
||
// NaN.
|
||
func AdaptiveMaxPool3D(input *core.Array, outputD, outputH, outputW int) (*core.Array, error) {
|
||
if input.NDim() != 5 {
|
||
return nil, base.Errf("AdaptiveMaxPool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveMaxPool3D: complex arrays are not supported")
|
||
}
|
||
if outputD < 1 || outputH < 1 || outputW < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool3D: output sizes must be at least 1")
|
||
}
|
||
n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
|
||
if dIn < 1 || hIn < 1 || wIn < 1 {
|
||
return nil, base.Errf("AdaptiveMaxPool3D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputD, outputH, outputW}...)
|
||
strideHW := wIn
|
||
strideDHW := hIn * wIn
|
||
strideCDHW := dIn * hIn * wIn
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output depth-height row; rows are disjoint.
|
||
items := n * c * outputD * outputH
|
||
floor := workFloorFor(wIn * max(1, (hIn+outputH-1)/outputH) * max(1, (dIn+outputD-1)/outputD))
|
||
engine.ParallelMin(items, floor, func(bs, be int) {
|
||
for item := bs; item < be; item++ {
|
||
oh := item % outputH
|
||
t := item / outputH
|
||
od := t % outputD
|
||
t /= outputD
|
||
ch := t % c
|
||
batch := t / c
|
||
chanBase := batch*c*strideCDHW + ch*strideCDHW
|
||
dStart := od * dIn / outputD
|
||
dEnd := ((od+1)*dIn + outputD - 1) / outputD
|
||
hStart := oh * hIn / outputH
|
||
hEnd := ((oh+1)*hIn + outputH - 1) / outputH
|
||
for ow := range outputW {
|
||
wStart := ow * wIn / outputW
|
||
wEnd := ((ow+1)*wIn + outputW - 1) / outputW
|
||
best := math.Inf(-1)
|
||
for d := dStart; d < dEnd; d++ {
|
||
for h := hStart; h < hEnd; h++ {
|
||
inRow := chanBase + d*strideDHW + h*strideHW
|
||
for w := wStart; w < wEnd; w++ {
|
||
if v := inF[inRow+w]; v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
}
|
||
}
|
||
}
|
||
outF[batch*c*outputD*outputH*outputW+ch*outputD*outputH*outputW+od*outputH*outputW+oh*outputW+ow] = best
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// AdaptiveAvgPool1D pools to the requested output size by averaging
|
||
// each output window, with the same floor-start, ceil-end convention
|
||
// as AdaptiveMaxPool1D. Input is (N, C, L).
|
||
func AdaptiveAvgPool1D(input *core.Array, outputL int) (*core.Array, error) {
|
||
if input.NDim() != 3 {
|
||
return nil, base.Errf("AdaptiveAvgPool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveAvgPool1D: complex arrays are not supported")
|
||
}
|
||
if outputL < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool1D: output size must be at least 1")
|
||
}
|
||
n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
|
||
if lIn < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool1D: the spatial dimension must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputL}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per (batch, channel) row; rows are disjoint and
|
||
// every window sums its elements in ascending order.
|
||
engine.ParallelMin(n*c, workFloorFor(lIn), func(bs, be int) {
|
||
for row := bs; row < be; row++ {
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*lIn + ch*lIn
|
||
for ol := range outputL {
|
||
start := ol * lIn / outputL
|
||
end := ((ol+1)*lIn + outputL - 1) / outputL
|
||
var sum float64
|
||
for l := start; l < end; l++ {
|
||
sum += inF[chanBase+l]
|
||
}
|
||
outF[batch*c*outputL+ch*outputL+ol] = sum / float64(end-start)
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// AdaptiveAvgPool3D pools to the requested output size by averaging
|
||
// each output window, with the same floor-start, ceil-end convention
|
||
// as AdaptiveMaxPool1D. Input is (N, C, D, H, W).
|
||
func AdaptiveAvgPool3D(input *core.Array, outputD, outputH, outputW int) (*core.Array, error) {
|
||
if input.NDim() != 5 {
|
||
return nil, base.Errf("AdaptiveAvgPool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveAvgPool3D: complex arrays are not supported")
|
||
}
|
||
if outputD < 1 || outputH < 1 || outputW < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool3D: output sizes must be at least 1")
|
||
}
|
||
n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
|
||
if dIn < 1 || hIn < 1 || wIn < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool3D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputD, outputH, outputW}...)
|
||
strideHW := wIn
|
||
strideDHW := hIn * wIn
|
||
strideCDHW := dIn * hIn * wIn
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output element here: the 3-D windows share no
|
||
// row structure worth exploiting, and every element sums its own
|
||
// window in ascending order.
|
||
items := n * c * outputD * outputH * outputW
|
||
floor := workFloorFor(max(1, (dIn+outputD-1)/outputD) * max(1, (hIn+outputH-1)/outputH) * max(1, (wIn+outputW-1)/outputW))
|
||
engine.ParallelMin(items, floor, func(bs, be int) {
|
||
for item := bs; item < be; item++ {
|
||
ow := item % outputW
|
||
t := item / outputW
|
||
oh := t % outputH
|
||
t /= outputH
|
||
od := t % outputD
|
||
t /= outputD
|
||
ch := t % c
|
||
batch := t / c
|
||
chanBase := batch*c*strideCDHW + ch*strideCDHW
|
||
dStart := od * dIn / outputD
|
||
dEnd := ((od+1)*dIn + outputD - 1) / outputD
|
||
hStart := oh * hIn / outputH
|
||
hEnd := ((oh+1)*hIn + outputH - 1) / outputH
|
||
wStart := ow * wIn / outputW
|
||
wEnd := ((ow+1)*wIn + outputW - 1) / outputW
|
||
var sum float64
|
||
count := 0
|
||
for d := dStart; d < dEnd; d++ {
|
||
for h := hStart; h < hEnd; h++ {
|
||
inRow := chanBase + d*strideDHW + h*strideHW
|
||
for w := wStart; w < wEnd; w++ {
|
||
sum += inF[inRow+w]
|
||
count++
|
||
}
|
||
}
|
||
}
|
||
if count > 0 {
|
||
outF[batch*c*outputD*outputH*outputW+ch*outputD*outputH*outputW+od*outputH*outputW+oh*outputW+ow] = sum / float64(count)
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// AdaptiveAvgPool2D pools the input to the requested output size by
|
||
// averaging each output window, with the same floor-start, ceil-end
|
||
// convention as AdaptiveMaxPool2D. outputH and outputW must both be
|
||
// at least 1.
|
||
func AdaptiveAvgPool2D(input *core.Array, outputH, outputW int) (*core.Array, error) {
|
||
if input.NDim() != 4 {
|
||
return nil, base.Errf("AdaptiveAvgPool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("AdaptiveAvgPool2D: complex arrays are not supported")
|
||
}
|
||
if outputH < 1 || outputW < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool2D: output size must be at least 1, got %dx%d", outputH, outputW)
|
||
}
|
||
n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
|
||
// A zero spatial dimension leaves every window empty: avg would
|
||
// divide by zero, so it is an error like every other shape refusal.
|
||
if hIn < 1 || wIn < 1 {
|
||
return nil, base.Errf("AdaptiveAvgPool2D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
out := core.New(core.Float, []int{n, c, outputH, outputW}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output row; rows are disjoint and every window
|
||
// sums its elements in ascending order.
|
||
items := n * c * outputH
|
||
engine.ParallelMin(items, workFloorFor(max(1, (hIn+outputH-1)/outputH)*wIn), func(bs, be int) {
|
||
for item := bs; item < be; item++ {
|
||
oy := item % outputH
|
||
row := item / outputH
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*hIn*wIn + ch*hIn*wIn
|
||
// Window in the input: floor start, ceil end, so
|
||
// every input sample lands in some window.
|
||
yStart := oy * hIn / outputH
|
||
yEnd := ((oy+1)*hIn + outputH - 1) / outputH
|
||
for ox := range outputW {
|
||
xStart := ox * wIn / outputW
|
||
xEnd := ((ox+1)*wIn + outputW - 1) / outputW
|
||
var sum float64
|
||
for y := yStart; y < yEnd; y++ {
|
||
inRow := chanBase + y*wIn
|
||
for x := xStart; x < xEnd; x++ {
|
||
sum += inF[inRow+x]
|
||
}
|
||
}
|
||
count := float64((yEnd - yStart) * (xEnd - xStart))
|
||
off := batch*c*outputH*outputW + ch*outputH*outputW + oy*outputW + ox
|
||
outF[off] = sum / count
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// pool2D is the shared implementation behind MaxPool2D and AvgPool2D.
|
||
func pool2D(input *core.Array, kernel, stride, padding int, isMax, countIncludePad bool) (*core.Array, error) {
|
||
if input.NDim() != 4 {
|
||
return nil, base.Errf("Pool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("Pool2D: complex arrays are not supported")
|
||
}
|
||
if kernel < 1 {
|
||
return nil, base.Errf("Pool2D: kernel must be at least 1, got %d", kernel)
|
||
}
|
||
if stride < 1 {
|
||
return nil, base.Errf("Pool2D: stride must be at least 1, got %d", stride)
|
||
}
|
||
if padding < 0 {
|
||
return nil, base.Errf("Pool2D: padding must be non-negative, got %d", padding)
|
||
}
|
||
// Padding from the kernel upward leaves output windows entirely
|
||
// inside the padding (the first window covers −padding ..
|
||
// −padding+kernel−1): a max would answer −Inf and an average 0,
|
||
// so such configurations are refused, in the spirit of the
|
||
// common pooling implementations' padding bound.
|
||
if padding >= kernel {
|
||
return nil, base.Errf("Pool2D: padding %d must stay below the kernel %d, some windows would hold no data", padding, kernel)
|
||
}
|
||
n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
|
||
// Guard the numerators before the division: a negative numerator
|
||
// truncates toward zero and masquerades as a 1-wide output.
|
||
hNum := hIn + 2*padding - kernel
|
||
wNum := wIn + 2*padding - kernel
|
||
if hNum < 0 || wNum < 0 {
|
||
return nil, base.Errf("Pool2D: kernel %d with padding %d does not fit the %dx%d input", kernel, padding, hIn, wIn)
|
||
}
|
||
hOut := hNum/stride + 1
|
||
wOut := wNum/stride + 1
|
||
if hOut < 1 || wOut < 1 {
|
||
return nil, base.Errf("Pool2D: output is empty (H_out=%d, W_out=%d)", hOut, wOut)
|
||
}
|
||
out := core.New(core.Float, []int{n, c, hOut, wOut}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output row (batch, channel, oy). Rows are
|
||
// disjoint. The max and average variants walk the same windows in
|
||
// the same (kernel height, kernel width) order, but as two loop
|
||
// bodies: the max/sum choice leaves the innermost loop, which would
|
||
// otherwise pay for a perfectly predicted branch on every tap, and
|
||
// neither variant's arithmetic changes a bit.
|
||
items := n * c * hOut
|
||
engine.ParallelMin(items, workFloorFor(wOut*kernel*kernel), func(bs, be int) {
|
||
if isMax {
|
||
for item := bs; item < be; item++ {
|
||
oy := item % hOut
|
||
row := item / hOut
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*hIn*wIn + ch*hIn*wIn
|
||
for ox := range wOut {
|
||
best := math.Inf(-1)
|
||
for kh := range kernel {
|
||
iy := oy*stride + kh - padding
|
||
if iy < 0 || iy >= hIn {
|
||
continue
|
||
}
|
||
inRow := chanBase + iy*wIn
|
||
for kw := range kernel {
|
||
ix := ox*stride + kw - padding
|
||
if ix < 0 || ix >= wIn {
|
||
continue
|
||
}
|
||
if v := inF[inRow+ix]; v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
}
|
||
}
|
||
off := batch*c*hOut*wOut + ch*hOut*wOut + oy*wOut + ox
|
||
outF[off] = best
|
||
}
|
||
}
|
||
return
|
||
}
|
||
for item := bs; item < be; item++ {
|
||
oy := item % hOut
|
||
row := item / hOut
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*hIn*wIn + ch*hIn*wIn
|
||
for ox := range wOut {
|
||
var acc float64
|
||
var count int
|
||
for kh := range kernel {
|
||
iy := oy*stride + kh - padding
|
||
if iy < 0 || iy >= hIn {
|
||
continue
|
||
}
|
||
inRow := chanBase + iy*wIn
|
||
for kw := range kernel {
|
||
ix := ox*stride + kw - padding
|
||
if ix < 0 || ix >= wIn {
|
||
continue
|
||
}
|
||
acc += inF[inRow+ix]
|
||
count++
|
||
}
|
||
}
|
||
if countIncludePad {
|
||
acc /= float64(kernel * kernel)
|
||
} else if count > 0 {
|
||
acc /= float64(count)
|
||
}
|
||
off := batch*c*hOut*wOut + ch*hOut*wOut + oy*wOut + ox
|
||
outF[off] = acc
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// pool1D is the shared implementation behind MaxPool1D and AvgPool1D.
|
||
// countIncludePad is honoured only by the average variant.
|
||
func pool1D(input *core.Array, kernel, stride, padding int, isMax, countIncludePad bool) (*core.Array, error) {
|
||
if input.NDim() != 3 {
|
||
return nil, base.Errf("Pool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("Pool1D: complex arrays are not supported")
|
||
}
|
||
if kernel < 1 || stride < 1 || padding < 0 {
|
||
return nil, base.Errf("Pool1D: kernel/stride ≥1, padding ≥0")
|
||
}
|
||
// Padding from the kernel upward leaves output windows entirely
|
||
// inside the padding (the first window covers −padding ..
|
||
// −padding+kernel−1): a max would answer −Inf and an average 0,
|
||
// so such configurations are refused, in the spirit of the
|
||
// common pooling implementations' padding bound.
|
||
if padding >= kernel {
|
||
return nil, base.Errf("Pool1D: padding %d must stay below the kernel %d, some windows would hold no data", padding, kernel)
|
||
}
|
||
n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
|
||
lNum := lIn + 2*padding - kernel
|
||
if lNum < 0 {
|
||
return nil, base.Errf("Pool1D: kernel %d with padding %d does not fit the length-%d input", kernel, padding, lIn)
|
||
}
|
||
lOut := lNum/stride + 1
|
||
if lOut < 1 {
|
||
return nil, base.Errf("Pool1D: output is empty (L_out=%d)", lOut)
|
||
}
|
||
out := core.New(core.Float, []int{n, c, lOut}...)
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per (batch, channel) row; rows are disjoint and
|
||
// every window walks its elements in ascending order.
|
||
engine.ParallelMin(n*c, workFloorFor(lOut*kernel), func(bs, be int) {
|
||
for row := bs; row < be; row++ {
|
||
ch := row % c
|
||
batch := row / c
|
||
chanBase := batch*c*lIn + ch*lIn
|
||
for ol := range lOut {
|
||
var acc float64
|
||
var best float64
|
||
if isMax {
|
||
best = math.Inf(-1)
|
||
}
|
||
var count int
|
||
for kl := range kernel {
|
||
il := ol*stride + kl - padding
|
||
if il < 0 || il >= lIn {
|
||
continue
|
||
}
|
||
v := inF[chanBase+il]
|
||
if isMax {
|
||
if v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
} else {
|
||
acc += v
|
||
}
|
||
count++
|
||
}
|
||
if isMax {
|
||
outF[batch*c*lOut+ch*lOut+ol] = best
|
||
} else if countIncludePad {
|
||
outF[batch*c*lOut+ch*lOut+ol] = acc / float64(kernel)
|
||
} else if count > 0 {
|
||
outF[batch*c*lOut+ch*lOut+ol] = acc / float64(count)
|
||
}
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// pool3D is the shared implementation behind MaxPool3D and AvgPool3D.
|
||
// countIncludePad is honoured only by the average variant.
|
||
func pool3D(input *core.Array, kernel, stride, padding [3]int, isMax, countIncludePad bool) (*core.Array, error) {
|
||
if input.NDim() != 5 {
|
||
return nil, base.Errf("Pool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
|
||
}
|
||
if input.Dtype() == core.Complex {
|
||
return nil, base.Errf("Pool3D: complex arrays are not supported")
|
||
}
|
||
for d := range kernel {
|
||
if kernel[d] < 1 || stride[d] < 1 || padding[d] < 0 {
|
||
return nil, base.Errf("Pool3D: kernel/stride ≥1, padding ≥0")
|
||
}
|
||
// Padding from the kernel upward leaves output windows
|
||
// entirely inside the padding: a max would answer −Inf and
|
||
// an average 0, so such configurations are refused, in the
|
||
// spirit of the common pooling implementations' padding
|
||
// bound.
|
||
if padding[d] >= kernel[d] {
|
||
return nil, base.Errf("Pool3D: padding %d must stay below the kernel %d in dimension %d, some windows would hold no data", padding[d], kernel[d], d)
|
||
}
|
||
}
|
||
n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
|
||
// Guard the numerators before the division: a negative numerator
|
||
// truncates toward zero and masquerades as a 1-wide output.
|
||
dNum := dIn + 2*padding[0] - kernel[0]
|
||
hNum := hIn + 2*padding[1] - kernel[1]
|
||
wNum := wIn + 2*padding[2] - kernel[2]
|
||
if dNum < 0 || hNum < 0 || wNum < 0 {
|
||
return nil, base.Errf("Pool3D: kernel %v with padding %v does not fit the input", kernel, padding)
|
||
}
|
||
dOut := dNum/stride[0] + 1
|
||
hOut := hNum/stride[1] + 1
|
||
wOut := wNum/stride[2] + 1
|
||
if dOut < 1 || hOut < 1 || wOut < 1 {
|
||
return nil, base.Errf("Pool3D: output is empty")
|
||
}
|
||
out := core.New(core.Float, []int{n, c, dOut, hOut, wOut}...)
|
||
strideHW := wIn
|
||
strideDHW := hIn * wIn
|
||
strideCDHW := dIn * hIn * wIn
|
||
inF := input.RawFloats()
|
||
if inF == nil || input.Strided() {
|
||
inF = widenFloats(input)
|
||
}
|
||
outF := out.RawFloats()
|
||
// One work item per output depth-height row (batch, channel, od,
|
||
// oh); rows are disjoint and each window walks its elements in the
|
||
// (kernel depth, kernel height, kernel width) order.
|
||
items := n * c * dOut * hOut
|
||
engine.ParallelMin(items, workFloorFor(wOut*kernel[0]*kernel[1]*kernel[2]), func(bs, be int) {
|
||
for item := bs; item < be; item++ {
|
||
oh := item % hOut
|
||
t := item / hOut
|
||
od := t % dOut
|
||
t /= dOut
|
||
ch := t % c
|
||
batch := t / c
|
||
chanBase := batch*c*strideCDHW + ch*strideCDHW
|
||
for ow := range wOut {
|
||
var acc float64
|
||
var best float64
|
||
if isMax {
|
||
best = math.Inf(-1)
|
||
}
|
||
var count int
|
||
for kd := range kernel[0] {
|
||
id := od*stride[0] + kd - padding[0]
|
||
if id < 0 || id >= dIn {
|
||
continue
|
||
}
|
||
inDepth := chanBase + id*strideDHW
|
||
for kh := range kernel[1] {
|
||
ih := oh*stride[1] + kh - padding[1]
|
||
if ih < 0 || ih >= hIn {
|
||
continue
|
||
}
|
||
inRow := inDepth + ih*strideHW
|
||
for kw := range kernel[2] {
|
||
iw := ow*stride[2] + kw - padding[2]
|
||
if iw < 0 || iw >= wIn {
|
||
continue
|
||
}
|
||
v := inF[inRow+iw]
|
||
if isMax {
|
||
if v > best || math.IsNaN(v) {
|
||
best = v
|
||
}
|
||
} else {
|
||
acc += v
|
||
}
|
||
count++
|
||
}
|
||
}
|
||
}
|
||
off := batch*c*dOut*hOut*wOut + ch*dOut*hOut*wOut + od*hOut*wOut + oh*wOut + ow
|
||
if isMax {
|
||
outF[off] = best
|
||
} else if countIncludePad {
|
||
outF[off] = acc / float64(kernel[0]*kernel[1]*kernel[2])
|
||
} else if count > 0 {
|
||
outF[off] = acc / float64(count)
|
||
}
|
||
}
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|