cwenzi/neuroflow-cpp

1#include "neuroflow/grad_scaler.hpp"
2
3#include <cmath>
4#include <iostream>
5
6namespace neuroflow {
7
8GradScaler::GradScaler(float init_scale, float growth_factor,
9                       float backoff_factor, size_t growth_interval)
10    : scale_(init_scale), growth_factor_(growth_factor),
11      backoff_factor_(backoff_factor), growth_interval_(growth_interval),
12      growth_tracker_(0) {}
13
14bool GradScaler::has_inf_or_nan(const std::vector<Tensor*>& grads) const {
15    for (const auto* grad : grads) {
16        if (!grad || grad->numel() == 0) continue;
17        const float* data = grad->as_fp32();
18        for (size_t i = 0; i < grad->numel(); ++i) {
19            if (!std::isfinite(data[i])) return true;
20        }
21    }
22    return false;
23}
24
25void GradScaler::unscale(std::vector<Tensor*>& grads) {
26    float inv_scale = 1.0f / scale_;
27    for (auto* grad : grads) {
28        if (!grad || grad->numel() == 0) continue;
29        float* data = grad->as_fp32();
30        for (size_t i = 0; i < grad->numel(); ++i) {
31            data[i] *= inv_scale;
32        }
33    }
34}
35
36void GradScaler::scale_loss(Tensor& loss) {
37    float* d = loss.as_fp32();
38    d[0] *= scale_;
39}
40
41void GradScaler::update(bool found_inf) {
42    if (found_inf) {
43        scale_ *= backoff_factor_;
44        growth_tracker_ = 0;
45    } else {
46        growth_tracker_++;
47        if (growth_tracker_ >= growth_interval_) {
48            scale_ *= growth_factor_;
49            growth_tracker_ = 0;
50        }
51    }
52}
53
54} // namespace neuroflow