moonxi-net

A deep learning training framework built with MoonBit, featuring tape-based autograd and PyTorch-like API (CPU backend)

deep-learning
resnet
neural-network
autograd
tensor
moon add chnlkw/moonxi-net@0.1.1
Download zip
Author
Version
0.1.1
License
MIT
Last updated
3 months ago
Downloads
33

Dependencies

README

#moonxi-net

A deep learning training framework built with MoonBit, featuring tape-based autograd and a PyTorch-like API on the CPU backend (NpArray). Write your model once using the Tensor trait, and it works with both CPU and GPU backends via tagless-final style.

#Installation

moon add chnlkw/moonxi-net

Note: This package works with the standard MoonBit toolchain — no CUDA or GPU required.

#Quick Example

A minimal linear regression: learn y = 3x + 1 from data.

///|
fn[T : @tensor.Tensor + @tensor.BlasTensor] train_linear() -> Array[T] {
let w = @grad.no_grad(T::zeros([1, 1]))
let b = @grad.no_grad(T::zeros([1, 1]))
let x = @grad.no_grad(
T::from_host(FixedArray::makei(5, i => Float::from_int(i + 1)), [5, 1]),
)
let y = @grad.no_grad(
T::from_host(FixedArray::makei(5, i => Float::from_int((i + 1) * 3 + 1)), [
5, 1,
]),
)
let params : Array[@grad.Grad[T]] = [w, b]
for _epoch in 0500 {
@grad.clear_tape()
for p in params { p.grad = Some(None) }
let loss = x.matmul(w).add(b).sub(y).square().mean()
loss.backward()
for p in params {
match p.grad {
Some(Some(g)) => p.value = p.value.sub(g.scale(0.001))
_ => ()
}
}
}
params.map(g => g.value)
}

moon run moonxi-net/examples/linear --target native # Learned: w=3.00, b=1.00

#Usage

#Import in moon.pkg

{ "import": [ "chnlkw/moonxi-net" @tensor, "chnlkw/moonxi-net/nparray" @nparray, "chnlkw/moonxi-net/grad" @grad, "chnlkw/moonxi-net/model" @model, "chnlkw/moonxi-net/optimizer" @optimizer, "chnlkw/moonxi-net/loss" @loss, "chnlkw/moonxi-net/train" @train, "chnlkw/moonxi-net/dataloader" @dl ] }

#Core Packages

PackageAliasDescription
moonxi-net@tensorCore traits: Tensor, BlasTensor, ImageTensor, ImageBackwardOps
moonxi-net/nparray@nparrayCPU tensor backend (NpArray implements all tensor traits)
moonxi-net/grad@gradTape-based autograd engine (Grad[T], backward(), clear_tape())
moonxi-net/model@modelNeural network layers: Linear[T], Conv2d[T], ResNet18[T], MLP[T]
moonxi-net/optimizer@optimizerOptimizers: Momentum SGD, Adam, RMSprop (with gradient clipping)
moonxi-net/loss@lossLoss functions: cross-entropy, MSE
moonxi-net/train@trainTraining loop utilities, Experiment config, run_cpu
moonxi-net/dataloader@dlDataLoader with Fisher-Yates shuffle and mini-batch iteration
moonxi-net/datasets/mnistMNIST dataset loader
moonxi-net/datasets/cifar10CIFAR-10 dataset loader

#Basic Training Workflow

///|
fn main {
// Create tensors using NpArray (CPU backend)
let w = @grad.no_grad(@nparray.NpArray::zeros([10, 10]))
let b = @grad.no_grad(@nparray.NpArray::zeros([10]))

// Build computation graph, compute loss, backprop
@grad.clear_tape()
let loss = /* ... your forward pass ... */
loss.backward()

// Access gradients
match w.grad {
Some(Some(g)) => w.value = w.value.sub(g.scale(0.01))
_ => ()
}
}

#Build & Test

# From this directory moon test --target native # Build moon build --target native

#License

#
BlasTensor

pub(open) trait BlasTensor {
fn matmul(Self, Self) -> Self
fn transpose(Self) -> Self
}

#
ImageBackwardOps

