Refactor autograd::Variable, option to disable grad calculaios

- autograd::Variable::Shared now a thin layer without methods
- Variable::BackwardFunc_t renamed to Variable::GradFunc_t
- Variable::getData renamed to Variable::array
- Variable::getGrad renamed to Variable::grad
- Variable::backward renamed to Variable::calcGradInputs

Author: pavanky

examples/FFNet.cpp

@@ -9,7 +9,6 @@
 
 #include <af/nn.h>
 
-using namespace af;
 using namespace af;
 using namespace af::nn;
 

examples/autograd.cpp

@@ -13,32 +13,54 @@ using af::autograd::Variable;
 using af::autograd::backward;
 void test1()
 {
-    auto x = Variable(af::randu(5));
-    af_print(x.getData());
+    auto x = Variable(af::randu(5), true);
+    af_print(x.array());
     auto y = x * x;
-    af_print(y.getData());
-    auto dy = Variable(af::constant(1.0, 5));
+    af_print(y.array());
+    auto dy = Variable(af::constant(1.0, 5), false);
     backward(y, dy);
-    af_print(x.getGrad().getData() - 2 * x.getData());
+    auto dx = x.grad();
+    af_print(dx.array() - 2 * x.array());
 }
 
 void test2()
 {
-    auto x = Variable(af::randu(5));
-    af_print(x.getData());
-    auto y = Variable(af::randu(5));
-    af_print(y.getData());
+    auto x = Variable(af::randu(5), true);
+    af_print(x.array());
+    auto y = Variable(af::randu(5), true);
+    af_print(y.array());
     auto z = x * x + x * y + y * y;
-    auto dz = Variable(af::constant(1.0, 5));
+    auto dz = Variable(af::constant(1.0, 5), false);
     backward(z, dz);
-    af_print(x.getGrad().getData() - 2 * x.getData() - y.getData());
-    af_print(y.getGrad().getData() - 2 * y.getData() - x.getData());
+    auto dx = x.grad();
+    auto dy = y.grad();
+    af_print(dx.array() - 2 * x.array() - y.array());
+    af_print(dy.array() - 2 * y.array() - x.array());
+}
+
+void test3()
+{
+    auto x = Variable(af::randu(5), false);
+    af_print(x.array());
+    auto y = Variable(af::randu(5), true);
+    af_print(y.array());
+    auto z = x * x + x * y + y * y;
+    auto dz = Variable(af::constant(1.0, 5), false);
+    backward(z, dz);
+    auto dy = y.grad();
+    af_print(dy.array() - 2 * y.array() - x.array());
+    try {
+        auto dx = x.grad();
+    } catch(af::exception &ex) {
+        std::cout << ex.what() << std::endl;
+    }
 }
 
 int main()
 {
     af::info();
     test1();
     test2();
+    test3();
     return 0;
 }

include/af/autograd/Functions.hpp

