WIP Indexing autograd and tests, new loss functions

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "extern/googletest"]
+	path = extern/googletest
+	url = https://github.com/google/googletest.git

CMakeLists.txt

@@ -35,3 +35,9 @@ set_target_properties(afml
   CXX_STANDARD 11)
 
 add_subdirectory(examples)
+
+option(PACKAGE_TESTS "Build tests" ON)
+if(PACKAGE_TESTS)
+    enable_testing()
+    add_subdirectory(tests)
+endif()

examples/CMakeLists.txt

@@ -15,4 +15,3 @@ endfunction(build_example)
 # build_example(Node.cpp)
 build_example(xor.cpp)
 # build_example(Weights.cpp)
-build_example(autograd.cpp)

include/af/autograd/Functions.hpp

@@ -44,6 +44,8 @@ namespace af {
         Variable operator <=(const Variable &lhs, const double &rhs);
 
         Variable operator !(const Variable &input);
+        Variable lookup(const Variable &input, const Variable &idx);
+        Variable assign(const Variable &input, const Variable &idx, const Variable &vals);
 
         Variable negate(const Variable &input);
         Variable reciprocal(const Variable &input);
@@ -54,6 +56,7 @@ namespace af {
         Variable cos(const Variable &input);
         Variable tanh(const Variable &input);
         Variable sigmoid(const Variable &input);
+        Variable softmax(const Variable &input);
 
         Variable max(const Variable &lhs, const Variable &rhs);
         Variable max(const Variable &lhs, const double &rhs);

include/af/nn/Modules/Loss.hpp

@@ -59,9 +59,36 @@ namespace af
                                        const autograd::Variable &weights);
         };
 
+        class CrossEntropyLoss : public Loss
+        {
+        public:
+            CrossEntropyLoss() {}
+
+            autograd::Variable forward(const autograd::Variable &inputs,
+                                       const autograd::Variable &targets);
+
+            autograd::Variable forward(const autograd::Variable &inputs,
+                                       const autograd::Variable &targets,
+                                       const autograd::Variable &weights);
+        };
+
+        class MultiMarginLoss : public Loss
+        {
+        public:
+            MultiMarginLoss() {}
+
+            autograd::Variable forward(const autograd::Variable &inputs,
+                                       const autograd::Variable &targets);
+
+            autograd::Variable forward(const autograd::Variable &inputs,
+                                       const autograd::Variable &targets,
+                                       const autograd::Variable &weights);
+        };
+
         typedef MeanSquaredError MSE;
         typedef MeanAbsoluteError MAE;
         typedef MeanAbsoluteError L1Loss;
         typedef BinaryCrossEntropyLoss BCELoss;
+        typedef CrossEntropyLoss CELoss;
     }
 }

src/autograd/Functions.cpp

@@ -117,6 +117,33 @@ namespace af {
             return Variable(result, false);
         }
 
+        Variable lookup(const Variable &input, const Variable &idx)
+        {
+            af::array result = input.array()(idx.array());
+
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                af::array grad = af::constant(0, inputs[0].dims());
+                grad(inputs[1].array()) = grad_output.array();
+
+                inputs[0].addGrad(Variable(grad, false));
+            };
+            return Variable(result, {input, idx}, grad_func);
+        }
+
+        Variable assign(const Variable &input, const Variable &idx, const Variable &vals)
+        {
+            af::array result = input.array();
+            result(idx.array()) = vals.array();
+
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                af::array mask = af::constant(1, inputs[0].dims(), s32);
+                mask(inputs[1].array()) = 0;
+
+                inputs[0].addGrad(Variable(mask, false) * grad_output);
+            };
+            return Variable(result, {input, idx, vals}, grad_func);
+        }
+
         Variable max(const Variable &lhs, const Variable &rhs)
         {
             auto mask = lhs > rhs;
@@ -241,6 +268,24 @@ namespace af {
             return Variable(result, {input}, grad_func);
         }
 
+        Variable softmax(const Variable &input)
+        {
+            //todo: add axis to apply?
+            auto exps = exp(input.array());
+            auto result = exps / tile(sum(exps, 0), exps.dims(0));
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                auto exps = exp(inputs[0]);
+                auto tmp = exps / tileAs(sum(exps, {0}), exps);
+
+                auto ps_j = tile(tmp, { 1, (int)tmp.dims()[0] } );
+                auto ps_i = transpose(tile(tmp, {1,(int)tmp.dims()[0] } ));
+                Variable I(identity((int)tmp.dims()[0], (int)tmp.dims()[0]), false);
+                auto jac = (sum(ps_i * (I - ps_j), { 1 }));
+                inputs[0].addGrad(grad_output * jac);
+            };
+            return Variable(result, {input}, grad_func);
+        }
+
         Variable transpose(const Variable &input)
         {
             auto result = transpose(input.array());
@@ -281,7 +326,7 @@ namespace af {
 
         Variable tile(const Variable &input, const std::vector<int> &repeats)
         {
-            dim4 dims;
+            dim4 dims(0);
             for (size_t i = 0; i < repeats.size(); i++) {
                 dims[i] = repeats[i];
             }