pub(open) trait ImageBackwardOps {
fn conv2d_backward_data(grad_output : Self, weight : Self, input : Self, stride : Int, padding : Int) -> Self
fn conv2d_backward_weight(grad_output : Self, input : Self, weight : Self, stride : Int, padding : Int) -> Self
fn conv2d_backward_bias(grad_output : Self) -> Self
fn relu_backward(grad_output : Self, input : Self) -> Self
fn batchnorm_backward(grad_output : Self, input : Self, gamma : Self, save_mean : Self, save_inv_var : Self, eps : Float) -> (Self, Self, Self)
fn maxpool2d_backward(grad_output : Self, input : Self, kernel_size : Int, stride : Int) -> Self
fn adaptive_avg_pool2d_backward(grad_output : Self, input : Self) -> Self
fn softmax_ce_backward(logits : Self, targets : Self) -> Self
fn softmax_ce_backward_labels(logits : Self, labels : Self, num_classes : Int) -> Self
}

#
ImageTensor

pub(open) trait ImageTensor {
fn conv2d(Self, weight : Self, bias : Self, stride : Int, padding : Int) -> Self
fn relu(Self) -> Self
fn maxpool2d(Self, kernel_size : Int, stride : Int) -> Self
fn adaptive_avg_pool2d(Self, output_size : Int) -> Self
fn batchnorm_training(Self, gamma : Self, beta : Self, running_mean : Self, running_var : Self, momentum : Float, eps : Float) -> (Self, Self, Self)
fn batchnorm_inference(Self, gamma : Self, beta : Self, running_mean : Self, running_var : Self, eps : Float) -> Self
fn softmax_cross_entropy(Self, targets : Self) -> Self
fn cross_entropy_with_labels(Self, labels : Self, num_classes : Int) -> Self
}

#
Tensor

pub(open) trait Tensor {
fn dims(Self) -> FixedArray[Int]
fn zeros(dims : FixedArray[Int]) -> Self
fn zeros_like(Self) -> Self
fn from_host(data : FixedArray[Float], shape : Array[Int]) -> Self
fn square(Self) -> Self
fn sqrt(Self) -> Self
fn mean(Self) -> Self
fn scale(Self, Float) -> Self
fn mul_elem(Self, Self) -> Self
fn div_elem(Self, Self) -> Self
fn scalar(Float) -> Self
fn size(Self) -> Int
fn reduce_sum_to(Self, target_dims : FixedArray[Int]) -> Self
fn value(Self) -> TensorData
fn view(Self, new_shape : FixedArray[Int]) -> Self
fn add(Self, Self) -> Self
fn add_into(Self, Self) -> Unit
fn sub(Self, Self) -> Self
fn broadcast_to(Self, target_shape : FixedArray[Int]) -> Self
}

#
CudaError

pub(all) suberror CudaError {
CudaError(String)
} derive(
Debug
)

#
IoError

pub(all) suberror IoError {
IoError(String)
} derive(
Debug
)

#
ShapeError

pub(all) suberror ShapeError {
ShapeError(String)
} derive(
Debug
)

#
UnsupportedError

pub(all) suberror UnsupportedError {
UnsupportedError(String)
} derive(
Debug
)

#
ShapeTensor

pub(all) struct ShapeTensor {
shape : FixedArray[Int]
} derive(
Debug
)

A shape-only tensor that computes output shapes without actual data. Aborts when shape constraints are violated.

#
ShapeTensor::dim

fn ShapeTensor::dim(self : ShapeTensor, i : Int) -> Int

#
ShapeTensor::from_array

fn ShapeTensor::from_array(shape : Array[Int]) -> ShapeTensor

#
ShapeTensor::linear

fn ShapeTensor::linear(input : ShapeTensor, weight : ShapeTensor) -> ShapeTensor

Linear layer output shape: [batch, out_features] from input [batch, in_f] and weight [out_f, in_f].

#
ShapeTensor::ndim

fn ShapeTensor::ndim(self : ShapeTensor) -> Int

#
ShapeTensor::pool2d

fn ShapeTensor::pool2d(self : ShapeTensor, kernel_size : Int, stride : Int, padding : Int) -> ShapeTensor

Pool2d output shape with explicit padding (not in trait).

#
ShapeTensor::to_array

fn ShapeTensor::to_array(self : ShapeTensor) -> Array[Int]

#
TensorData

pub(all) struct TensorData {
dims : FixedArray[Int]
data : FixedArray[Float]
} derive(
Debug
)