cwenzi/neuroflow-cpp

1#include "test_framework.hpp"
2#include "neuroflow/adamw.hpp"
3#include <cmath>
4
5using namespace neuroflow;
6
7TEST(AdamW, Construction) {
8    AdamW opt(0.001f, 0.9f, 0.999f, 1e-8f, 0.01f);
9    EXPECT_NEAR(opt.lr_, 0.001f, 1e-8f);
10    EXPECT_NEAR(opt.beta1_, 0.9f, 1e-8f);
11    EXPECT_EQ(opt.step_, 0u);
12}
13
14TEST(AdamW, SingleStepUpdate) {
15    AdamW opt(0.01f);
16    Tensor param({4}, QuantType::FP32);
17    float* pp = param.as_fp32();
18    pp[0] = 1.0f; pp[1] = 2.0f; pp[2] = 3.0f; pp[3] = 4.0f;
19
20    Tensor grad({4}, QuantType::FP32);
21    float* gp = grad.as_fp32();
22    gp[0] = 0.1f; gp[1] = 0.2f; gp[2] = 0.3f; gp[3] = 0.4f;
23
24    ParamGroup pg;
25    pg.params = {&param};
26    pg.grads = {&grad};
27    pg.lr = 0.01f;
28    pg.weight_decay = 0.01f;
29    opt.add_param_group(pg);
30
31    float orig_p0 = pp[0];
32    opt.step();
33
34    EXPECT_NE(pp[0], orig_p0);
35    EXPECT_EQ(opt.step_, 1u);
36}
37
38TEST(AdamW, BiasCorrectionStep1) {
39    AdamW opt(0.01f);
40    Tensor param({2}, QuantType::FP32);
41    float* pp = param.as_fp32();
42    pp[0] = 1.0f; pp[1] = 1.0f;
43
44    Tensor grad({2}, QuantType::FP32);
45    float* gp = grad.as_fp32();
46    gp[0] = 1.0f; gp[1] = 1.0f;
47
48    ParamGroup pg;
49    pg.params = {&param};
50    pg.grads = {&grad};
51    pg.lr = 0.01f;
52    pg.weight_decay = 0.0f;
53    opt.add_param_group(pg);
54
55    opt.step();
56
57    float m_hat = 0.1f / (1.0f - 0.9f);
58    float v_hat = 0.01f / (1.0f - 0.999f);
59    float expected = 1.0f - 0.01f * m_hat / (std::sqrt(v_hat) + 1e-8f);
60    EXPECT_NEAR(pp[0], expected, 0.01f);
61}
62
63TEST(AdamW, WeightDecayApplied) {
64    AdamW opt_no_wd(0.01f, 0.9f, 0.999f, 1e-8f, 0.0f);
65    AdamW opt_wd(0.01f, 0.9f, 0.999f, 1e-8f, 0.1f);
66
67    Tensor p1({2}, QuantType::FP32);
68    Tensor p2({2}, QuantType::FP32);
69    float* p1p = p1.as_fp32();
70    float* p2p = p2.as_fp32();
71    p1p[0] = 5.0f; p1p[1] = 5.0f;
72    p2p[0] = 5.0f; p2p[1] = 5.0f;
73
74    Tensor g({2}, QuantType::FP32);
75    float* gp = g.as_fp32();
76    gp[0] = 0.0f; gp[1] = 0.0f;
77
78    ParamGroup pg1;
79    pg1.params = {&p1}; pg1.grads = {&g}; pg1.lr = 0.01f; pg1.weight_decay = 0.0f;
80    opt_no_wd.add_param_group(pg1);
81
82    ParamGroup pg2;
83    pg2.params = {&p2}; pg2.grads = {&g}; pg2.lr = 0.01f; pg2.weight_decay = 0.1f;
84    opt_wd.add_param_group(pg2);
85
86    opt_no_wd.step();
87    opt_wd.step();
88
89    EXPECT_GT(std::abs(p1p[0] - p2p[0]), 1e-6f);
90}
91
92TEST(AdamW, SetLr) {
93    AdamW opt(0.001f);
94    opt.set_lr(0.01f);
95    EXPECT_NEAR(opt.get_lr(), 0.01f, 1e-8f);
96}
97
98TEST(AdamW, NaNGradientSkipped) {
99    AdamW opt(0.01f);
100    Tensor param({2}, QuantType::FP32);
101    float* pp = param.as_fp32();
102    pp[0] = 1.0f; pp[1] = 2.0f;
103
104    Tensor grad({2}, QuantType::FP32);
105    float* gp = grad.as_fp32();
106    gp[0] = std::nanf(""); gp[1] = 0.1f;
107
108    ParamGroup pg;
109    pg.params = {&param}; pg.grads = {&grad}; pg.lr = 0.01f; pg.weight_decay = 0.0f;
110    opt.add_param_group(pg);
111
112    opt.step();
113
114    EXPECT_NEAR(pp[0], 1.0f, 1e-6f);
115}
116
117int main() { RUN_ALL_TESTS(); }
118