torsh-nn

Neural network modules for ToRSh with PyTorch-compatible API, powered by scirs2-neural.

Overview

This crate provides comprehensive neural network building blocks including:

• Common layers (Linear, Conv2d, BatchNorm, etc.)
• Activation functions (ReLU, Sigmoid, GELU, etc.)
• Container modules (Sequential, ModuleList, ModuleDict)
• Parameter initialization utilities
• Functional API for stateless operations

Usage

Basic Module Definition

use torsh_nn::prelude::*;
use torsh_tensor::prelude::*;

// Define a simple neural network
struct SimpleNet {
    fc1: Linear,
    fc2: Linear,
    fc3: Linear,
}

impl SimpleNet {
    fn new() -> Self {
        Self {
            fc1: Linear::new(784, 128, true),
            fc2: Linear::new(128, 64, true),
            fc3: Linear::new(64, 10, true),
        }
    }
}

impl Module for SimpleNet {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let x = self.fc1.forward(input)?;
        let x = F::relu(&x);
        let x = self.fc2.forward(&x)?;
        let x = F::relu(&x);
        self.fc3.forward(&x)
    }
    
    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
        let mut params = self.fc1.parameters();
        params.extend(self.fc2.parameters());
        params.extend(self.fc3.parameters());
        params
    }
    
    // ... other trait methods
}

Using Sequential Container

use torsh_nn::prelude::*;

let model = Sequential::new()
    .add(Linear::new(784, 128, true))
    .add(ReLU::new(false))
    .add(Dropout::new(0.5, false))
    .add(Linear::new(128, 64, true))
    .add(ReLU::new(false))
    .add(Linear::new(64, 10, true));

let output = model.forward(&input)?;

Functional API

use torsh_nn::functional as F;

// Activation functions
let x = F::relu(&input);
let x = F::gelu(&x);
let x = F::softmax(&x, -1)?;

// Pooling
let x = F::max_pool2d(&input, (2, 2), None, None, None)?;
let x = F::global_avg_pool2d(&x)?;

// Loss functions
let loss = F::cross_entropy(&logits, &targets, None, "mean", None)?;
let loss = F::mse_loss(&predictions, &targets, "mean")?;

Parameter Initialization

use torsh_nn::init;

// Xavier/Glorot initialization
let weight = init::xavier_uniform(&[128, 784]);

// Kaiming/He initialization for ReLU
let weight = init::kaiming_normal(&[64, 128], "fan_out");

// Initialize existing tensor
let mut tensor = zeros(&[10, 10]);
init::init_tensor(&mut tensor, "orthogonal", Some(1.0), None);

Common Layers

Linear Layer

let linear = Linear::new(in_features, out_features, bias);

Convolutional Layer

let conv = Conv2d::new(
    in_channels,
    out_channels,
    (3, 3),        // kernel_size
    Some((1, 1)),  // stride
    Some((1, 1)),  // padding
    None,          // dilation
    None,          // groups
    true,          // bias
);

Batch Normalization

let bn = BatchNorm2d::new(
    num_features,
    Some(1e-5),    // eps
    Some(0.1),     // momentum
    true,          // affine
    true,          // track_running_stats
);

LSTM

let lstm = LSTM::new(
    input_size,
    hidden_size,
    Some(2),       // num_layers
    true,          // bias
    false,         // batch_first
    Some(0.2),     // dropout
    false,         // bidirectional
);

Container Modules

ModuleList

let mut layers = ModuleList::new();
layers.append(Linear::new(10, 20, true));
layers.append(Linear::new(20, 30, true));

// Access by index
if let Some(layer) = layers.get(0) {
    let output = layer.forward(&input)?;
}

ModuleDict

let mut blocks = ModuleDict::new();
blocks.insert("encoder".to_string(), Linear::new(784, 128, true));
blocks.insert("decoder".to_string(), Linear::new(128, 784, true));

// Access by key
if let Some(encoder) = blocks.get("encoder") {
    let encoded = encoder.forward(&input)?;
}

Parameter Management

use torsh_nn::parameter::utils;

// Count parameters
let total = utils::count_parameters(&model.parameters());
let trainable = utils::count_trainable_parameters(&model.parameters());

// Freeze/unfreeze parameters
utils::freeze_parameters(&encoder.parameters());
utils::unfreeze_parameters(&decoder.parameters());

// Get parameter statistics
let stats = utils::parameter_stats(&model.parameters());
println!("{}", stats);

// Gradient clipping
utils::clip_grad_norm_(&mut model.parameters(), 1.0, 2.0);
utils::clip_grad_value_(&mut model.parameters(), 0.5);

Integration with SciRS2

This crate leverages scirs2-neural for:

• Optimized layer implementations
• Automatic differentiation support
• Hardware acceleration
• Memory-efficient operations

All modules are designed to work seamlessly with ToRSh's autograd system while benefiting from scirs2's performance optimizations.

License

Licensed under the Apache License, Version 2.0. See LICENSE for details.