Conv2D

CMakeLists.txt

@@ -18,6 +18,7 @@ target_sources(afml
   src/nn/Modules/Loss.cpp
   src/nn/Modules/Module.cpp
   src/nn/Init.cpp
+  src/nn/Modules/Conv2D.cpp
   )
 
 target_include_directories(afml

examples/CMakeLists.txt

@@ -13,6 +13,7 @@ endfunction(build_example)
 # build_example(Activations.cpp)
 # build_example(FFNet.cpp)
 # build_example(Node.cpp)
-build_example(perceptron.cpp)
+# build_example(perceptron.cpp)
+build_example(alexnet.cpp)
 # build_example(Weights.cpp)
 build_example(autograd.cpp)

include/af/autograd/Functions.hpp

@@ -75,6 +75,13 @@ namespace af {
         Variable matmulTN(const Variable &lhs, const Variable &rhs);
         Variable matmulNT(const Variable &lhs, const Variable &rhs);
 
+        Variable unwrap(const Variable &input, int wx, int wy, int sx, int sy, int px, int py);
+        Variable wrap(const Variable &input, int ox, int oy, int wx, int wy, int sx, int sy, int px, int py);
+
+        Variable moddims(const Variable &input, int d0, int d1=-1, int d2=-1, int d3=-1);
+        Variable reorder(const Variable &input, int d0, int d1, int d2, int d3);
+        Variable conv2d(const Variable &input, const Variable &weights, int, int, int, int, int, int);
+
         Variable abs(const Variable &input);
 
         Variable flat(const Variable &input);

include/af/nn/Modules.hpp

@@ -13,3 +13,4 @@
 #include <af/nn/Modules/Container.hpp>
 #include <af/nn/Modules/Activations.hpp>
 #include <af/nn/Modules/Loss.hpp>
+#include <af/nn/Modules/Conv2D.hpp>

include/af/nn/Modules/Conv2D.hpp

@@ -0,0 +1,37 @@
+/*******************************************************
+ * 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
+ ********************************************************/
+#pragma once
+
+#include <af/nn/Modules/Module.hpp>
+
+namespace af
+{
+    namespace nn
+    {
+        class Conv2D : public Module
+        {
+        private:
+            bool m_bias;
+            int m_wx;
+            int m_wy;
+            int m_sx;
+            int m_sy;
+            int m_px;
+            int m_py;
+        public:
+            Conv2D(int wx, int wy, int sx, int sy, int px, int py, int n_in, int n_out, bool bias = true, float spread = 0.05);
+
+            Conv2D(const autograd::Variable &w, int sx = 1, int sy = 1, int px = 0, int py = 0);
+
+            Conv2D(const autograd::Variable &w, const autograd::Variable &b, int sx = 1, int sy = 1, int px = 0, int py = 0);
+
+            autograd::Variable forward(const autograd::Variable &input);
+        };
+    }
+}

src/autograd/Functions.cpp

@@ -135,8 +135,8 @@ namespace af {
             auto result = min(lhs.array(), rhs.array());
 
             auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-              inputs[0].addGrad( inputs[2] * grad_output);
-              inputs[1].addGrad(!inputs[2] * grad_output);
+                inputs[0].addGrad( inputs[2] * grad_output);
+                inputs[1].addGrad(!inputs[2] * grad_output);
             };
             return Variable(result, {lhs, rhs, mask}, grad_func);
         }
@@ -161,12 +161,12 @@ namespace af {
         }
 
 
-      INSTANTIATE_FUNCTION(max);
-      INSTANTIATE_FUNCTION(min);
+        INSTANTIATE_FUNCTION(max);
+        INSTANTIATE_FUNCTION(min);
 
 #undef INSTANTIATE_FUNCTION
 
-      Variable negate(const Variable &input)
+        Variable negate(const Variable &input)
         {
             auto result = 0.0 - input.array();
             auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
@@ -194,15 +194,6 @@ namespace af {
             return Variable(result, {input}, grad_func);
         }
 
-        Variable log(const Variable &input)
-        {
-            auto result = log(input.array());
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(grad_output / inputs[0]);
-            };
-            return Variable(result, {input}, grad_func);
-        }
-
         Variable sin(const Variable &input)
         {
             auto result = sin(input.array());
@@ -250,74 +241,31 @@ namespace af {
             return Variable(result, {input}, grad_func);
         }
 
-        Variable tileAs(const Variable &input, const Variable &reference)
+        Variable expandAs(const Variable &input, const Variable &reference)
         {
             dim4 dims(1,1,1,1);
-            dim4 rdims = reference.dims();
-            dim4 idims = input.dims();
+            dim4 idims = input.array().dims();
+            dim4 rdims = reference.array().dims();
             for (int i = 0; i < 4; i++) {
                 dims[i] = rdims[i] / idims[i];
             }
             auto result = tile(input.array(), dims);
             auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(sumAs(grad_output, inputs[0]));
+                inputs[0].addGrad(reduceAs(grad_output, inputs[0]));
             };
             return Variable(result, {input}, grad_func);
         }
 
-        Variable sumAs(const Variable &input, const Variable &reference)
+        Variable reduceAs(const Variable &input, const Variable &reference)
         {
-            dim4 rdims = reference.dims();
-            dim4 idims = input.dims();
+            dim4 idims = input.array().dims();
+            dim4 rdims = reference.array().dims();
             auto result = input.array();
             for (int i = 0; i < 4; i++) {
                 if (idims[i] != rdims[i]) result = sum(result, i);
             }
             auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(tileAs(grad_output, inputs[0]));
-            };
-            return Variable(result, {input}, grad_func);
-        }
-
-        Variable tile(const Variable &input, const std::vector<int> &repeats)
-        {
-            dim4 dims;
-            for (size_t i = 0; i < repeats.size(); i++) {
-                dims[i] = repeats[i];
-            }
-            auto result = tile(input.array(), dims);
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(sumAs(grad_output, inputs[0]));
-            };
-            return Variable(result, {input}, grad_func);
-        }
-
-        Variable sum(const Variable &input, const std::vector<int> &axes)
-        {
-            auto result = input.array();
-            for (size_t i = 0; i < axes.size(); i++) {
-                result = sum(result, axes[i]);
-            }
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(tileAs(grad_output, inputs[0]));
-            };
-            return Variable(result, {input}, grad_func);
-        }
-
-        Variable mean(const Variable &input, const std::vector<int> &axes)
-        {
-            auto result = input.array();
-            for (size_t i = 0; i < axes.size(); i++) {
-                result = mean(result, axes[i]);
-            }
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                dim4 odims = grad_output.dims();
-                dim4 idims = inputs[0].dims();
-                dim_t count = 1;
-                for (int i = 0; i < 4; i++) {
-                    count *= idims[i] / odims[i];
-                }
-                inputs[0].addGrad(count * tileAs(grad_output, inputs[0]));
+                inputs[0].addGrad(expandAs(grad_output, inputs[0]));
             };
             return Variable(result, {input}, grad_func);
         }
