MNIST Benchmark Parity: Sorix vs. TensorFlow vs. PyTorch¶

This notebook compares the performance of Sorix with industry-leading frameworks like PyTorch and TensorFlow using the MNIST (Digit Recognizer) dataset. We will measure:

Training Time (Total across multiple epochs).
Accuracy on the test set.
Inference Speed (Average time per batch).

Experimental Setup¶

Personal Workstation Specs: Intel Core i9 (32 cores), 64GB RAM (This justifies the extremely fast CPU computation).
Architecture: 3-layer MLP with BatchNorm and Dropout.
Optimizer: RMSprop (lr=1e-3, alpha/rho=0.99).
Loss: CrossEntropy.
Hardware: Comparisons explicitly performed on both CPU and GPU.

In [9]:

Copied!

# Uncomment the next line and run this cell to install sorix
#!pip install 'sorix @ git+https://github.com/Mitchell-Mirano/sorix.git@main'
# Uncomment the next line and run this cell to install sorix
#!pip install 'sorix @ git+https://github.com/Mitchell-Mirano/sorix.git@main'

In [10]:

Copied!





import os
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime

# Set seeds for reproducibility
def seed_everything(seed=42):
    np.random.seed(seed)
    try: import torch; torch.manual_seed(seed)
    except: pass
    try: import tensorflow as tf; tf.random.set_seed(seed)
    except: pass

DATA_PATH = "../data/digit-recognizer/train.csv"
print(f"Using data from: {os.path.abspath(DATA_PATH)}")
data = pd.read_csv(DATA_PATH)
SEED = 42
EPOCHS = 10
TRAIN_BATCH_SIZE = 128

# Logging CPU Info
import multiprocessing
print(f"CPU Cores available: {multiprocessing.cpu_count()}")
import os
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime

# Set seeds for reproducibility
def seed_everything(seed=42):
    np.random.seed(seed)
    try: import torch; torch.manual_seed(seed)
    except: pass
    try: import tensorflow as tf; tf.random.set_seed(seed)
    except: pass

DATA_PATH = "../data/digit-recognizer/train.csv"
print(f"Using data from: {os.path.abspath(DATA_PATH)}")
data = pd.read_csv(DATA_PATH)
SEED = 42
EPOCHS = 10
TRAIN_BATCH_SIZE = 128

# Logging CPU Info
import multiprocessing
print(f"CPU Cores available: {multiprocessing.cpu_count()}")

Using data from: /home/mitchellmirano/Desktop/MitchellProjects/sorix/docs/examples/data/digit-recognizer/train.csv
CPU Cores available: 32

0. Data Preparation¶

We use the same data split for all frameworks to ensure a fair comparison.

In [11]:

Copied!

from sklearn.model_selection import train_test_split

X = data.drop("label", axis=1).values.astype('float32') / 255.0
y = data["label"].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=SEED)
print(f"Train size: {len(X_train)}, Test size: {len(X_test)}")
from sklearn.model_selection import train_test_split

X = data.drop("label", axis=1).values.astype('float32') / 255.0
y = data["label"].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=SEED)
print(f"Train size: {len(X_train)}, Test size: {len(X_test)}")

Train size: 33600, Test size: 8400

1. Unified Device-Aware Benchmark¶

We evaluate each framework on both CPU and GPU (if available) to differentiate performance accurately.

In [12]:

Copied!





import torch
import tensorflow as tf
import cupy as cp
import time
import sorix
import numpy as np
from sorix import tensor
from sorix.nn import Module, Linear, CrossEntropyLoss, ReLU, BatchNorm1d, Dropout
from sorix.optim import RMSprop
from sorix.utils.data import Dataset, DataLoader

# Results storage
all_results = []
INFERENCE_BATCH_CPU = 4096  # Using large batch for CPU as requested
INFERENCE_BATCH_GPU = 1024

