neuralforge 0.0.5

An educational framework similar to PyTorch, built to be interpretable and easy to implement.

Deep Learning Framework From Scratch

• Unit-tested and documented educational framework. Similar to PyTorch, but with simpler sintax.
• The autograd engine is in tensor_operations.py. I got a lot of inspiration from Andrej Karpathy's micrograd videos.
• The deep learning model layers are in nn.py.

Check out the implemented basic operations:

• Addition
• Subtraction
• Multiplication
• Division
• Matrix multiplication
• Exponentiation
• Log
• Square Root

The implemented statistics:

• Sum
• Mean
• Max
• Variance

And the implemented tensor operations:

• Reshape
• Transpose
• Concatenate
• Stack
• MaskedFill
• Slice

1. Project Structure

• neuralforge/ : Framework with python files.
  • neuralforge/tensor_operations.py: File with the Tensor class and all of the tensor Operations.
  • neuralforge/nn.py: Most deep learning layers, and nn.Module class.
  • neuralforge/utils.py: File with operations and helper functions.
  • neuralforge/test_framework.py: File with unit tests.
  • neuralforge/optim.py : File with optimizers.
• data/ : Folder to store the text file used to test the Transformer. Currently holds shakespeare.txt.
• setup.py : Setup file for the framework.

2. Running it Yourself

Simple Autograd Example:

import neuralforge as forge

# Instantiate Tensors:
x = forge.randn((8,4,5), requires_grad = True)
w = forge.randn((8,5,4), requires_grad = True)
b = forge.randint((5), requires_grad = True)

# Make calculations:
x = x @ w
x += b

# Compute gradients on whole graph:
x.backward()

# Get gradients from specific Tensors:
print(w.grad)
print(b.grad)

Complex Autograd Example (Transformer):

import neuralforge as forge
import neuralforge.nn as nn

# Implement Transformer class inheriting from forge.nn.Module:
class Transformer(nn.Module):
    def __init__(self, vocab_size: int, hidden_size: int, n_timesteps: int, n_heads: int, p: float):
        super().__init__()
        # Instantiate Transformer's Layers:
        self.embed = nn.Embedding(vocab_size, hidden_size)
        self.pos_embed = nn.PositionalEmbedding(n_timesteps, hidden_size)
        self.b1 = nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_prob=p) 
        self.b2 = nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_prob=p)
        self.ln = nn.LayerNorm(hidden_size)
        self.linear = nn.Linear(hidden_size, vocab_size)

    def forward(self, x):
        z = self.embed(x) + self.pos_embed(x)
        z = self.b1(z)
        z = self.b2(z)
        z = self.ln(z)
        z = self.linear(z)

        return z

# Get tiny Shakespeare test data:
text = load_text_data(f'{PATH}/data/shakespeare.txt')

# Create Transformer instance:
model = Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p)

# Define loss function and optimizer:
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01, reg=0)
        
# Training Loop:
for _ in range(n_iters):
    x, y = get_batch(test_data, n_timesteps, batch_size)

    z = model.forward(x)

    # Get loss:
    loss = loss_func(z, y)

    # Backpropagate the loss using forge.tensor's backward() method:
    loss.backward()

    # Update the weights:
    optimizer.step()

    # Reset the gradients to zero after each training step:
    optimizer.zero_grad()

Note: You can use the framework locally running on the terminal: pip install neuralforge

Requirements

• The required packages are listed in requirements.txt.
• The requirements can be installed on a virtual environment with the command:

pip install -r requirements.txt

Note: The framework is built around numpy, so there is no CUDA availability.

Build a Custom Model

• To create a custom model class, you can use the exact same syntax as you would in PyTorch, inheriting from nn.Module.

You may chose among the following layers:

- nn.Embedding (first layer, turns input indexes into vectors)
- nn.PositionalEmbedding (second layer, adds position information to every timestep of the input)
- nn.Linear (simple fully-connected layer)
- nn.MultiHeadSelfAttention (core of the transformer, calculates weighted sum of inputs)
- nn.RNN (Recurrent Neural Network layer)
- nn.Block (full transformer block - connects MHSA and Dense layers with residuals and LayerNorm)
- nn.CrossEntropyLoss (last layer, returns probabilities for next generated character)

And the following functions:

- nn.Dropout (can be added to apply dropout)
- nn.LayerNorm (normalizes the tensors)
- nn.Softmax (scales the values between 0 and 1)
- nn.Tanh (scales the values between -1 and 1)
- nn.Relu (zeroes all negative values)

3. Results

• The models implemented in test_framework.py all converged to near-zero losses.
• This framework is not as fast or as optimized as PyTorch, but I tried making it more interpretable.
• Hope you enjoy!