@@ -385,34 +333,149 @@ namespace af {
             return Variable(result, {lhs, rhs}, grad_func);
         }
 
-        Variable abs(const Variable &input)
+        Variable conv2d(const Variable &input, const Variable &weights, int wx, int wy, int sx, int sy, int px, int py)
         {
-            auto result = af::abs(input.array());
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                // af::sign returns signbit
-                // Convert it into -1, 1
-                auto sign = Variable(1 - 2 * af::sign(inputs[0].array()), false);
-                inputs[0].addGrad(sign * grad_output);
+            dim4 idims = input.array().dims();      // (x_i, y_i, c_i,  n  )
+            dim4 wdims = weights.array().dims();    // (wx,  wy,  c_i,  c_o)
+
+            int x_i = idims[0];                     //size of x dim of input
+            int y_i = idims[1];                     //size of y dim of input
+            int c_i = idims[2];                     //number of input channels
+            int n   = idims[3];                     //batch size (1 for now)
+
+            int x_o = (x_i + 2 * px - wx) / sx + 1; //size of x dim of output
+            int y_o = (y_i + 2 * py - wy) / sy + 1; //size of x dim of output
+            int c_o = wdims[3];                     //number of output channels
+
+            array windows = unwrap(input.array(), wx, wy, sx, sy, px, py);
+
+            array lhs = moddims(
+                reorder(windows, 1, 0, 2, 3),
+                x_o * y_o, wx * wy * c_i, n, 1);
+            array rhs = moddims(weights.array(), wx * wy * c_i, c_o, 1, 1);
+
+            //TODO: This loop can be replaced with a batched matmult as soon as
+            //that is added to arrayfire
+            std::vector<array> out;
+            for(int i = 0; i < n; i++){
+                array res = matmul(lhs(span, span, i), rhs);
+                out.push_back(moddims(res , x_o, y_o, c_o, 1));
+            }
+
+            //LOL @ C++ API - need this loop to have arbitrary batch size
+            array result = out[0];
+            for(int i = 1; i < n; i+=3){
+                int rem = n - i;
+                if(rem >= 3){
+                    result = join(3, result, out[i], out[i+1], out[i+2]);
+                }else if(rem == 2){
+                    result = join(3, result, out[i], out[i+1]);
+                    break;
+                }else if(rem == 1){
+                    result = join(3, result, out[i]);
+                    break;
+                }else{
+                    break;
+                }
+            }
+
+            auto grad_func = [wx, wy, sx, sy, px, py, c_i, n](std::vector<Variable> &inputs, const Variable &grad_output) {
+                dim4 odims = grad_output.array().dims();
+                dim4 wdims = inputs[1].array().dims();
+                dim4 idims = inputs[0].array().dims();
+
+                auto grad_out_reshape = moddims(grad_output, odims[0]*odims[1], odims[2], odims[3], 1);
+
+                auto weights_reshape = moddims(inputs[1], wdims[0]*wdims[1]*wdims[2], wdims[3], 1, 1);
+
+                //TODO: This really needs batched matmul...
+                //TODO: This doesn't work for n > 1
+                //TODO: Can these lines be shortened? - This seems like a large grad function - perhaps this
+                // could all be implemented in Conv2D::forward(). I had to implement the helper functions anyways
+                /*
+                std::vector<array> out;
+                for(int i = 0; i < n; i++){
+                auto a = matmulNT(grad_out_reshape(span, span, i), weights_reshape); //Problem is here - can't call () on Variable
+                auto adims = a.array().dims();
+                auto b = moddims(a, adims[0], wx*wy, c_i, adims[3]);
+                auto c = reorder(b, 1, 0, 2, 3);
+                out.push_pack(wrap(c, idims[0], idims[1], wx, wy, sx, sy, px, py));
+                }
+
+                array result = out[0];
+                for(int i = 1; i < n; i+=3){
+                    int rem = n - i;
+                    if(rem >= 3){
+                        result = join(3, result, out[i], out[i+1], out[i+2]);
+                    }else if(rem == 2){
+                        result = join(3, result, out[i], out[i+1]);
+                        break;
+                    }else if(rem == 1){
+                        result = join(3, result, out[i]);
+                        break;
+                    }else{
+                        break;
+                    }
+                }
+                */
+                auto a = matmulNT(grad_out_reshape, weights_reshape);
+                auto adims = a.array().dims();
+                auto b = moddims(a, adims[0], wx*wy, c_i, adims[3]);
+                auto c = reorder(b, 1, 0, 2, 3);
+                inputs[0].addGrad(wrap(c, idims[0], idims[1], wx, wy, sx, sy, px, py));
+
+                auto d = matmulTN(inputs[2],grad_out_reshape);
+                inputs[1].addGrad(moddims(d, wx, wy, c_i, d.array().dims()[1]));
+
+                /*
+                  for(int i = 0; i < odims[3]; i++){
+                  inputs[0].addGrad(wrap(), p[0], p[1], p[2], p[3], p[4], p[5]);
+                  inputs[0].addGrad(wrap(matmulNT(Variable(lhs(span, span, span, i), false), p_tmp), p[0], p[1], p[2], p[3], p[4], p[5]));
+                  inputs[1].addGrad(matmulTN(inputs[0], Variable(lhs(span, span, span, i), true)));
+                  }
+                */
+
             };
-            return Variable(result, {input}, grad_func);
+            return Variable(result, {input, weights, Variable(lhs, false)}, grad_func);
+
         }
 
