From 9ebe1e6d108c4a13f7f47063064f89876a90b945 Mon Sep 17 00:00:00 2001 From: dellaert Date: Wed, 8 Oct 2014 23:50:17 +0200 Subject: [PATCH] Super-speedup by specializing to 2-dimensional output (for now). Using some template magic. --- gtsam_unstable/nonlinear/Expression-inl.h | 132 +++++++++++++----- gtsam_unstable/nonlinear/Expression.h | 4 +- .../nonlinear/tests/testExpression.cpp | 2 +- 3 files changed, 100 insertions(+), 38 deletions(-) diff --git a/gtsam_unstable/nonlinear/Expression-inl.h b/gtsam_unstable/nonlinear/Expression-inl.h index 31e70e11b..3a5dd0afc 100644 --- a/gtsam_unstable/nonlinear/Expression-inl.h +++ b/gtsam_unstable/nonlinear/Expression-inl.h @@ -33,7 +33,12 @@ typedef std::map JacobianMap; //----------------------------------------------------------------------------- /// The JacobinaTrace class records a tree-structured expression's execution +template struct JacobianTrace { + + // Some fixed matrix sizes we need + typedef Eigen::Matrix Jacobian2T; + /** * The Pointer class is a tagged union that obviates the need to create * a JacobianTrace subclass for Constants and Leaf Expressions. Instead @@ -70,7 +75,6 @@ struct JacobianTrace { content.ptr = trace; } // Either insert identity into Jacobians (Leaf) or propagate (Function) - template void reverseAD(JacobianMap& jacobians) const { if (type == Leaf) { size_t n = T::Dim(); @@ -89,17 +93,45 @@ struct JacobianTrace { } else if (type == Function) content.ptr->reverseAD(dTdA, jacobians); } + // Either add to Jacobians (Leaf) or propagate (Function) + void reverseAD2(const Jacobian2T& dTdA, JacobianMap& jacobians) const { + if (type == Leaf) { + JacobianMap::iterator it = jacobians.find(content.key); + if (it != jacobians.end()) + it->second += dTdA; + else + jacobians[content.key] = dTdA; + } else if (type == Function) + content.ptr->reverseAD2(dTdA, jacobians); + } }; + /// Make sure destructor is virtual virtual ~JacobianTrace() { } virtual void reverseAD(JacobianMap& jacobians) const = 0; virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const = 0; -// template -// void reverseAD(const JacobianFT& dFdT, JacobianMap& jacobians) const { + virtual void reverseAD2(const Jacobian2T& dFdT, + JacobianMap& jacobians) const = 0; }; -typedef JacobianTrace::Pointer TracePtr; +template +struct Select { + typedef Eigen::Matrix Jacobian; + static void reverseAD(const typename JacobianTrace::Pointer& trace, + const Jacobian& dTdA, JacobianMap& jacobians) { + trace.reverseAD(dTdA, jacobians); + } +}; + +template +struct Select<2, A> { + typedef Eigen::Matrix Jacobian; + static void reverseAD(const typename JacobianTrace::Pointer& trace, + const Jacobian& dTdA, JacobianMap& jacobians) { + trace.reverseAD2(dTdA, jacobians); + } +}; //template //struct TypedTrace { @@ -243,7 +275,8 @@ public: virtual Augmented forward(const Values& values) const = 0; /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const = 0; + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const = 0; }; //----------------------------------------------------------------------------- @@ -280,7 +313,8 @@ public: } /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const { + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const { return constant_; } }; @@ -320,7 +354,8 @@ public: } /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const { + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const { p.setLeaf(key_); return values.at(key_); } @@ -329,22 +364,22 @@ public: //----------------------------------------------------------------------------- /// Unary Function Expression -template +template class UnaryExpression: public ExpressionNode { public: - typedef Eigen::Matrix JacobianTA; - typedef boost::function)> Function; + typedef Eigen::Matrix JacobianTA; + typedef boost::function)> Function; private: Function function_; - boost::shared_ptr > expressionA_; + boost::shared_ptr > expressionA1_; /// Constructor with a unary function f, and input argument e - UnaryExpression(Function f, const Expression& e) : - function_(f), expressionA_(e.root()) { + UnaryExpression(Function f, const Expression& e) : + function_(f), expressionA1_(e.root()) { } friend class Expression ; @@ -353,18 +388,18 @@ public: /// Return keys that play in this expression virtual std::set keys() const { - return expressionA_->keys(); + return expressionA1_->keys(); } /// Return value virtual T value(const Values& values) const { - return function_(this->expressionA_->value(values), boost::none); + return function_(this->expressionA1_->value(values), boost::none); } /// Return value and derivatives virtual Augmented forward(const Values& values) const { using boost::none; - Augmented argument = this->expressionA_->forward(values); + Augmented argument = this->expressionA1_->forward(values); JacobianTA dTdA; T t = function_(argument.value(), argument.constant() ? none : boost::optional(dTdA)); @@ -372,26 +407,33 @@ public: } /// Trace structure for reverse AD - struct Trace: public JacobianTrace { - TracePtr trace; - JacobianTA dTdA; + struct Trace: public JacobianTrace { + typename JacobianTrace::Pointer trace1; + JacobianTA dTdA1; /// Start the reverse AD process virtual void reverseAD(JacobianMap& jacobians) const { - trace.reverseAD(dTdA, jacobians); + Select::reverseAD(trace1, dTdA1, jacobians); } /// Given df/dT, multiply in dT/dA and continue reverse AD process virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const { - trace.reverseAD(dFdT * dTdA, jacobians); + trace1.reverseAD(dFdT * dTdA1, jacobians); + } + /// Version specialized to 2-dimensional output + typedef Eigen::Matrix Jacobian2T; + virtual void reverseAD2(const Jacobian2T& dFdT, + JacobianMap& jacobians) const { + trace1.reverseAD2(dFdT * dTdA1, jacobians); } }; /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const { + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const { Trace* trace = new Trace(); p.setFunction(trace); - A a = this->expressionA_->traceExecution(values, trace->trace); - return function_(a, trace->dTdA); + A1 a = this->expressionA1_->traceExecution(values, trace->trace1); + return function_(a, trace->dTdA1); } }; @@ -454,25 +496,34 @@ public: } /// Trace structure for reverse AD - struct Trace: public JacobianTrace { - TracePtr trace1, trace2; + struct Trace: public JacobianTrace { + typename JacobianTrace::Pointer trace1; + typename JacobianTrace::Pointer trace2; JacobianTA1 dTdA1; JacobianTA2 dTdA2; /// Start the reverse AD process virtual void reverseAD(JacobianMap& jacobians) const { - trace1.reverseAD(dTdA1, jacobians); - trace2.reverseAD(dTdA2, jacobians); + Select::reverseAD(trace1, dTdA1, jacobians); + Select::reverseAD(trace2, dTdA2, jacobians); } /// Given df/dT, multiply in dT/dA and continue reverse AD process virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const { trace1.reverseAD(dFdT * dTdA1, jacobians); trace2.reverseAD(dFdT * dTdA2, jacobians); } + /// Version specialized to 2-dimensional output + typedef Eigen::Matrix Jacobian2T; + virtual void reverseAD2(const Jacobian2T& dFdT, + JacobianMap& jacobians) const { + trace1.reverseAD2(dFdT * dTdA1, jacobians); + trace2.reverseAD2(dFdT * dTdA2, jacobians); + } }; /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const { + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const { Trace* trace = new Trace(); p.setFunction(trace); A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1); @@ -553,17 +604,19 @@ public: } /// Trace structure for reverse AD - struct Trace: public JacobianTrace { - TracePtr trace1, trace2, trace3; + struct Trace: public JacobianTrace { + typename JacobianTrace::Pointer trace1; + typename JacobianTrace::Pointer trace2; + typename JacobianTrace::Pointer trace3; JacobianTA1 dTdA1; JacobianTA2 dTdA2; JacobianTA3 dTdA3; /// Start the reverse AD process virtual void reverseAD(JacobianMap& jacobians) const { - trace1.reverseAD(dTdA1, jacobians); - trace2.reverseAD(dTdA2, jacobians); - trace3.reverseAD(dTdA3, jacobians); + Select::reverseAD(trace1, dTdA1, jacobians); + Select::reverseAD(trace2, dTdA2, jacobians); + Select::reverseAD(trace3, dTdA3, jacobians); } /// Given df/dT, multiply in dT/dA and continue reverse AD process virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const { @@ -571,10 +624,19 @@ public: trace2.reverseAD(dFdT * dTdA2, jacobians); trace3.reverseAD(dFdT * dTdA3, jacobians); } + /// Version specialized to 2-dimensional output + typedef Eigen::Matrix Jacobian2T; + virtual void reverseAD2(const Jacobian2T& dFdT, + JacobianMap& jacobians) const { + trace1.reverseAD2(dFdT * dTdA1, jacobians); + trace2.reverseAD2(dFdT * dTdA2, jacobians); + trace3.reverseAD2(dFdT * dTdA3, jacobians); + } }; /// Construct an execution trace for reverse AD - virtual T traceExecution(const Values& values, TracePtr& p) const { + virtual T traceExecution(const Values& values, + typename JacobianTrace::Pointer& p) const { Trace* trace = new Trace(); p.setFunction(trace); A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1); diff --git a/gtsam_unstable/nonlinear/Expression.h b/gtsam_unstable/nonlinear/Expression.h index 772e6e88e..210cdcea8 100644 --- a/gtsam_unstable/nonlinear/Expression.h +++ b/gtsam_unstable/nonlinear/Expression.h @@ -117,10 +117,10 @@ public: Augmented augmented(const Values& values) const { #define REVERSE_AD #ifdef REVERSE_AD - JacobianTrace::Pointer pointer; + typename JacobianTrace::Pointer pointer; T value = root_->traceExecution(values,pointer); Augmented augmented(value); - pointer.reverseAD(augmented.jacobians()); + pointer.reverseAD(augmented.jacobians()); return augmented; #else return root_->forward(values); diff --git a/gtsam_unstable/nonlinear/tests/testExpression.cpp b/gtsam_unstable/nonlinear/tests/testExpression.cpp index 188394e0a..3747ed078 100644 --- a/gtsam_unstable/nonlinear/tests/testExpression.cpp +++ b/gtsam_unstable/nonlinear/tests/testExpression.cpp @@ -81,7 +81,7 @@ TEST(Expression, leaf) { TEST(Expression, test) { // Test Constant expression - Expression c(0); + Expression c(Rot3::identity()); // Create leaves Expression x(1);