348 lines
10 KiB
C
348 lines
10 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* Foundry Training Example: MLP with Adam Optimizer
|
|
*
|
|
* Demonstrates the complete training pipeline:
|
|
* - Tensor operations with autograd
|
|
* - Adam optimizer with learning rate scheduling
|
|
* - MSE loss with automatic differentiation
|
|
* - Progress tracking
|
|
*/
|
|
|
|
#include "../tensor/tensor.h"
|
|
#include "../optim/optim.h"
|
|
#include "../progress/progress.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <math.h>
|
|
#include <time.h>
|
|
|
|
/* ============================================================= */
|
|
/* Simple MLP Model */
|
|
/* ============================================================= */
|
|
|
|
typedef struct {
|
|
/* Layer 1: 2 -> 16 */
|
|
tm_tensor* W1;
|
|
tm_tensor* b1;
|
|
|
|
/* Layer 2: 16 -> 1 */
|
|
tm_tensor* W2;
|
|
tm_tensor* b2;
|
|
|
|
/* Parameter array for optimizer */
|
|
tm_tensor* params[4];
|
|
int num_params;
|
|
} mlp_model;
|
|
|
|
/* Note: We now use tm_add with broadcasting for bias addition */
|
|
|
|
mlp_model* mlp_new(int hidden_size) {
|
|
mlp_model* m = calloc(1, sizeof(mlp_model));
|
|
if (!m) return NULL;
|
|
|
|
int64_t w1_shape[] = {2, hidden_size};
|
|
int64_t b1_shape[] = {hidden_size};
|
|
int64_t w2_shape[] = {hidden_size, 1};
|
|
int64_t b2_shape[] = {1};
|
|
|
|
m->W1 = tm_randn(w1_shape, 2, TM_F32);
|
|
m->b1 = tm_zeros(b1_shape, 1, TM_F32);
|
|
m->W2 = tm_randn(w2_shape, 2, TM_F32);
|
|
m->b2 = tm_zeros(b2_shape, 1, TM_F32);
|
|
|
|
/* Scale weights (Xavier initialization) */
|
|
float scale1 = sqrtf(2.0f / 2.0f);
|
|
float scale2 = sqrtf(2.0f / hidden_size);
|
|
float* w1_data = tm_data_f32(m->W1);
|
|
float* w2_data = tm_data_f32(m->W2);
|
|
for (int64_t i = 0; i < m->W1->numel; i++) w1_data[i] *= scale1;
|
|
for (int64_t i = 0; i < m->W2->numel; i++) w2_data[i] *= scale2;
|
|
|
|
/* Enable gradients */
|
|
tm_set_requires_grad(m->W1, true);
|
|
tm_set_requires_grad(m->b1, true);
|
|
tm_set_requires_grad(m->W2, true);
|
|
tm_set_requires_grad(m->b2, true);
|
|
|
|
/* Setup params array for optimizer */
|
|
m->params[0] = m->W1;
|
|
m->params[1] = m->b1;
|
|
m->params[2] = m->W2;
|
|
m->params[3] = m->b2;
|
|
m->num_params = 4;
|
|
|
|
return m;
|
|
}
|
|
|
|
void mlp_free(mlp_model* m) {
|
|
if (!m) return;
|
|
tm_release(m->W1);
|
|
tm_release(m->b1);
|
|
tm_release(m->W2);
|
|
tm_release(m->b2);
|
|
free(m);
|
|
}
|
|
|
|
/* Forward pass: x -> linear -> relu -> linear -> sigmoid */
|
|
tm_tensor* mlp_forward(mlp_model* m, tm_tensor* x) {
|
|
/* Layer 1: h = relu(x @ W1 + b1) - using broadcasting for bias */
|
|
tm_tensor* z1_pre = tm_matmul(x, m->W1);
|
|
tm_tensor* z1 = tm_add(z1_pre, m->b1); /* Broadcasting: [batch, hidden] + [hidden] */
|
|
tm_tensor* h = tm_relu(z1);
|
|
|
|
/* Layer 2: out = sigmoid(h @ W2 + b2) - using broadcasting for bias */
|
|
tm_tensor* z2_pre = tm_matmul(h, m->W2);
|
|
tm_tensor* z2 = tm_add(z2_pre, m->b2); /* Broadcasting: [batch, 1] + [1] */
|
|
tm_tensor* out = tm_sigmoid(z2);
|
|
|
|
tm_release(z1_pre);
|
|
tm_release(z1);
|
|
tm_release(h);
|
|
tm_release(z2_pre);
|
|
tm_release(z2);
|
|
|
|
return out;
|
|
}
|
|
|
|
/* Manual backward pass (since we don't have full autograd for matmul+bias yet) */
|
|
void mlp_backward(mlp_model* m, tm_tensor* x, tm_tensor* pred, tm_tensor* target) {
|
|
int batch_size = (int)x->shape[0];
|
|
int hidden_size = (int)m->W1->shape[1];
|
|
|
|
float* pred_data = tm_data_f32(pred);
|
|
float* target_data = tm_data_f32(target);
|
|
|
|
/* d_loss/d_pred for MSE: 2 * (pred - target) / n */
|
|
float* d_out = malloc(batch_size * sizeof(float));
|
|
for (int i = 0; i < batch_size; i++) {
|
|
d_out[i] = 2.0f * (pred_data[i] - target_data[i]) / batch_size;
|
|
}
|
|
|
|
/* Sigmoid backward: d_sigmoid = sigmoid * (1 - sigmoid) * d_out */
|
|
float* d_z2 = malloc(batch_size * sizeof(float));
|
|
for (int i = 0; i < batch_size; i++) {
|
|
float s = pred_data[i];
|
|
d_z2[i] = s * (1.0f - s) * d_out[i];
|
|
}
|
|
|
|
/* Recompute hidden layer */
|
|
tm_tensor* z1_pre = tm_matmul(x, m->W1);
|
|
tm_tensor* z1 = tm_add(z1_pre, m->b1); /* Broadcasting */
|
|
tm_tensor* h = tm_relu(z1);
|
|
float* h_data = tm_data_f32(h);
|
|
|
|
/* Ensure gradients exist */
|
|
if (!m->W2->grad) m->W2->grad = tm_zeros(m->W2->shape, m->W2->ndim, TM_F32);
|
|
if (!m->b2->grad) m->b2->grad = tm_zeros(m->b2->shape, m->b2->ndim, TM_F32);
|
|
if (!m->W1->grad) m->W1->grad = tm_zeros(m->W1->shape, m->W1->ndim, TM_F32);
|
|
if (!m->b1->grad) m->b1->grad = tm_zeros(m->b1->shape, m->b1->ndim, TM_F32);
|
|
|
|
/* Gradient for W2: h^T @ d_z2 */
|
|
float* W2_grad = tm_data_f32(m->W2->grad);
|
|
for (int j = 0; j < hidden_size; j++) {
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < batch_size; i++) {
|
|
sum += h_data[i * hidden_size + j] * d_z2[i];
|
|
}
|
|
W2_grad[j] += sum;
|
|
}
|
|
|
|
/* Gradient for b2 */
|
|
float* b2_grad = tm_data_f32(m->b2->grad);
|
|
for (int i = 0; i < batch_size; i++) {
|
|
b2_grad[0] += d_z2[i];
|
|
}
|
|
|
|
/* Gradient for hidden layer: d_z2 @ W2^T */
|
|
float* W2_data = tm_data_f32(m->W2);
|
|
float* d_h = malloc(batch_size * hidden_size * sizeof(float));
|
|
for (int i = 0; i < batch_size; i++) {
|
|
for (int j = 0; j < hidden_size; j++) {
|
|
d_h[i * hidden_size + j] = d_z2[i] * W2_data[j];
|
|
}
|
|
}
|
|
|
|
/* ReLU backward */
|
|
float* z1_data = tm_data_f32(z1);
|
|
float* d_z1 = malloc(batch_size * hidden_size * sizeof(float));
|
|
for (int i = 0; i < batch_size * hidden_size; i++) {
|
|
d_z1[i] = z1_data[i] > 0 ? d_h[i] : 0.0f;
|
|
}
|
|
|
|
/* Gradient for W1: x^T @ d_z1 */
|
|
float* x_data = tm_data_f32(x);
|
|
float* W1_grad = tm_data_f32(m->W1->grad);
|
|
for (int i = 0; i < 2; i++) {
|
|
for (int j = 0; j < hidden_size; j++) {
|
|
float sum = 0.0f;
|
|
for (int b = 0; b < batch_size; b++) {
|
|
sum += x_data[b * 2 + i] * d_z1[b * hidden_size + j];
|
|
}
|
|
W1_grad[i * hidden_size + j] += sum;
|
|
}
|
|
}
|
|
|
|
/* Gradient for b1 */
|
|
float* b1_grad = tm_data_f32(m->b1->grad);
|
|
for (int j = 0; j < hidden_size; j++) {
|
|
for (int i = 0; i < batch_size; i++) {
|
|
b1_grad[j] += d_z1[i * hidden_size + j];
|
|
}
|
|
}
|
|
|
|
free(d_out);
|
|
free(d_z2);
|
|
free(d_h);
|
|
free(d_z1);
|
|
tm_release(z1_pre);
|
|
tm_release(z1);
|
|
tm_release(h);
|
|
}
|
|
|
|
/* Compute MSE loss */
|
|
float compute_mse(tm_tensor* pred, tm_tensor* target) {
|
|
float* p = tm_data_f32(pred);
|
|
float* t = tm_data_f32(target);
|
|
float mse = 0.0f;
|
|
for (int64_t i = 0; i < pred->numel; i++) {
|
|
float diff = p[i] - t[i];
|
|
mse += diff * diff;
|
|
}
|
|
return mse / pred->numel;
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Main Training Loop */
|
|
/* ============================================================= */
|
|
|
|
int main(void) {
|
|
printf("\n");
|
|
printf("===========================================\n");
|
|
printf(" Foundry Training Example: XOR with Adam\n");
|
|
printf("===========================================\n\n");
|
|
|
|
srand((unsigned int)time(NULL));
|
|
|
|
/* XOR dataset */
|
|
static float X_data[] = {
|
|
0.0f, 0.0f,
|
|
0.0f, 1.0f,
|
|
1.0f, 0.0f,
|
|
1.0f, 1.0f
|
|
};
|
|
static float Y_data[] = {0.0f, 1.0f, 1.0f, 0.0f};
|
|
|
|
int64_t x_shape[] = {4, 2};
|
|
int64_t y_shape[] = {4, 1};
|
|
|
|
tm_tensor* X = tm_from_data(X_data, x_shape, 2, TM_F32);
|
|
tm_tensor* Y = tm_from_data(Y_data, y_shape, 2, TM_F32);
|
|
|
|
/* Create model */
|
|
int hidden_size = 16;
|
|
mlp_model* model = mlp_new(hidden_size);
|
|
if (!model) {
|
|
printf("Failed to create model\n");
|
|
return 1;
|
|
}
|
|
|
|
/* Create Adam optimizer */
|
|
fo_adam_config adam_cfg = fo_adam_default(0.01f);
|
|
fo_optimizer* optimizer = fo_adam_new(model->params, model->num_params, &adam_cfg);
|
|
|
|
/* Training hyperparameters */
|
|
int num_epochs = 500;
|
|
int print_every = 50;
|
|
|
|
printf("Model: 2-layer MLP (2 -> %d -> 1)\n", hidden_size);
|
|
printf("Optimizer: Adam (lr=%.3f, beta1=%.1f, beta2=%.3f)\n",
|
|
adam_cfg.lr, adam_cfg.beta1, adam_cfg.beta2);
|
|
printf("Epochs: %d\n\n", num_epochs);
|
|
|
|
/* Training loop */
|
|
fp_progress* progress = fp_new(num_epochs, "Training");
|
|
|
|
float initial_loss = 0.0f;
|
|
float final_loss = 0.0f;
|
|
|
|
for (int epoch = 0; epoch < num_epochs; epoch++) {
|
|
/* Zero gradients */
|
|
fo_zero_grad(optimizer);
|
|
|
|
/* Forward pass */
|
|
tm_tensor* pred = mlp_forward(model, X);
|
|
|
|
/* Compute loss */
|
|
float loss = compute_mse(pred, Y);
|
|
|
|
if (epoch == 0) initial_loss = loss;
|
|
if (epoch == num_epochs - 1) final_loss = loss;
|
|
|
|
/* Backward pass */
|
|
mlp_backward(model, X, pred, Y);
|
|
|
|
/* Optimizer step */
|
|
fo_step(optimizer);
|
|
|
|
/* Progress update */
|
|
fp_update(progress, 1);
|
|
|
|
if ((epoch + 1) % print_every == 0 || epoch == 0) {
|
|
printf("Epoch %4d | Loss: %.6f | LR: %.4f\n",
|
|
epoch + 1, loss, fo_get_lr(optimizer));
|
|
}
|
|
|
|
tm_release(pred);
|
|
}
|
|
|
|
fp_close(progress);
|
|
|
|
/* Final evaluation */
|
|
printf("\n");
|
|
printf("===========================================\n");
|
|
printf(" Training Complete!\n");
|
|
printf("===========================================\n\n");
|
|
|
|
printf("Initial loss: %.6f\n", initial_loss);
|
|
printf("Final loss: %.6f\n", final_loss);
|
|
printf("Improvement: %.2fx\n\n", initial_loss / (final_loss + 1e-8f));
|
|
|
|
printf("Predictions vs Targets:\n");
|
|
printf("-----------------------\n");
|
|
|
|
tm_tensor* final_pred = mlp_forward(model, X);
|
|
float* pred_data = tm_data_f32(final_pred);
|
|
|
|
const char* inputs[] = {"0 XOR 0", "0 XOR 1", "1 XOR 0", "1 XOR 1"};
|
|
int correct = 0;
|
|
for (int i = 0; i < 4; i++) {
|
|
float rounded = pred_data[i] > 0.5f ? 1.0f : 0.0f;
|
|
bool is_correct = (rounded == Y_data[i]);
|
|
if (is_correct) correct++;
|
|
printf(" %s = %.3f (target: %.0f) %s\n",
|
|
inputs[i], pred_data[i], Y_data[i],
|
|
is_correct ? "[OK]" : "[MISS]");
|
|
}
|
|
|
|
printf("\nAccuracy: %d/4 (%.0f%%)\n", correct, 100.0f * correct / 4.0f);
|
|
|
|
if (correct == 4) {
|
|
printf("\n[SUCCESS] Model learned XOR perfectly!\n");
|
|
} else if (final_loss < initial_loss * 0.1f) {
|
|
printf("\n[PROGRESS] Significant improvement. May need more epochs.\n");
|
|
}
|
|
|
|
/* Cleanup */
|
|
tm_release(final_pred);
|
|
tm_release(X);
|
|
tm_release(Y);
|
|
fo_free(optimizer);
|
|
mlp_free(model);
|
|
|
|
printf("\n");
|
|
return 0;
|
|
}
|