@@ -15,22 +15,22 @@ namespace af {
 
         Variable operator +(const Variable lhs, const Variable rhs)
         {
-            auto result = lhs.getData() + rhs.getData();
-            auto backward = [](std::vector<Variable> inputs, Variable grad_output) {
+            auto result = lhs.array() + rhs.array();
+            auto grad_func = [](std::vector<Variable> inputs, Variable grad_output) {
                 inputs[0].addGrad(grad_output);
                 inputs[1].addGrad(grad_output);
             };
-            return Variable(result, {lhs, rhs}, backward);
+            return Variable(result, {lhs, rhs}, grad_func);
         }
 
         Variable operator *(const Variable lhs, const Variable rhs)
         {
-            auto result = lhs.getData() * rhs.getData();
-            auto backward = [](std::vector<Variable> inputs, Variable grad_output) {
+            auto result = lhs.array() * rhs.array();
+            auto grad_func = [](std::vector<Variable> inputs, Variable grad_output) {
                 inputs[0].addGrad(grad_output * inputs[1]);
                 inputs[1].addGrad(grad_output * inputs[0]);
             };
-            return Variable(result, {lhs, rhs}, backward);
+            return Variable(result, {lhs, rhs}, grad_func);
         }
 
     }

include/af/autograd/Grad.hpp

@@ -18,7 +18,7 @@ namespace af {
             var.addGrad(grad);
             Variable::DAG_t dag = var.build();
             for (auto iter = dag.rbegin(); iter != dag.rend(); iter++) {
-                iter->backward();
+                iter->calcGradInputs();
             }
         }
     }

include/af/autograd/Variable.hpp

@@ -28,80 +28,44 @@ namespace af {
         class Variable
         {
         public:
-            typedef std::function<void(std::vector<Variable>, Variable)> BackwardFunc_t;
+            typedef std::function<void(std::vector<Variable>, Variable)> GradFunc_t;
             typedef std::unordered_map<std::ptrdiff_t, bool> Cache_t;
             typedef std::vector<Variable> DAG_t;
 
         private:
-            class Shared {
-            public:
+            struct Shared {
                 Shared() :
+                    m_calc_grad(true),
                     m_data(),
                     m_inputs(),
                     m_grads(),
-                    m_backward(nullptr)
+                    m_grad_func(nullptr)
                 {}
 
-                Shared(af::array data) :
+                Shared(af::array data, bool calc_grad) :
+                    m_calc_grad(calc_grad),
                     m_data(data),
                     m_inputs(),
                     m_grads(),
-                    m_backward(nullptr)
+                    m_grad_func(nullptr)
                 {}
 
-                Shared(af::array data, std::vector<Variable> inputs, BackwardFunc_t backward) :
+                Shared(af::array data,
+                       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_backward(backward)
+                    m_grad_func(grad_func)
                 {}
 
-                af::array getData() const
-                {
-                    return m_data;
-                }
-
-                Variable getGrad() const
-                {
-                    if (m_grads.size() == 0) {
-                        throw std::runtime_error("Gradient hasn't been calculated");
-                    }
-                    return m_grads[0];
-                }
-
-                void addGrad(Variable grad)
-                {
-                    m_grads.push_back(grad);
-                }
-
-                std::vector<Variable> getInputs()
-                {
-                    return m_inputs;
-                }
-
-                void evalGrad()
-                {
-                    if (m_grads.size() == 1) return;
-                    Variable grad = m_grads[0];
-                    for (int i = 1; i < (int)m_grads.size(); i++) {
-                        grad = grad + m_grads[i];
-                    }
-                    grad.getData().eval();
-                    m_grads.clear();
-                    m_grads.push_back(grad);
-                }
-
-                void backward()
-                {
-                    this->evalGrad();
-                    if (m_backward) m_backward(m_inputs, m_grads[0]);
-                }
-
-            private:
+                bool m_calc_grad;
                 af::array m_data;
                 std::vector<Variable> m_inputs;
                 std::vector<Variable> m_grads;
-                BackwardFunc_t m_backward;
+                GradFunc_t m_grad_func;
             };
 
             public:
@@ -111,62 +75,106 @@ namespace af {
             {
             }
 
-            Variable(af::array data) :
-                m_shared(new Shared(data))
+            Variable(af::array data, bool calc_grad) :
+                m_shared(new Shared(data, calc_grad))
             {}
 
             Variable(af::array data,
                      std::vector<Variable> inputs,
-                     BackwardFunc_t backward) :
-                m_shared(new Shared(data, inputs, backward))
-            {}
+                     GradFunc_t grad_func) :
+                m_shared(nullptr)
+            {
+                bool calc_grad = false;
+                for (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 array() const
+            {
+                return m_shared->m_data;
+            }
 
-            af::array getData() const
+            Variable grad() const
             {
-                return m_shared->getData();
+                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];
             }
 
-            Variable getGrad() const
+            bool isCalcGrad()
             {
-                return m_shared->getGrad();
+                return m_shared->m_calc_grad;
+            }
+
+            void 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 addGrad(Variable child_grad)
             {
-                m_shared->addGrad(child_grad);
+                if (m_shared->m_calc_grad) {
+                    m_shared->m_grads.push_back(child_grad);
+                }
             }
 
             std::vector<Variable> getInputs() const
             {
-                return m_shared->getInputs();
+                return m_shared->m_inputs;
             }
 
             void evalGrad()
             {
-                m_shared->evalGrad();
+                // Flag asking not to calculate gradients
+                if (!m_shared->m_calc_grad) return;
+                Variable grad = m_shared->m_grads[0];
+                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.clear();
+                m_shared->m_grads.push_back(grad);
             }
 
-            void backward()
+            void calcGradInputs()
             {
-                m_shared->backward();
+                evalGrad();
+                if (m_shared->m_grad_func) {
+                    m_shared->m_grad_func(m_shared->m_inputs, m_shared->m_grads[0]);
+                }
             }
 
             DAG_t build()
             {
                 Cache_t cache;
                 DAG_t dag;
-                this->buildGraph(cache, dag);
+                this->buildSubGraph(cache, dag);
                 return dag;
             }
 
-            void buildGraph(Cache_t &cache, DAG_t &dag)
+            void buildSubGraph(Cache_t &cache, DAG_t &dag)
             {
                 std::ptrdiff_t id = (std::ptrdiff_t)m_shared.get();
                 if (cache.find(id) != cache.end()) {
                     return;
                 }
-                for (auto input : m_shared->getInputs()) {
-                    input.buildGraph(cache, dag);
+                for (auto input : m_shared->m_inputs) {
+                    input.buildSubGraph(cache, dag);
                 }
                 cache[id] = true;
                 dag.push_back(*this);