def run_sorix(device_name='cpu'):
    print(f"--- Running Sorix on {device_name} ---")
    seed_everything(SEED)
    
    class SorixModel(Module):
        def __init__(self):
            super().__init__()
            self.linear1 = Linear(784, 128, bias=False)
            self.bn1 = BatchNorm1d(128)
            self.linear2 = Linear(128, 64)
            self.linear3 = Linear(64, 10)
            self.relu = ReLU()
            self.dropout = Dropout(p=0.2)
        def forward(self, x):
            x = self.linear1(x); x = self.bn1(x); x = self.relu(x)
            x = self.linear2(x); x = self.relu(x)
            x = self.dropout(x); x = self.linear3(x)
            return x

    model = SorixModel().to(device_name)
    loss_fn = CrossEntropyLoss()
    optimizer = RMSprop(model.parameters(), lr=1e-3, alpha=0.99)
    train_ds = Dataset(X_train, y_train.reshape(-1, 1))
    train_loader = DataLoader(train_ds, batch_size=TRAIN_BATCH_SIZE, shuffle=True)

    # 1. Training Time
    start_train = time.time()
    for epoch in range(EPOCHS):
        model.train()
        for xb, yb in train_loader:
            xb, yb = xb.to(device_name), yb.to(device_name)
            optimizer.zero_grad()
            logits = model(xb)
            loss = loss_fn(logits, yb)
            loss.backward()
            optimizer.step()
    train_time = time.time() - start_train

    # 2. Inference Time
    model.eval()
    batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
    inf_loader = DataLoader(Dataset(X_test, y_test.reshape(-1, 1)), batch_size=batch_size)
    
    # Warmup
    dummy = tensor(X_test[:100]).to(device_name)
    with sorix.no_grad(): _ = model(dummy)
    
    start_inf = time.time()
    with sorix.no_grad():
        for xb, _ in inf_loader:
            xb = xb.to(device_name)
            _ = model(xb)
    inf_time = time.time() - start_inf

    # Accuracy check
    tx = tensor(X_test[:1000]).to(device_name)
    with sorix.no_grad():
        out = model(tx)
        preds = sorix.argmax(out, axis=1, keepdims=True)
    acc = (preds.cpu().data.flatten() == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'Sorix', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

def run_pytorch(device_name='cpu'):
    print(f"--- Running PyTorch on {device_name} ---")
    pt_device = torch.device(device_name)
    seed_everything(SEED)
    
    class PyTorchModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.net = torch.nn.Sequential(
                torch.nn.Linear(784, 128, bias=False),
                torch.nn.BatchNorm1d(128),
                torch.nn.ReLU(),
                torch.nn.Linear(128, 64),
                torch.nn.ReLU(),
                torch.nn.Dropout(0.2),
                torch.nn.Linear(64, 10)
            )
        def forward(self, x): return self.net(x)

    model = PyTorchModel().to(pt_device)
    loss_fn = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-3, alpha=0.99)
    ds = torch.utils.data.TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train).long())
    loader = torch.utils.data.DataLoader(ds, batch_size=TRAIN_BATCH_SIZE, shuffle=True)

    # 1. Training Time
    start_train = time.time()
    for epoch in range(EPOCHS):
        model.train()
        for xb, yb in loader:
            xb, yb = xb.to(pt_device), yb.to(pt_device)
            optimizer.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            optimizer.step()
    train_time = time.time() - start_train

    # 2. Inference Time
    model.eval()
    batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
    inf_loader = torch.utils.data.DataLoader(
        torch.utils.data.TensorDataset(torch.from_numpy(X_test)), 
        batch_size=batch_size
    )
    
    # Warmup
    dummy = torch.from_numpy(X_test[:100]).to(pt_device)
    with torch.no_grad(): _ = model(dummy)
    
    start_inf = time.time()
    with torch.no_grad():
        for xb, in inf_loader:
            xb = xb.to(pt_device)
            _ = model(xb)
    inf_time = time.time() - start_inf

    # Accuracy
    model.eval()
    tx = torch.from_numpy(X_test[:1000]).to(pt_device)
    with torch.no_grad():
        acc = (model(tx).argmax(1).cpu().numpy() == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'PyTorch', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

def run_tensorflow(device_name='cpu'):
    tf_dev = f"/{device_name.upper()}:0"
    print(f"--- Running TensorFlow on {tf_dev} ---")
    seed_everything(SEED)
    
    with tf.device(tf_dev):
        model = tf.keras.models.Sequential([
            tf.keras.layers.Dense(128, input_shape=(784,), use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(),
            tf.keras.layers.Dense(64),
            tf.keras.layers.ReLU(),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(10)
        ])
        model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-3, rho=0.99),
                      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
        
        # 1. Training Time
        start_train = time.time()
        model.fit(X_train, y_train, batch_size=TRAIN_BATCH_SIZE, epochs=EPOCHS, verbose=0)
        train_time = time.time() - start_train
        
        # 2. Inference Time
        batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
        # Warmup
        _ = model.predict(X_test[:100], verbose=0)
        
        start_inf = time.time()
        _ = model.predict(X_test, batch_size=batch_size, verbose=0)
        inf_time = time.time() - start_inf
        
        preds = model.predict(X_test[:1000], verbose=0).argmax(axis=1)
        acc = (preds == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'TensorFlow', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

# Run all
sorix_model = run_sorix('cpu')
pytorch_model = run_pytorch('cpu')
tensorflow_model = run_tensorflow('cpu')

if sorix.cuda.is_available():
    sorix_model_gpu = run_sorix('cuda')
    pytorch_model_gpu = run_pytorch('cuda')
    tensorflow_model_gpu = run_tensorflow('gpu')

df_results = pd.DataFrame(all_results)
display(df_results)

import torch
import tensorflow as tf
import cupy as cp
import time
import sorix
import numpy as np
from sorix import tensor
from sorix.nn import Module, Linear, CrossEntropyLoss, ReLU, BatchNorm1d, Dropout
from sorix.optim import RMSprop
from sorix.utils.data import Dataset, DataLoader

# Results storage
all_results = []
INFERENCE_BATCH_CPU = 4096  # Using large batch for CPU as requested
INFERENCE_BATCH_GPU = 1024

def run_sorix(device_name='cpu'):
    print(f"--- Running Sorix on {device_name} ---")
    seed_everything(SEED)
    
    class SorixModel(Module):
        def __init__(self):
            super().__init__()
            self.linear1 = Linear(784, 128, bias=False)
            self.bn1 = BatchNorm1d(128)
            self.linear2 = Linear(128, 64)
            self.linear3 = Linear(64, 10)
            self.relu = ReLU()
            self.dropout = Dropout(p=0.2)
        def forward(self, x):
            x = self.linear1(x); x = self.bn1(x); x = self.relu(x)
            x = self.linear2(x); x = self.relu(x)
            x = self.dropout(x); x = self.linear3(x)
            return x

    model = SorixModel().to(device_name)
    loss_fn = CrossEntropyLoss()
    optimizer = RMSprop(model.parameters(), lr=1e-3, alpha=0.99)
    train_ds = Dataset(X_train, y_train.reshape(-1, 1))
    train_loader = DataLoader(train_ds, batch_size=TRAIN_BATCH_SIZE, shuffle=True)

    # 1. Training Time
    start_train = time.time()
    for epoch in range(EPOCHS):
        model.train()
        for xb, yb in train_loader:
            xb, yb = xb.to(device_name), yb.to(device_name)
            optimizer.zero_grad()
            logits = model(xb)
            loss = loss_fn(logits, yb)
            loss.backward()
            optimizer.step()
    train_time = time.time() - start_train

    # 2. Inference Time
    model.eval()
    batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
    inf_loader = DataLoader(Dataset(X_test, y_test.reshape(-1, 1)), batch_size=batch_size)
    
    # Warmup
    dummy = tensor(X_test[:100]).to(device_name)
    with sorix.no_grad(): _ = model(dummy)
    
    start_inf = time.time()
    with sorix.no_grad():
        for xb, _ in inf_loader:
            xb = xb.to(device_name)
            _ = model(xb)
    inf_time = time.time() - start_inf

    # Accuracy check
    tx = tensor(X_test[:1000]).to(device_name)
    with sorix.no_grad():
        out = model(tx)
        preds = sorix.argmax(out, axis=1, keepdims=True)
    acc = (preds.cpu().data.flatten() == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'Sorix', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

def run_pytorch(device_name='cpu'):
    print(f"--- Running PyTorch on {device_name} ---")
    pt_device = torch.device(device_name)
    seed_everything(SEED)
    
    class PyTorchModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.net = torch.nn.Sequential(
                torch.nn.Linear(784, 128, bias=False),
                torch.nn.BatchNorm1d(128),
                torch.nn.ReLU(),
                torch.nn.Linear(128, 64),
                torch.nn.ReLU(),
                torch.nn.Dropout(0.2),
                torch.nn.Linear(64, 10)
            )
        def forward(self, x): return self.net(x)

    model = PyTorchModel().to(pt_device)
    loss_fn = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-3, alpha=0.99)
    ds = torch.utils.data.TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train).long())
    loader = torch.utils.data.DataLoader(ds, batch_size=TRAIN_BATCH_SIZE, shuffle=True)

    # 1. Training Time
    start_train = time.time()
    for epoch in range(EPOCHS):
        model.train()
        for xb, yb in loader:
            xb, yb = xb.to(pt_device), yb.to(pt_device)
            optimizer.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            optimizer.step()
    train_time = time.time() - start_train

    # 2. Inference Time
    model.eval()
    batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
    inf_loader = torch.utils.data.DataLoader(
        torch.utils.data.TensorDataset(torch.from_numpy(X_test)), 
        batch_size=batch_size
    )
    
    # Warmup
    dummy = torch.from_numpy(X_test[:100]).to(pt_device)
    with torch.no_grad(): _ = model(dummy)
    
    start_inf = time.time()
    with torch.no_grad():
        for xb, in inf_loader:
            xb = xb.to(pt_device)
            _ = model(xb)
    inf_time = time.time() - start_inf

    # Accuracy
    model.eval()
    tx = torch.from_numpy(X_test[:1000]).to(pt_device)
    with torch.no_grad():
        acc = (model(tx).argmax(1).cpu().numpy() == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'PyTorch', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

def run_tensorflow(device_name='cpu'):
    tf_dev = f"/{device_name.upper()}:0"
    print(f"--- Running TensorFlow on {tf_dev} ---")
    seed_everything(SEED)
    
    with tf.device(tf_dev):
        model = tf.keras.models.Sequential([
            tf.keras.layers.Dense(128, input_shape=(784,), use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(),
            tf.keras.layers.Dense(64),
            tf.keras.layers.ReLU(),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(10)
        ])
        model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-3, rho=0.99),
                      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))
        
        # 1. Training Time
        start_train = time.time()
        model.fit(X_train, y_train, batch_size=TRAIN_BATCH_SIZE, epochs=EPOCHS, verbose=0)
        train_time = time.time() - start_train
        
        # 2. Inference Time
        batch_size = INFERENCE_BATCH_CPU if device_name == 'cpu' else INFERENCE_BATCH_GPU
        # Warmup
        _ = model.predict(X_test[:100], verbose=0)
        
        start_inf = time.time()
        _ = model.predict(X_test, batch_size=batch_size, verbose=0)
        inf_time = time.time() - start_inf
        
        preds = model.predict(X_test[:1000], verbose=0).argmax(axis=1)
        acc = (preds == y_test[:1000]).mean()
    
    all_results.append({
        'Framework': 'TensorFlow', 
        'Device': 'GPU' if device_name.lower() in ['cuda', 'gpu'] else 'CPU', 
        'train_batch_size':TRAIN_BATCH_SIZE,
        'Train Time': train_time, 
        'test_batch_size':batch_size,
        'Inference Time': inf_time,
        'Accuracy': acc
    })
    return model

