-
Notifications
You must be signed in to change notification settings - Fork 22
Expand file tree
/
Copy pathVariable.cpp
More file actions
200 lines (173 loc) · 5.87 KB
/
Copy pathVariable.cpp
File metadata and controls
200 lines (173 loc) · 5.87 KB
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
/*******************************************************
* Copyright (c) 2017, ArrayFire
* All rights reserved.
*
* This file is distributed under 3-clause BSD license.
* The complete license agreement can be obtained at:
* http://arrayfire.com/licenses/BSD-3-Clause
********************************************************/
#include <af/autograd/Variable.hpp>
#include <af/autograd/Functions.hpp>
namespace af {
namespace autograd {
Variable::Shared::Shared() :
m_calc_grad(true),
m_data(),
m_inputs(),
m_grads(),
m_grad_func(nullptr)
{}
Variable::Shared::Shared(const af::array &data, bool calc_grad) :
m_calc_grad(calc_grad),
m_data(data),
m_inputs(),
m_grads(),
m_grad_func(nullptr)
{}
Variable::Shared::Shared(const af::array &data,
const std::vector<Variable> &inputs,
GradFunc_t grad_func,
bool calc_grad) :
m_calc_grad(calc_grad),
m_data(data),
m_inputs(inputs.begin(), inputs.end()),
m_grads(),
m_grad_func(grad_func)
{}
Variable::Variable() :
m_shared(new Shared())
{
}
Variable::Variable(const af::array &data, bool calc_grad) :
m_shared(new Shared(data, calc_grad))
{}
Variable::Variable(const af::array &data,
const std::vector<Variable> &inputs,
GradFunc_t grad_func) :
m_shared(nullptr)
{
bool calc_grad = false;
for (const auto &input : inputs) {
calc_grad |= input.isCalcGrad();
}
if (calc_grad) {
m_shared = std::shared_ptr<Shared>(new Shared(data, inputs, grad_func, true));
} else {
m_shared = std::shared_ptr<Shared>(new Shared(data, false));
}
}
af::array& Variable::array() const
{
return m_shared->m_data;
}
Variable& Variable::grad() const
{
if (!m_shared->m_calc_grad) {
throw af::exception("Gradient calclation disabled.");
}
if (m_shared->m_grads.size() == 0) {
throw af::exception("Gradient hasn't been calculated yet.");
}
return m_shared->m_grads[0];
}
std::ptrdiff_t Variable::id() const
{
return (std::ptrdiff_t)m_shared.get();
}
std::vector<Variable>& Variable::getInputs() const
{
return m_shared->m_inputs;
}
bool Variable::isCalcGrad() const
{
return m_shared->m_calc_grad;
}
bool Variable::isGradAvailable() const
{
if (!m_shared->m_calc_grad) return false;
return m_shared->m_grads.size() >= 1;
}
af::dim4 Variable::dims() const
{
return m_shared->m_data.dims();
}
af::dtype Variable::type() const
{
return m_shared->m_data.type();
}
void Variable::zeroGrad()
{
m_shared->m_grads.clear();
}
void Variable::setCalcGrad(bool calc_grad)
{
m_shared->m_calc_grad = calc_grad;
if (!calc_grad) {
m_shared->m_grad_func = nullptr;
m_shared->m_inputs.clear();
m_shared->m_grads.clear();
}
}
void Variable::addGrad(const Variable &child_grad)
{
if (m_shared->m_calc_grad) {
m_shared->m_grads.push_back(child_grad);
}
}
void Variable::evalGrad(bool retain_grad_graph)
{
// Flag asking not to calculate gradients
if (!m_shared->m_calc_grad) return;
// Best not to evaluate the JIT immediately if theres only a single gradient
Variable grad = m_shared->m_grads[0];
if (m_shared->m_grads.size() > 1) {
for (unsigned i = 1; i < m_shared->m_grads.size(); i++) {
grad = grad + m_shared->m_grads[i];
}
grad.array().eval();
m_shared->m_grads.resize(1);
}
grad.setCalcGrad(retain_grad_graph);
m_shared->m_grads[0] = grad;
}
void Variable::calcGradInputs(bool retain_grad_graph)
{
evalGrad();
if (m_shared->m_grad_func) {
m_shared->m_grad_func(m_shared->m_inputs, m_shared->m_grads[0]);
}
}
void Variable::backward(const Variable &grad, bool retain_grad_graph)
{
this->addGrad(grad);
Variable::DAG_t dag = Variable::build(*this);
for (auto iter = dag.rbegin(); iter != dag.rend(); iter++) {
iter->calcGradInputs(retain_grad_graph);
}
}
void Variable::backward(bool retain_grad_graph)
{
auto ones = Variable(af::constant(1, this->dims()), false);
this->backward(ones, retain_grad_graph);
}
Variable::DAG_t Variable::build(const Variable &var)
{
Cache_t cache;
Variable::DAG_t dag;
Variable::buildSubGraph(cache, dag, var);
return dag;
}
void Variable::buildSubGraph(Cache_t &cache, Variable::DAG_t &dag, const Variable &var)
{
std::ptrdiff_t id = var.id();
if (cache.find(id) != cache.end()) {
return;
}
for (const auto &input : var.getInputs()) {
Variable::buildSubGraph(cache, dag, input);
}
cache[id] = true;
dag.push_back(var);
}
}
}