385 lines
11 KiB
C
385 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#include "../../lib/test.h"
|
|
#include "../optim.h"
|
|
|
|
#include <string.h>
|
|
#include <stdlib.h>
|
|
#include <math.h>
|
|
|
|
/* ============================================================= */
|
|
/* SGD Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_sgd_default_config(void) {
|
|
fo_sgd_config cfg = fo_sgd_default(0.01f);
|
|
|
|
ASSERT_FLT_NEAR(0.01f, cfg.lr, 0.0001f);
|
|
ASSERT_FLT_NEAR(0.0f, cfg.momentum, 0.0001f);
|
|
ASSERT_FLT_NEAR(0.0f, cfg.weight_decay, 0.0001f);
|
|
ASSERT_FALSE(cfg.nesterov);
|
|
}
|
|
|
|
static void test_sgd_create(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[2];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
params[1] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
tm_set_requires_grad(params[1], true);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 2, &cfg);
|
|
|
|
ASSERT_NOT_NULL(opt);
|
|
ASSERT_FLT_NEAR(0.1f, fo_get_lr(opt), 0.0001f);
|
|
ASSERT_INT_EQ(0, fo_get_step(opt));
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
tm_release(params[1]);
|
|
}
|
|
|
|
static void test_sgd_step(void) {
|
|
int64_t shape[] = {2};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_full(shape, 1, TM_F32, 1.0); /* Initial value = 1 */
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_full(shape, 1, TM_F32, 0.5); /* Gradient = 0.5 */
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
fo_step(opt);
|
|
|
|
/* After step: param = 1.0 - 0.1 * 0.5 = 0.95 */
|
|
float* data = tm_data_f32(params[0]);
|
|
ASSERT_FLT_NEAR(0.95f, data[0], 0.0001f);
|
|
ASSERT_INT_EQ(1, fo_get_step(opt));
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
static void test_sgd_momentum(void) {
|
|
int64_t shape[] = {1};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_full(shape, 1, TM_F32, 1.0);
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_full(shape, 1, TM_F32, 1.0);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
cfg.momentum = 0.9f;
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
/* First step */
|
|
fo_step(opt);
|
|
float* data = tm_data_f32(params[0]);
|
|
ASSERT_FLT_NEAR(0.9f, data[0], 0.0001f); /* 1.0 - 0.1 * 1.0 = 0.9 */
|
|
|
|
/* Second step with same gradient */
|
|
tm_data_f32(params[0]->grad)[0] = 1.0f;
|
|
fo_step(opt);
|
|
/* momentum_buf = 0.9 * 1.0 + 1.0 = 1.9 */
|
|
/* param = 0.9 - 0.1 * 1.9 = 0.71 */
|
|
ASSERT_FLT_NEAR(0.71f, data[0], 0.0001f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Adam Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_adam_default_config(void) {
|
|
fo_adam_config cfg = fo_adam_default(0.001f);
|
|
|
|
ASSERT_FLT_NEAR(0.001f, cfg.lr, 0.0001f);
|
|
ASSERT_FLT_NEAR(0.9f, cfg.beta1, 0.0001f);
|
|
ASSERT_FLT_NEAR(0.999f, cfg.beta2, 0.0001f);
|
|
ASSERT_FLT_NEAR(1e-8f, cfg.eps, 1e-10f);
|
|
ASSERT_FALSE(cfg.amsgrad);
|
|
}
|
|
|
|
static void test_adam_create(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
|
|
fo_adam_config cfg = fo_adam_default(0.001f);
|
|
fo_optimizer* opt = fo_adam_new(params, 1, &cfg);
|
|
|
|
ASSERT_NOT_NULL(opt);
|
|
ASSERT_FLT_NEAR(0.001f, fo_get_lr(opt), 0.0001f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
static void test_adam_step(void) {
|
|
int64_t shape[] = {1};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_full(shape, 1, TM_F32, 0.0);
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_full(shape, 1, TM_F32, 1.0);
|
|
|
|
fo_adam_config cfg = fo_adam_default(0.1f);
|
|
fo_optimizer* opt = fo_adam_new(params, 1, &cfg);
|
|
|
|
fo_step(opt);
|
|
|
|
/* After one step, parameter should have changed */
|
|
float* data = tm_data_f32(params[0]);
|
|
ASSERT_TRUE(data[0] != 0.0f);
|
|
ASSERT_TRUE(data[0] < 0.0f); /* Should be negative since grad was positive */
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
static void test_adamw_create(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
|
|
fo_adam_config cfg = fo_adam_default(0.001f);
|
|
cfg.weight_decay = 0.01f;
|
|
fo_optimizer* opt = fo_adamw_new(params, 1, &cfg);
|
|
|
|
ASSERT_NOT_NULL(opt);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Zero Grad Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_zero_grad(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_full(shape, 1, TM_F32, 5.0);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
/* Verify grad is not zero */
|
|
float* grad = tm_data_f32(params[0]->grad);
|
|
ASSERT_FLT_NEAR(5.0f, grad[0], 0.0001f);
|
|
|
|
fo_zero_grad(opt);
|
|
|
|
/* Verify grad is now zero */
|
|
ASSERT_FLT_NEAR(0.0f, grad[0], 0.0001f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Learning Rate Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_set_lr(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
ASSERT_FLT_NEAR(0.1f, fo_get_lr(opt), 0.0001f);
|
|
|
|
fo_set_lr(opt, 0.01f);
|
|
ASSERT_FLT_NEAR(0.01f, fo_get_lr(opt), 0.0001f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Scheduler Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_scheduler_cosine(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(1.0f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
fo_scheduler_config sched = fo_scheduler_default();
|
|
sched.type = FO_SCHED_COSINE;
|
|
sched.total_steps = 100;
|
|
sched.min_lr = 0.0f;
|
|
fo_set_scheduler(opt, &sched);
|
|
|
|
/* At step 0, should be base_lr */
|
|
float lr0 = fo_get_scheduled_lr(opt, 0);
|
|
ASSERT_FLT_NEAR(1.0f, lr0, 0.01f);
|
|
|
|
/* At step 50 (halfway), should be ~0.5 */
|
|
float lr50 = fo_get_scheduled_lr(opt, 50);
|
|
ASSERT_FLT_NEAR(0.5f, lr50, 0.1f);
|
|
|
|
/* At step 100, should be min_lr */
|
|
float lr100 = fo_get_scheduled_lr(opt, 100);
|
|
ASSERT_FLT_NEAR(0.0f, lr100, 0.01f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
static void test_scheduler_warmup_cosine(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(1.0f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
fo_scheduler_config sched = fo_scheduler_default();
|
|
sched.type = FO_SCHED_WARMUP_COSINE;
|
|
sched.warmup_steps = 10;
|
|
sched.total_steps = 100;
|
|
sched.min_lr = 0.0f;
|
|
fo_set_scheduler(opt, &sched);
|
|
|
|
/* During warmup, should increase linearly */
|
|
float lr5 = fo_get_scheduled_lr(opt, 4); /* Step 5 out of 10 warmup */
|
|
ASSERT_TRUE(lr5 > 0.0f && lr5 < 1.0f);
|
|
|
|
/* At end of warmup, should be at base_lr */
|
|
float lr10 = fo_get_scheduled_lr(opt, 10);
|
|
ASSERT_FLT_NEAR(1.0f, lr10, 0.1f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Gradient Clipping Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_clip_grad_norm(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_full(shape, 1, TM_F32, 2.0); /* Norm = 4 */
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
float norm = fo_clip_grad_norm(opt, 1.0f);
|
|
|
|
/* Original norm should be 4 (sqrt(4 * 4)) = 4 */
|
|
ASSERT_FLT_NEAR(4.0f, norm, 0.01f);
|
|
|
|
/* After clipping to max_norm=1, gradients should be scaled down */
|
|
float* grad = tm_data_f32(params[0]->grad);
|
|
float new_norm = 0.0f;
|
|
for (int i = 0; i < 4; i++) {
|
|
new_norm += grad[i] * grad[i];
|
|
}
|
|
new_norm = sqrtf(new_norm);
|
|
ASSERT_FLT_NEAR(1.0f, new_norm, 0.01f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
static void test_clip_grad_value(void) {
|
|
int64_t shape[] = {4};
|
|
tm_tensor* params[1];
|
|
params[0] = tm_zeros(shape, 1, TM_F32);
|
|
tm_set_requires_grad(params[0], true);
|
|
params[0]->grad = tm_zeros(shape, 1, TM_F32);
|
|
float* grad = tm_data_f32(params[0]->grad);
|
|
grad[0] = 5.0f;
|
|
grad[1] = -5.0f;
|
|
grad[2] = 0.5f;
|
|
grad[3] = -0.5f;
|
|
|
|
fo_sgd_config cfg = fo_sgd_default(0.1f);
|
|
fo_optimizer* opt = fo_sgd_new(params, 1, &cfg);
|
|
|
|
fo_clip_grad_value(opt, 1.0f);
|
|
|
|
ASSERT_FLT_NEAR(1.0f, grad[0], 0.0001f);
|
|
ASSERT_FLT_NEAR(-1.0f, grad[1], 0.0001f);
|
|
ASSERT_FLT_NEAR(0.5f, grad[2], 0.0001f);
|
|
ASSERT_FLT_NEAR(-0.5f, grad[3], 0.0001f);
|
|
|
|
fo_free(opt);
|
|
tm_release(params[0]);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Null Safety Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_null_safety(void) {
|
|
fo_step(NULL);
|
|
fo_zero_grad(NULL);
|
|
fo_free(NULL);
|
|
fo_set_lr(NULL, 0.1f);
|
|
fo_scheduler_step(NULL);
|
|
fo_clip_grad_norm(NULL, 1.0f);
|
|
fo_clip_grad_value(NULL, 1.0f);
|
|
|
|
ASSERT_FLT_NEAR(0.0f, fo_get_lr(NULL), 0.0001f);
|
|
ASSERT_INT_EQ(0, fo_get_step(NULL));
|
|
|
|
ASSERT_TRUE(1); /* Didn't crash */
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Main */
|
|
/* ============================================================= */
|
|
|
|
int main(void) {
|
|
printf("\n Running optim tests...\n\n");
|
|
|
|
/* SGD tests */
|
|
RUN_TEST(sgd_default_config);
|
|
RUN_TEST(sgd_create);
|
|
RUN_TEST(sgd_step);
|
|
RUN_TEST(sgd_momentum);
|
|
|
|
/* Adam tests */
|
|
RUN_TEST(adam_default_config);
|
|
RUN_TEST(adam_create);
|
|
RUN_TEST(adam_step);
|
|
RUN_TEST(adamw_create);
|
|
|
|
/* Common ops tests */
|
|
RUN_TEST(zero_grad);
|
|
RUN_TEST(set_lr);
|
|
|
|
/* Scheduler tests */
|
|
RUN_TEST(scheduler_cosine);
|
|
RUN_TEST(scheduler_warmup_cosine);
|
|
|
|
/* Gradient clipping tests */
|
|
RUN_TEST(clip_grad_norm);
|
|
RUN_TEST(clip_grad_value);
|
|
|
|
/* Null safety */
|
|
RUN_TEST(null_safety);
|
|
|
|
TEST_REPORT();
|
|
|
|
return __ft_failed > 0 ? 1 : 0;
|
|
}
|