# Run all
sorix_model = run_sorix('cpu')
pytorch_model = run_pytorch('cpu')
tensorflow_model = run_tensorflow('cpu')

if sorix.cuda.is_available():
    sorix_model_gpu = run_sorix('cuda')
    pytorch_model_gpu = run_pytorch('cuda')
    tensorflow_model_gpu = run_tensorflow('gpu')

df_results = pd.DataFrame(all_results)
display(df_results)

--- Running Sorix on cpu ---
--- Running PyTorch on cpu ---
--- Running TensorFlow on /CPU:0 ---

/home/mitchellmirano/Desktop/MitchellProjects/sorix/.venv/lib/python3.13/site-packages/keras/src/layers/core/dense.py:106: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(activity_regularizer=activity_regularizer, **kwargs)

✅ GPU basic operation passed
✅ GPU available: NVIDIA GeForce RTX 4070 Laptop GPU
CUDA runtime version: 13000
CuPy version: 13.6.0
--- Running Sorix on cuda ---
--- Running PyTorch on cuda ---
--- Running TensorFlow on /GPU:0 ---

I0000 00:00:1781022202.438069   24508 dot_merger.cc:481] Merging Dots in computation: a_inference_one_step_on_data_30026__.14
I0000 00:00:1781022204.279871   24512 dot_merger.cc:481] Merging Dots in computation: a_inference_one_step_on_data_30026__.14

	Framework	Device	train_batch_size	Train Time	test_batch_size	Inference Time	Accuracy
