Super-speedup by specializing to 2-dimensional output (for now). Using some template magic.
parent
31c138d0d6
commit
9ebe1e6d10
|
|
@ -33,7 +33,12 @@ typedef std::map<Key, Matrix> JacobianMap;
|
||||||
|
|
||||||
//-----------------------------------------------------------------------------
|
//-----------------------------------------------------------------------------
|
||||||
/// The JacobinaTrace class records a tree-structured expression's execution
|
/// The JacobinaTrace class records a tree-structured expression's execution
|
||||||
|
template<class T>
|
||||||
struct JacobianTrace {
|
struct JacobianTrace {
|
||||||
|
|
||||||
|
// Some fixed matrix sizes we need
|
||||||
|
typedef Eigen::Matrix<double, 2, T::dimension> Jacobian2T;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* The Pointer class is a tagged union that obviates the need to create
|
* The Pointer class is a tagged union that obviates the need to create
|
||||||
* a JacobianTrace subclass for Constants and Leaf Expressions. Instead
|
* a JacobianTrace subclass for Constants and Leaf Expressions. Instead
|
||||||
|
|
@ -70,7 +75,6 @@ struct JacobianTrace {
|
||||||
content.ptr = trace;
|
content.ptr = trace;
|
||||||
}
|
}
|
||||||
// Either insert identity into Jacobians (Leaf) or propagate (Function)
|
// Either insert identity into Jacobians (Leaf) or propagate (Function)
|
||||||
template<class T>
|
|
||||||
void reverseAD(JacobianMap& jacobians) const {
|
void reverseAD(JacobianMap& jacobians) const {
|
||||||
if (type == Leaf) {
|
if (type == Leaf) {
|
||||||
size_t n = T::Dim();
|
size_t n = T::Dim();
|
||||||
|
|
@ -89,17 +93,45 @@ struct JacobianTrace {
|
||||||
} else if (type == Function)
|
} else if (type == Function)
|
||||||
content.ptr->reverseAD(dTdA, jacobians);
|
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
|
/// Make sure destructor is virtual
|
||||||
virtual ~JacobianTrace() {
|
virtual ~JacobianTrace() {
|
||||||
}
|
}
|
||||||
virtual void reverseAD(JacobianMap& jacobians) const = 0;
|
virtual void reverseAD(JacobianMap& jacobians) const = 0;
|
||||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const = 0;
|
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const = 0;
|
||||||
// template<class JacobianFT>
|
virtual void reverseAD2(const Jacobian2T& dFdT,
|
||||||
// void reverseAD(const JacobianFT& dFdT, JacobianMap& jacobians) const {
|
JacobianMap& jacobians) const = 0;
|
||||||
};
|
};
|
||||||
|
|
||||||
typedef JacobianTrace::Pointer TracePtr;
|
template<size_t M, class A>
|
||||||
|
struct Select {
|
||||||
|
typedef Eigen::Matrix<double, M, A::dimension> Jacobian;
|
||||||
|
static void reverseAD(const typename JacobianTrace<A>::Pointer& trace,
|
||||||
|
const Jacobian& dTdA, JacobianMap& jacobians) {
|
||||||
|
trace.reverseAD(dTdA, jacobians);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template<class A>
|
||||||
|
struct Select<2, A> {
|
||||||
|
typedef Eigen::Matrix<double, 2, A::dimension> Jacobian;
|
||||||
|
static void reverseAD(const typename JacobianTrace<A>::Pointer& trace,
|
||||||
|
const Jacobian& dTdA, JacobianMap& jacobians) {
|
||||||
|
trace.reverseAD2(dTdA, jacobians);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
//template <class Derived>
|
//template <class Derived>
|
||||||
//struct TypedTrace {
|
//struct TypedTrace {
|
||||||
|
|
@ -243,7 +275,8 @@ public:
|
||||||
virtual Augmented<T> forward(const Values& values) const = 0;
|
virtual Augmented<T> forward(const Values& values) const = 0;
|
||||||
|
|
||||||
/// Construct an execution trace for reverse AD
|
/// 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<T>::Pointer& p) const = 0;
|
||||||
};
|
};
|
||||||
|
|
||||||
//-----------------------------------------------------------------------------
|
//-----------------------------------------------------------------------------
|
||||||
|
|
@ -280,7 +313,8 @@ public:
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Construct an execution trace for reverse AD
|
/// Construct an execution trace for reverse AD
|
||||||
virtual T traceExecution(const Values& values, TracePtr& p) const {
|
virtual T traceExecution(const Values& values,
|
||||||
|
typename JacobianTrace<T>::Pointer& p) const {
|
||||||
return constant_;
|
return constant_;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
@ -320,7 +354,8 @@ public:
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Construct an execution trace for reverse AD
|
/// Construct an execution trace for reverse AD
|
||||||
virtual T traceExecution(const Values& values, TracePtr& p) const {
|
virtual T traceExecution(const Values& values,
|
||||||
|
typename JacobianTrace<T>::Pointer& p) const {
|
||||||
p.setLeaf(key_);
|
p.setLeaf(key_);
|
||||||
return values.at<T>(key_);
|
return values.at<T>(key_);
|
||||||
}
|
}
|
||||||
|
|
@ -329,22 +364,22 @@ public:
|
||||||
|
|
||||||
//-----------------------------------------------------------------------------
|
//-----------------------------------------------------------------------------
|
||||||
/// Unary Function Expression
|
/// Unary Function Expression
|
||||||
template<class T, class A>
|
template<class T, class A1>
|
||||||
class UnaryExpression: public ExpressionNode<T> {
|
class UnaryExpression: public ExpressionNode<T> {
|
||||||
|
|
||||||
public:
|
public:
|
||||||
|
|
||||||
typedef Eigen::Matrix<double, T::dimension, A::dimension> JacobianTA;
|
typedef Eigen::Matrix<double, T::dimension, A1::dimension> JacobianTA;
|
||||||
typedef boost::function<T(const A&, boost::optional<JacobianTA&>)> Function;
|
typedef boost::function<T(const A1&, boost::optional<JacobianTA&>)> Function;
|
||||||
|
|
||||||
private:
|
private:
|
||||||
|
|
||||||
Function function_;
|
Function function_;
|
||||||
boost::shared_ptr<ExpressionNode<A> > expressionA_;
|
boost::shared_ptr<ExpressionNode<A1> > expressionA1_;
|
||||||
|
|
||||||
/// Constructor with a unary function f, and input argument e
|
/// Constructor with a unary function f, and input argument e
|
||||||
UnaryExpression(Function f, const Expression<A>& e) :
|
UnaryExpression(Function f, const Expression<A1>& e) :
|
||||||
function_(f), expressionA_(e.root()) {
|
function_(f), expressionA1_(e.root()) {
|
||||||
}
|
}
|
||||||
|
|
||||||
friend class Expression<T> ;
|
friend class Expression<T> ;
|
||||||
|
|
@ -353,18 +388,18 @@ public:
|
||||||
|
|
||||||
/// Return keys that play in this expression
|
/// Return keys that play in this expression
|
||||||
virtual std::set<Key> keys() const {
|
virtual std::set<Key> keys() const {
|
||||||
return expressionA_->keys();
|
return expressionA1_->keys();
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Return value
|
/// Return value
|
||||||
virtual T value(const Values& values) const {
|
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
|
/// Return value and derivatives
|
||||||
virtual Augmented<T> forward(const Values& values) const {
|
virtual Augmented<T> forward(const Values& values) const {
|
||||||
using boost::none;
|
using boost::none;
|
||||||
Augmented<A> argument = this->expressionA_->forward(values);
|
Augmented<A1> argument = this->expressionA1_->forward(values);
|
||||||
JacobianTA dTdA;
|
JacobianTA dTdA;
|
||||||
T t = function_(argument.value(),
|
T t = function_(argument.value(),
|
||||||
argument.constant() ? none : boost::optional<JacobianTA&>(dTdA));
|
argument.constant() ? none : boost::optional<JacobianTA&>(dTdA));
|
||||||
|
|
@ -372,26 +407,33 @@ public:
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Trace structure for reverse AD
|
/// Trace structure for reverse AD
|
||||||
struct Trace: public JacobianTrace {
|
struct Trace: public JacobianTrace<T> {
|
||||||
TracePtr trace;
|
typename JacobianTrace<A1>::Pointer trace1;
|
||||||
JacobianTA dTdA;
|
JacobianTA dTdA1;
|
||||||
|
|
||||||
/// Start the reverse AD process
|
/// Start the reverse AD process
|
||||||
virtual void reverseAD(JacobianMap& jacobians) const {
|
virtual void reverseAD(JacobianMap& jacobians) const {
|
||||||
trace.reverseAD(dTdA, jacobians);
|
Select<T::dimension, A1>::reverseAD(trace1, dTdA1, jacobians);
|
||||||
}
|
}
|
||||||
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
||||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
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<double, 2, T::dimension> Jacobian2T;
|
||||||
|
virtual void reverseAD2(const Jacobian2T& dFdT,
|
||||||
|
JacobianMap& jacobians) const {
|
||||||
|
trace1.reverseAD2(dFdT * dTdA1, jacobians);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
/// Construct an execution trace for reverse AD
|
/// Construct an execution trace for reverse AD
|
||||||
virtual T traceExecution(const Values& values, TracePtr& p) const {
|
virtual T traceExecution(const Values& values,
|
||||||
|
typename JacobianTrace<T>::Pointer& p) const {
|
||||||
Trace* trace = new Trace();
|
Trace* trace = new Trace();
|
||||||
p.setFunction(trace);
|
p.setFunction(trace);
|
||||||
A a = this->expressionA_->traceExecution(values, trace->trace);
|
A1 a = this->expressionA1_->traceExecution(values, trace->trace1);
|
||||||
return function_(a, trace->dTdA);
|
return function_(a, trace->dTdA1);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
@ -454,25 +496,34 @@ public:
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Trace structure for reverse AD
|
/// Trace structure for reverse AD
|
||||||
struct Trace: public JacobianTrace {
|
struct Trace: public JacobianTrace<T> {
|
||||||
TracePtr trace1, trace2;
|
typename JacobianTrace<A1>::Pointer trace1;
|
||||||
|
typename JacobianTrace<A2>::Pointer trace2;
|
||||||
JacobianTA1 dTdA1;
|
JacobianTA1 dTdA1;
|
||||||
JacobianTA2 dTdA2;
|
JacobianTA2 dTdA2;
|
||||||
|
|
||||||
/// Start the reverse AD process
|
/// Start the reverse AD process
|
||||||
virtual void reverseAD(JacobianMap& jacobians) const {
|
virtual void reverseAD(JacobianMap& jacobians) const {
|
||||||
trace1.reverseAD(dTdA1, jacobians);
|
Select<T::dimension, A1>::reverseAD(trace1, dTdA1, jacobians);
|
||||||
trace2.reverseAD(dTdA2, jacobians);
|
Select<T::dimension, A2>::reverseAD(trace2, dTdA2, jacobians);
|
||||||
}
|
}
|
||||||
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
||||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
||||||
trace1.reverseAD(dFdT * dTdA1, jacobians);
|
trace1.reverseAD(dFdT * dTdA1, jacobians);
|
||||||
trace2.reverseAD(dFdT * dTdA2, jacobians);
|
trace2.reverseAD(dFdT * dTdA2, jacobians);
|
||||||
}
|
}
|
||||||
|
/// Version specialized to 2-dimensional output
|
||||||
|
typedef Eigen::Matrix<double, 2, T::dimension> 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
|
/// Construct an execution trace for reverse AD
|
||||||
virtual T traceExecution(const Values& values, TracePtr& p) const {
|
virtual T traceExecution(const Values& values,
|
||||||
|
typename JacobianTrace<T>::Pointer& p) const {
|
||||||
Trace* trace = new Trace();
|
Trace* trace = new Trace();
|
||||||
p.setFunction(trace);
|
p.setFunction(trace);
|
||||||
A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1);
|
A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1);
|
||||||
|
|
@ -553,17 +604,19 @@ public:
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Trace structure for reverse AD
|
/// Trace structure for reverse AD
|
||||||
struct Trace: public JacobianTrace {
|
struct Trace: public JacobianTrace<T> {
|
||||||
TracePtr trace1, trace2, trace3;
|
typename JacobianTrace<A1>::Pointer trace1;
|
||||||
|
typename JacobianTrace<A2>::Pointer trace2;
|
||||||
|
typename JacobianTrace<A3>::Pointer trace3;
|
||||||
JacobianTA1 dTdA1;
|
JacobianTA1 dTdA1;
|
||||||
JacobianTA2 dTdA2;
|
JacobianTA2 dTdA2;
|
||||||
JacobianTA3 dTdA3;
|
JacobianTA3 dTdA3;
|
||||||
|
|
||||||
/// Start the reverse AD process
|
/// Start the reverse AD process
|
||||||
virtual void reverseAD(JacobianMap& jacobians) const {
|
virtual void reverseAD(JacobianMap& jacobians) const {
|
||||||
trace1.reverseAD(dTdA1, jacobians);
|
Select<T::dimension, A1>::reverseAD(trace1, dTdA1, jacobians);
|
||||||
trace2.reverseAD(dTdA2, jacobians);
|
Select<T::dimension, A2>::reverseAD(trace2, dTdA2, jacobians);
|
||||||
trace3.reverseAD(dTdA3, jacobians);
|
Select<T::dimension, A3>::reverseAD(trace3, dTdA3, jacobians);
|
||||||
}
|
}
|
||||||
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
/// Given df/dT, multiply in dT/dA and continue reverse AD process
|
||||||
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
virtual void reverseAD(const Matrix& dFdT, JacobianMap& jacobians) const {
|
||||||
|
|
@ -571,10 +624,19 @@ public:
|
||||||
trace2.reverseAD(dFdT * dTdA2, jacobians);
|
trace2.reverseAD(dFdT * dTdA2, jacobians);
|
||||||
trace3.reverseAD(dFdT * dTdA3, jacobians);
|
trace3.reverseAD(dFdT * dTdA3, jacobians);
|
||||||
}
|
}
|
||||||
|
/// Version specialized to 2-dimensional output
|
||||||
|
typedef Eigen::Matrix<double, 2, T::dimension> 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
|
/// Construct an execution trace for reverse AD
|
||||||
virtual T traceExecution(const Values& values, TracePtr& p) const {
|
virtual T traceExecution(const Values& values,
|
||||||
|
typename JacobianTrace<T>::Pointer& p) const {
|
||||||
Trace* trace = new Trace();
|
Trace* trace = new Trace();
|
||||||
p.setFunction(trace);
|
p.setFunction(trace);
|
||||||
A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1);
|
A1 a1 = this->expressionA1_->traceExecution(values, trace->trace1);
|
||||||
|
|
|
||||||
|
|
@ -117,10 +117,10 @@ public:
|
||||||
Augmented<T> augmented(const Values& values) const {
|
Augmented<T> augmented(const Values& values) const {
|
||||||
#define REVERSE_AD
|
#define REVERSE_AD
|
||||||
#ifdef REVERSE_AD
|
#ifdef REVERSE_AD
|
||||||
JacobianTrace::Pointer pointer;
|
typename JacobianTrace<T>::Pointer pointer;
|
||||||
T value = root_->traceExecution(values,pointer);
|
T value = root_->traceExecution(values,pointer);
|
||||||
Augmented<T> augmented(value);
|
Augmented<T> augmented(value);
|
||||||
pointer.reverseAD<T>(augmented.jacobians());
|
pointer.reverseAD(augmented.jacobians());
|
||||||
return augmented;
|
return augmented;
|
||||||
#else
|
#else
|
||||||
return root_->forward(values);
|
return root_->forward(values);
|
||||||
|
|
|
||||||
|
|
@ -81,7 +81,7 @@ TEST(Expression, leaf) {
|
||||||
TEST(Expression, test) {
|
TEST(Expression, test) {
|
||||||
|
|
||||||
// Test Constant expression
|
// Test Constant expression
|
||||||
Expression<int> c(0);
|
Expression<Rot3> c(Rot3::identity());
|
||||||
|
|
||||||
// Create leaves
|
// Create leaves
|
||||||
Expression<Pose3> x(1);
|
Expression<Pose3> x(1);
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue