megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

/**
 * \file src/opr/impl/training/optimizer.cpp
 *
 * This file is part of MegBrain, a deep learning framework developed by Megvii.
 *
 * \copyright Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 */

#include "megbrain/opr/training/optimizer.h"
#include "megbrain/exception.h"
#include "megbrain/opr/training/utils.h"

namespace mgb {
namespace optimizer {
SymbolVarArray Optimizer::make_multiple(
        SymbolVarArray symbol_weights, SymbolVarArray symbol_grads,
        std::shared_ptr<mgb::cg::ComputingGraph> graph) {
    if (symbol_weights.size() != symbol_grads.size()) {
        mgb_throw(AssertionError, "The count of weights differs with that of grads.");
    }

    SymbolVarArray r;
    for (size_t i = 0; i < symbol_weights.size(); i++) {
        r.push_back(make(symbol_weights[i], symbol_grads[i], graph));
    }
    return r;
}

SGD::SGD(float lr, float weight_decay, float momentum)
        : m_lr(lr), m_weight_decay(weight_decay), m_momentum(momentum) {
    if (m_lr <= 0) {
        mgb_throw(AssertionError, "Invalid learning rate: negative value.");
    }
    if (m_weight_decay < 0) {
        mgb_throw(AssertionError, "Invalid weight_decay value: negative value.");
    }
    if (m_momentum < 0) {
        mgb_throw(AssertionError, "Invalid momentum value: negative value.");
    }
}

SymbolVar SGD::make(
        SymbolVar symbol_weight, SymbolVar symbol_grad,
        std::shared_ptr<cg::ComputingGraph> graph) {
    SymbolVar symbol_pre_grad;
    auto pre_grad = TensorGen::zeros<dtype::Float32>(
            symbol_grad.shape(), symbol_grad.node()->comp_node());
    m_pre_grads.push_back(pre_grad);
    symbol_pre_grad = opr::SharedDeviceTensor::make(*graph, *pre_grad);

    if (m_weight_decay != .0f) {
        symbol_grad = symbol_grad + m_weight_decay * symbol_weight;
    }

    if (m_momentum != .0f) {
        symbol_pre_grad =
                opr::AddUpdate::make(symbol_pre_grad, symbol_grad, {m_momentum, 1.0f});
        return opr::AddUpdate::make(symbol_weight, -symbol_pre_grad, {1.f, m_lr});
    } else {
        return opr::AddUpdate::make(symbol_weight, -symbol_grad, {1.f, m_lr});
    }
}

Adam::Adam(
        float lr, float weight_decay, std::pair<float, float> betas, float eps,
        bool amsgrad)
        : m_lr(lr),
          m_weight_decay(weight_decay),
          m_betas(betas),
          m_eps(eps),
          m_amsgrad(amsgrad) {
    mgb_assert(m_lr > 0, "Invalid learning rate: negative value.");
    mgb_assert(m_weight_decay >= 0, "Invalid weight_decay value: negative value.");
    mgb_assert(
            m_betas.first >= 0 && m_betas.second >= 0 && m_betas.first < 1 &&
                    m_betas.second < 1,
            "Invalid betas value: negative value or larger than 1.");
}

SymbolVar Adam::make(
        SymbolVar symbol_weight, SymbolVar symbol_grad,
        std::shared_ptr<cg::ComputingGraph> graph) {
    CompNode comp_node = symbol_grad.node()->comp_node();
    DType dt = symbol_grad.dtype();
    m_correction1 = TensorGen::ones<dtype::Float32>({1}, comp_node);
    m_correction2 = TensorGen::ones<dtype::Float32>({1}, comp_node);
    std::shared_ptr<DeviceTensorND> exp_avg =
            std::make_shared<DeviceTensorND>(comp_node, symbol_grad.shape(), dt);
    mgb::fill_zero_dev_tensor(*exp_avg);
    std::shared_ptr<DeviceTensorND> exp_avg_sq =
            std::make_shared<DeviceTensorND>(comp_node, symbol_grad.shape(), dt);
    mgb::fill_zero_dev_tensor(*exp_avg_sq);
    m_exp_avg.push_back(exp_avg);
    m_exp_avg_sq.push_back(exp_avg_sq);

    SymbolVar symbol_correction1 =
            opr::SharedDeviceTensor::make(*graph, *m_correction1);
    SymbolVar symbol_correction2 =
            opr::SharedDeviceTensor::make(*graph, *m_correction2);
    SymbolVar symbol_exp_avg = opr::SharedDeviceTensor::make(*graph, exp_avg);
    SymbolVar symbol_exp_avg_sq = opr::SharedDeviceTensor::make(*graph, exp_avg_sq);

    symbol_correction1 = opr::AddUpdate::make(
            symbol_correction1, symbol_correction1, {m_betas.first, .0f});
    symbol_correction2 = opr::AddUpdate::make(
            symbol_correction2, symbol_correction2, {m_betas.second, .0f});

    if (m_weight_decay != .0f) {
        symbol_grad = symbol_grad + m_weight_decay * symbol_weight;
    }
    symbol_exp_avg = opr::AddUpdate::make(
            symbol_exp_avg, symbol_grad, {m_betas.first, 1.f - m_betas.first});
    symbol_exp_avg_sq = opr::AddUpdate::make(
            symbol_exp_avg_sq, symbol_grad * symbol_grad,
            {m_betas.second, 1.f - m_betas.second});

    SymbolVar delta;
    if (m_amsgrad) {
        std::shared_ptr<DeviceTensorND> max_exp_avg_sq =
                std::make_shared<DeviceTensorND>(comp_node, symbol_grad.shape(), dt);
        mgb::fill_zero_dev_tensor(*max_exp_avg_sq);
        SymbolVar symbol_max_exp_avg_sq =
                opr::SharedDeviceTensor::make(*graph, max_exp_avg_sq);

        symbol_max_exp_avg_sq = opr::AddUpdate::make(
                symbol_exp_avg_sq, opr::max(symbol_max_exp_avg_sq, symbol_exp_avg_sq),
                {1.0f, 1.0f});
        delta = (symbol_exp_avg / (1.f - symbol_correction1)) /
                (opr::powf(symbol_max_exp_avg_sq / (1.f - symbol_correction2), 0.5f) +
                 m_eps);
    } else {
        delta = (symbol_exp_avg / (1.f - symbol_correction1)) /
                (opr::pow(
                         symbol_exp_avg_sq / (1.f - symbol_correction2),
                         symbol_exp_avg.make_scalar(0.5f)) +
                 m_eps);
    }

    return opr::AddUpdate::make(symbol_weight, -delta, {1.0f, m_lr});
}
}  // namespace optimizer
}  // namespace mgb