0	Sorix	CPU	128	8.570118	4096	0.028992	0.972
1	PyTorch	CPU	128	5.102359	4096	0.024793	0.974
2	TensorFlow	CPU	128	17.448998	4096	0.201416	0.967
3	Sorix	GPU	128	6.441508	1024	0.013910	0.974
4	PyTorch	GPU	128	3.875234	1024	0.037778	0.976
5	TensorFlow	GPU	128	9.963803	1024	0.492464	0.970

2. Model Export and Size Comparison¶

Size is measured using CPU-exported state dicts.

In [13]:

Copied!





import pickle
# We use the CPU models for size comparison to be fair
sorix_path = "model_sorix.sor"
with open(sorix_path, 'wb') as f: pickle.dump(sorix_model.state_dict(), f)
sorix_size = os.path.getsize(sorix_path) / 1024

pytorch_path = "model_pytorch.pt"
torch.save(pytorch_model.state_dict(), pytorch_path)
pytorch_size = os.path.getsize(pytorch_path) / 1024

tensorflow_path = "model_tf.keras"
tensorflow_model.save(tensorflow_path) # TF saves full model usually
tensorflow_size = os.path.getsize(tensorflow_path) / 1024

model_sizes = pd.DataFrame({
    'Framework': ['Sorix', 'PyTorch', 'TensorFlow'],
    'Size (KB)': [sorix_size, pytorch_size, tensorflow_size]
})
display(model_sizes)

import pickle
# We use the CPU models for size comparison to be fair
sorix_path = "model_sorix.sor"
with open(sorix_path, 'wb') as f: pickle.dump(sorix_model.state_dict(), f)
sorix_size = os.path.getsize(sorix_path) / 1024

pytorch_path = "model_pytorch.pt"
torch.save(pytorch_model.state_dict(), pytorch_path)
pytorch_size = os.path.getsize(pytorch_path) / 1024

tensorflow_path = "model_tf.keras"
tensorflow_model.save(tensorflow_path) # TF saves full model usually
tensorflow_size = os.path.getsize(tensorflow_path) / 1024

model_sizes = pd.DataFrame({
    'Framework': ['Sorix', 'PyTorch', 'TensorFlow'],
    'Size (KB)': [sorix_size, pytorch_size, tensorflow_size]
})
display(model_sizes)

	Framework	Size (KB)
0	Sorix	429.559570
1	PyTorch	432.911133
2	TensorFlow	890.761719