#pragma once
// SafeC Standard Library — trainable convolutional networks on std::ml's
// Autograd::get core.
//
// cnn.h runs convolutions forward on FeatureMaps; this file makes them
// differentiable on Tensors so they train with Tensor::backward and the
// optimizers in nn.h, alongside Mlp. Images travel as rows, the same as
// every other model in std::ml: a batch is a [batch, C*H*W] tensor, each
// row one image in CHW order (row[c*H*W + h*W + w]). The geometry (C, H,
// W) is passed to every op explicitly, since a 2D tensor does not carry it.
//
// * tensor_conv2d / tensor_conv2d_ex — convolution with stride, zero
// padding, dilation and groups (depthwise when groups == channels),
// weight [outC, (inC/groups)*kH*kW] (a PyTorch Conv2d weight,
// flattened), bias [outC].
// * tensor_conv_transpose2d / _ex — its transpose (the gradient of conv2d
// with respect to its input, run forward), weight [inC,
// (outC/groups)*kH*kW] (a PyTorch ConvTranspose2d weight, flattened),
// bias [outC], with output padding to reach every output size. The
// usual learned upsampler of decoders and generators.
// * Pooling: max / avg with any window, stride, padding and ceil mode;
// global average / max.
// * Resizing: nearest (integer factor, fast path), and nearest / bilinear
// / bicubic to any size, with or without aligned corners.
// * Normalization: batch norm with running statistics; group norm (layer
// norm and instance norm as its two ends); spectral norm of a weight.
// * Dropout, elementwise or by whole channels.
// * ConvNet — a stack of those layers plus fully connected ones, with
// saved activations for residual (add) and U-Net (concat) skips, built
// layer by layer and parallel to Mlp. autoencoder.h (ConvAutoencoder,
// ConvVae) and gan.h (Dcgan) take ConvNets as their encoders,
// decoders, generators and discriminators.
//
// Output sides along each axis: conv2d (dim + 2*pad - dil*(k-1) - 1)/stride
// + 1 (rounded down, as PyTorch); conv_transpose2d (dim - 1)*stride - 2*pad
// + dil*(k-1) + outPad + 1; pooling as conv2d with dilation 1, rounded up
// in ceil mode. Both convolutions compute through im2col + sgemm (3x3
// stride-1 forward passes with 8+ channels through cnn.h's Winograd
// kernel); a large enough batch is split across std/sync/parallel.h's
// thread pool, one chunk of images per thread, with results independent of
// scheduling. Pooling, resizing and the norms split large inputs across the
// pool by plane, row or channel the same way.
#include <std/ml/tensor.h>
#include <std/ml/nn.h>
#include <std/collections/vec.h>
namespace std {
// ── Convolution ──────────────────────────────────────────────────────────────
// Every field per axis; Conv2dSpec::new fills a square kernel with dilation 1,
// 1 group and no output padding. inC and outC must divide by groups;
// outPad (transposed convolution only) must be below stride or dilation.
struct Conv2dSpec {
unsigned long kH;
unsigned long kW;
unsigned long strideH;
unsigned long strideW;
unsigned long padH;
unsigned long padW;
unsigned long dilationH;
unsigned long dilationW;
unsigned long groups;
unsigned long outPadH;
unsigned long outPadW;
static struct Conv2dSpec new(unsigned long kernel, unsigned long stride, unsigned long padding);
};
// x is [batch, inC*inH*inW] (a 1D tensor is one image). Each returns
// [batch, outC*outH*outW] and records its backward edge to x, weight and
// bias, like tensor_add does. outC is bias's size.
&FloatTensor tensor_conv2d_ex(const &FloatTensor x, const &FloatTensor weight, const &FloatTensor bias,
unsigned long inC, unsigned long inH, unsigned long inW, const &Conv2dSpec spec);
&FloatTensor tensor_conv_transpose2d_ex(const &FloatTensor x, const &FloatTensor weight, const &FloatTensor bias,
unsigned long inC, unsigned long inH, unsigned long inW, const &Conv2dSpec spec);
// The common case: kH x kW kernel, one stride and padding for both axes.
&FloatTensor tensor_conv2d(const &FloatTensor x, const &FloatTensor weight, const &FloatTensor bias,
unsigned long inC, unsigned long inH, unsigned long inW,
unsigned long kH, unsigned long kW,
unsigned long stride, unsigned long padding);
&FloatTensor tensor_conv_transpose2d(const &FloatTensor x, const &FloatTensor weight, const &FloatTensor bias,
unsigned long inC, unsigned long inH, unsigned long inW,
unsigned long kH, unsigned long kW,
unsigned long stride, unsigned long padding);
// Output side along one axis.
unsigned long conv_out_side(unsigned long dim, unsigned long kernel, unsigned long stride, unsigned long padding);
unsigned long conv_transpose_out_side(unsigned long dim, unsigned long kernel, unsigned long stride, unsigned long padding);
unsigned long conv_out_side_ex(unsigned long dim, unsigned long kernel, unsigned long stride,
unsigned long padding, unsigned long dilation);
unsigned long conv_transpose_out_side_ex(unsigned long dim, unsigned long kernel, unsigned long stride,
unsigned long padding, unsigned long dilation, unsigned long outPad);
// ── Pooling ──────────────────────────────────────────────────────────────────
// Windows are clipped to the plane: max pooling ignores the padding (as
// PyTorch's MaxPool2d does) and sends each window's gradient to its
// (first) largest input. Avg pooling with countIncludePad divides by the
// window's extent inside the padded plane (PyTorch's AvgPool2d default),
// else by the positions inside the plane. padH / padW at most half the
// window. ceilMode rounds the output side up, as PyTorch's ceil_mode (a
// last window must still start inside the input or its left padding).
struct Pool2dSpec {
unsigned long kH;
unsigned long kW;
unsigned long strideH;
unsigned long strideW;
unsigned long padH;
unsigned long padW;
int ceilMode;
int countIncludePad; // avg pooling only
// Square window; floor mode, padding counted.
static struct Pool2dSpec new(unsigned long kernel, unsigned long stride, unsigned long padding);
};
unsigned long pool_out_side(unsigned long dim, unsigned long kernel, unsigned long stride,
unsigned long padding, int ceilMode);
&FloatTensor FloatTensor::maxpool2d_ex(unsigned long channels, unsigned long height, unsigned long width, const &Pool2dSpec spec) const;
&FloatTensor FloatTensor::avgpool2d_ex(unsigned long channels, unsigned long height, unsigned long width, const &Pool2dSpec spec) const;
&FloatTensor FloatTensor::maxpool2d(unsigned long channels, unsigned long height, unsigned long width, unsigned long kernel,
unsigned long stride, unsigned long padding) const;
&FloatTensor FloatTensor::avgpool2d(unsigned long channels, unsigned long height, unsigned long width, unsigned long kernel,
unsigned long stride, unsigned long padding, int countIncludePad) const;
// Global pooling: each channel's whole plane to one value, [batch, C*H*W]
// -> [batch, C] (a classifier head, or the squeeze of squeeze-and-excite).
// Max sends the gradient to the plane's (first) largest input.
&FloatTensor FloatTensor::global_avgpool2d(unsigned long channels, unsigned long height, unsigned long width) const;
&FloatTensor FloatTensor::global_maxpool2d(unsigned long channels, unsigned long height, unsigned long width) const;
// ── Resizing ─────────────────────────────────────────────────────────────────
// out[c][h*factor + i][w*factor + j] = in[c][h][w]; gradients sum back.
&FloatTensor FloatTensor::upsample_nearest(unsigned long channels, unsigned long height, unsigned long width,
unsigned long factor) const;
// To any outH x outW (up or down), as PyTorch's interpolate(size=...):
// RESIZE_NEAREST source floor(o * in / out) ('nearest'; alignCorners ignored)
// RESIZE_BILINEAR the four nearest pixels
// RESIZE_BICUBIC the sixteen nearest, cubic convolution with a = -0.75
// alignCorners maps the corner pixels' centres onto each other (source o *
// (in - 1) / (out - 1)); otherwise pixel areas line up (source (o + 0.5) *
// in / out - 0.5, clamped at 0 for bilinear). Edges are clamped. Gradients
// scatter back with the same weights.
#define RESIZE_NEAREST 0
#define RESIZE_BILINEAR 1
#define RESIZE_BICUBIC 2
&FloatTensor FloatTensor::resize2d(unsigned long channels, unsigned long height, unsigned long width, unsigned long outH,
unsigned long outW, int mode, int alignCorners) const;
// Tensor::resize2d to factor x the size, bilinear, corners not aligned.
&FloatTensor FloatTensor::upsample_bilinear(unsigned long channels, unsigned long height, unsigned long width,
unsigned long factor) const;
// ── Normalization and dropout ────────────────────────────────────────────────
// Batch normalization (Ioffe & Szegedy 2015) over x [batch, channels*plane]:
// each channel is normalized over the batch and its plane (plane = H*W for
// feature maps, 1 for [batch, features] rows), then scaled by gamma and
// shifted by beta ([channels] each, trained). When 'training', the batch's
// statistics are used and runMean / runVar ([channels], no gradient) move
// toward them by 'momentum' (PyTorch: 0.1, with the unbiased variance);
// otherwise the running statistics are used. eps is added to the variance.
&FloatTensor tensor_batch_norm(const &FloatTensor x, const &FloatTensor gamma, const &FloatTensor beta,
&FloatTensor runMean, &FloatTensor runVar,
unsigned long channels, unsigned long plane,
int training, float momentum, float eps);
// Group normalization (Wu & He 2018): each image's channels split into
// 'groups' groups, each normalized over its channels and planes, then
// gamma / beta per channel. Batch-independent, the same in training and
// inference. groups 1 is layer norm over (C, H, W); groups == channels is
// instance norm.
&FloatTensor tensor_group_norm(const &FloatTensor x, const &FloatTensor gamma, const &FloatTensor beta,
unsigned long channels, unsigned long plane, unsigned long groups, float eps);
// Spectral normalization (Miyato et al. 2018): weight / sigma, sigma its
// largest singular value (the weight seen as a matrix [rows, size/rows]),
// estimated by power iteration from u ([rows], no gradient). In training
// 'iters' iterations update u; in inference u is kept. The gradient treats
// u and v as constants, as PyTorch does.
&FloatTensor tensor_spectral_norm(const &FloatTensor weight, &FloatTensor u, unsigned long iters, int training);
// Inverted dropout: in training each value is zeroed with probability p
// and the rest scaled by 1 / (1 - p); in inference x passes unchanged.
// Tensor::dropout2d drops whole channel planes (spatial dropout).
&FloatTensor FloatTensor::dropout(float p, int training, &Rng rng) const;
&FloatTensor FloatTensor::dropout2d(unsigned long channels, unsigned long plane, float p, int training, &Rng rng) const;
// ── Conditioning ─────────────────────────────────────────────────────────────
// One constant plane per class, 1 in the plane of each row's label (m
// class indices, stored as floats) and 0 elsewhere -> [m,
// classes*height*width], no gradient. tensor_append_label_channels puts
// them after x's channels (x [m, C*height*width] -> [m, (C+classes)*...]):
// how a conditional ConvVae / Dcgan feeds the label to its image-side net.
&FloatTensor FloatTensor::label_planes(unsigned long classes, unsigned long height, unsigned long width) const;
&FloatTensor tensor_append_label_channels(const &FloatTensor x, const &FloatTensor labels, unsigned long classes,
unsigned long height, unsigned long width);
// ── ConvNet ──────────────────────────────────────────────────────────────────
#define CONV_LAYER_CONV 0
#define CONV_LAYER_DECONV 1 // transposed convolution
#define CONV_LAYER_MAXPOOL 2
#define CONV_LAYER_AVGPOOL 3
#define CONV_LAYER_UPSAMPLE 4 // nearest, by 'factor'
#define CONV_LAYER_DENSE 5 // fully connected over the flattened map -> outC x 1 x 1
#define CONV_LAYER_RESHAPE 6 // reinterpret the geometry (same size), no data change
#define CONV_LAYER_BATCHNORM 7
#define CONV_LAYER_RESIZE 8 // Tensor::resize2d to outH x outW ('mode', 'flag' = alignCorners)
#define CONV_LAYER_GLOBAL_AVGPOOL 9 // -> C x 1 x 1
#define CONV_LAYER_GLOBAL_MAXPOOL 10
#define CONV_LAYER_GROUPNORM 11 // 'factor' groups
#define CONV_LAYER_DROPOUT 12 // 'prob', 'flag' = by channel
#define CONV_LAYER_SAVE 13 // keep the activation in slot 'slot'
#define CONV_LAYER_ADD 14 // + slot 'slot' (through a 1x1 projection when 'flag')
#define CONV_LAYER_CONCAT 15 // channel concat with slot 'slot'
#define CONV_LAYER_IDENTITY 16 // no-op (a folded batch norm)
struct ConvLayer {
int kind; // CONV_LAYER_*
int act; // ACT_* after the layer
unsigned long inC;
unsigned long inH;
unsigned long inW;
unsigned long outC;
unsigned long outH;
unsigned long outW;
struct Conv2dSpec conv; // CONV / DECONV, and ADD's projection
struct Pool2dSpec pool; // MAXPOOL / AVGPOOL
unsigned long factor; // UPSAMPLE factor, GROUPNORM groups
int mode; // RESIZE mode
int flag; // RESIZE alignCorners, DROPOUT by channel, ADD projection
float prob; // DROPOUT
unsigned long slot; // SAVE / ADD / CONCAT
unsigned long skipC; // ADD / CONCAT: the saved activation's geometry
unsigned long skipH;
unsigned long skipW;
unsigned long param; // first trained tensor's index in params (weight / gamma; bias / beta next)
unsigned long buffer; // BATCHNORM running mean (variance next), spectral norm's u, DROPOUT's rng
int spectral; // CONV / DECONV / DENSE: weight spectrally normalized (u in buffers[buffer])
};
// Start with the input geometry, then append layers; each one's input is
// the previous one's output. Weights are initialized like PyTorch's:
// U(-1/sqrt(fanIn), 1/sqrt(fanIn)) for weights and biases; the norms start
// at gamma 1, beta 0. params holds every layer's trained tensors, all
// requiring grad, on the heap; buffers holds what is not trained (batch
// norms' running statistics, spectral norms' u). 'training' (1 from
// ConvNet::new) picks batch statistics in batch norm and updates the running
// ones, iterates spectral norms and applies dropout; set it to 0 for
// inference with ConvNet::set_training. The models in autoencoder.h and
// gan.h do that in their inference functions (encode / decode / sample /
// generate / score).
//
// struct ConvNet net = ConvNet::new(1UL, 28UL, 28UL);
// net.conv(32UL, 3UL, 1UL, 1UL, ACT_NONE, &rng); // 32 x 28 x 28
// net.batchnorm(ACT_RELU);
// net.residual_block(64UL, 2UL, 1, ACT_RELU, &rng); // 64 x 14 x 14
// net.global_avgpool(); // 64 x 1 x 1
// net.dense(10UL, ACT_NONE, &rng); // logits
// &Tensor logits = net.forward(batch); // [m, 784] -> [m, 10]
struct ConvNet {
struct Vec<struct ConvLayer> layers;
struct Vec<&FloatTensor> params;
struct Vec<&FloatTensor> buffers; // not trained by the optimizer
struct Vec<struct Rng> rngs; // one per dropout layer
unsigned long slots; // saved activations (CONV_LAYER_SAVE)
int training;
unsigned long inC;
unsigned long inH;
unsigned long inW;
unsigned long outC; // the last layer's output geometry
unsigned long outH;
unsigned long outW;
void conv(unsigned long outC, unsigned long kernel, unsigned long stride, unsigned long padding, int act, &Rng rng);
void conv_ex(unsigned long outC, const &Conv2dSpec spec, int act, &Rng rng);
void deconv(unsigned long outC, unsigned long kernel, unsigned long stride, unsigned long padding, int act,
&Rng rng);
void deconv_ex(unsigned long outC, const &Conv2dSpec spec, int act, &Rng rng);
void maxpool(unsigned long kernel, unsigned long stride, unsigned long padding);
void avgpool(unsigned long kernel, unsigned long stride, unsigned long padding, int countIncludePad);
void maxpool_ex(const &Pool2dSpec spec);
void avgpool_ex(const &Pool2dSpec spec);
void global_avgpool(); // -> C x 1 x 1
void global_maxpool();
void upsample(unsigned long factor); // nearest, integer factor
void upsample_bilinear(unsigned long factor);
void resize(unsigned long outH, unsigned long outW, int mode, int alignCorners);
// Norms over the current output's channels, then 'act': place them between
// a convolution (ACT_NONE) and its activation. Batch norm: momentum 0.1,
// eps 1e-5. Group norm: eps 1e-5; layer norm is 1 group, instance norm one
// group per channel.
void batchnorm(int act);
void groupnorm(unsigned long groups, int act);
void layernorm(int act);
void instancenorm(int act);
// Spectral normalization of the last CONV / DECONV / DENSE layer's weight
// (one power iteration per training forward pass): the usual way to keep a
// GAN critic Lipschitz.
void spectral_norm(&Rng rng);
// Dropout with probability p (byChannel: whole planes); its random stream
// is seeded from rng.
void dropout(float p, int byChannel, &Rng rng);
void dense(unsigned long outDim, int act, &Rng rng);
// channels*height*width must equal the current output size.
void reshape(unsigned long channels, unsigned long height, unsigned long width);
// Skips. ConvNet::save keeps the current activation and returns its slot.
// ConvNet::add_skip adds a slot to the current output, then 'act': through a
// learned 1x1 convolution (stride = the size ratio) when the channels or
// size differ -- ResNet's projection shortcut. ConvNet::concat_skip appends
// a slot's channels (same height and width) to the current output -- a
// U-Net skip. ConvNet::residual_block is ResNet's basic block: conv 3x3
// (stride) -> [batch norm] -> act -> conv 3x3 -> [batch norm] -> + shortcut
// -> act. Output side (dim - 1)/stride + 1.
unsigned long save();
void add_skip(unsigned long slot, int act, &Rng rng);
void concat_skip(unsigned long slot);
void residual_block(unsigned long outC, unsigned long stride, int batchNorm, int act, &Rng rng);
// An activation on the last layer's output (replaces the one it had).
void set_out_act(int act);
// Training (1) or inference (0) mode; returns the previous mode.
int set_training(int training);
// For inference: folds every batch norm that directly follows a CONV /
// DECONV / DENSE layer (no activation or spectral norm between) into that
// layer's weight and bias, using the running statistics, and turns the
// batch norm into a no-op. Inference output is unchanged; do not train the
// net afterwards. Returns how many were folded.
unsigned long fold_batchnorm();
unsigned long in_dim() const; // inC*inH*inW
unsigned long out_dim() const; // outC*outH*outW
&FloatTensor forward(const &FloatTensor x) const; // [batch, in_dim] -> [batch, out_dim]
// Same layers, a copy of the parameters and buffers (e.g. a GAN's averaged
// generator).
struct ConvNet clone() const;
// Parameterized (module.h): the trained parameters; the state adds the
// buffers (running statistics, spectral norms' u), so model_save_file /
// model_copy carry them too.
void collect_params(&Vec<&FloatTensor> out) const;
void collect_state(&Vec<&FloatTensor> out) const;
void drop();
void free() consume; // = drop()
static struct ConvNet new(unsigned long channels, unsigned long height, unsigned long width);
};
} // namespace std