-        Variable flat(const Variable &input)
+        Variable unwrap(const Variable &input, int wx, int wy, int sx, int sy, int px, int py)
         {
-            auto result = af::flat(input.array());
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(moddims(grad_output, inputs[0].dims()));
+            array res = unwrap(input.array(), wx, wy, sx, sy, px, py);
+            dim4 d = input.array().dims();
+            auto grad_func = [d, wx, wy, sx, sy, px, py](std::vector<Variable> &inputs, const Variable &grad_output) {
+                inputs[0].addGrad(wrap(grad_output, d[0], d[1], wx, wy, sx, sy, px, py));
             };
-            return Variable(result, {input}, grad_func);
+            return Variable(res, {input}, grad_func);
         }
 
-        Variable moddims(const Variable &input, const dim4 &dims)
+        Variable wrap(const Variable &input, int ox, int oy, int wx, int wy, int sx, int sy, int px, int py)
         {
-            auto result = af::moddims(input.array(), dims);
-            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
-                inputs[0].addGrad(moddims(grad_output, inputs[0].dims()));
+            array res = wrap(input.array(), ox, oy, wx, wy, sx, sy, px, py);
+            auto grad_func = [wx, wy, sx, sy, px, py](std::vector<Variable> &inputs, const Variable &grad_output) {
+                inputs[0].addGrad(unwrap(grad_output, wx, wy, sx, sy, px, py));
             };
-            return Variable(result, {input}, grad_func);
+            return Variable(res, {input}, grad_func);
+        }
+
+        Variable reorder(const Variable &input, int d0, int d1, int d2, int d3)
+        {
+            array res = reorder(input.array(), d0, d1, d2, d3);
+
+            int tmp[] = {d0, d1, d2, d3};
+            int tmp2[4];
+            for(int i = 0; i < 4; i++){
+                tmp2[tmp[i]] = i;
+            }
+            auto reverse = Variable(array(4, tmp2), false);
+
+            auto grad_func = [tmp2](std::vector<Variable> &inputs, const Variable &grad_output){
+                inputs[0].addGrad(reorder(grad_output, tmp2[0], tmp2[1], tmp2[2], tmp2[3]));
+            };
+            return Variable(res, {input, reverse}, grad_func);
         }
+
+
     }
 }