Advanced Rust ML: Custom Modules with Tch-rs | ML Engineering

 

Hi All

Another day, another Rust ML backend application. Today we're looking at defining custom modules in Rust in the ML backend. We're going to be creating a custom linear layer in Rust using Tch-rs (Rust PyTorch bindings).  

View full source below:  

Let's break the above code down. Block by block.

1. Imports

use tch::{nn, nn::Module, Tensor};

What it does: 

*     tch::nn - Neural network module containing layer definitions.

*    nn:Module - The trait that all neaural network modules must implement.

*    Tensor - The Tensor type used throughout tch-rs

 2. Struct Definition 

struct CustomLayer {
    weight: Tensor,
    bias: Tensor,
}

What it does:

*    weight - A tensor holding a layer's weights (matrix)

*    bias -  A tensor holding the layer's bias (vector)

This is essentially a linear layer (fully connected layer) that we're defining manually. 

 3.  Constructor Implementation

impl CustomLayer {
    fn new(vs: &nn::Path, in_features: i64, out_features: i64) -> Self {
        Self {
            weight: vs.zeros(&[in_features, out_features]),
            bias: vs.zeros(&[out_features]),
        }
    }
}

What it does: Implements a constructor that:

*    Takes a VarStore path (vs) to register parameters.

*    Takes in_features and out_features as dimensions.

*    Creates and initializes:

    *    weight tensor with shape [in_features, out_features] filled with zeros.

    *    bias tensor with weight [out_feature] filled with zeros. 

The vs.zeros() methods both creates the tensor AND registers it with the VarStore so it can be optimized during training.

 4. Manual Trait Manipulation

impl Module for CustomLayer {
    fn forward(&self, xs: &Tensor) -> Tensor {
        xs.matmul(&self.weight) + &self.bias
    }
}

What it does: Implements the Module trait for the CustomLayer:

*    Defines the forward method that all tch-rs modules must have

*    Takes an input tensor xs and returns an output tensor 

*    Computes: output = input * weights + bias

    *    xs.matmul(&self.weight) - Matrix Multiplication

    *    + &self.bias - add bias vector (broadcasts automatically) 

This is the standard linear transformation y = Wx + b 

Summary

 This code creates a custom linear layer from scratch in tch-rs by:

    1.    Defining a struct to hold parameters

    2.    Implementing a constructor that initializes and registers parameters

    3.    Implement the Module trait with the forward pass computation

This pattern can be extended to create any custom layer type (convolutions, attention, etc.) in Rust while leveraging PyTorch's backend through Tch-rs. 

 And that's all folks. Many more Rust ML examples and experiments coming soon.

 

Later

- Ed