C++
基本用法¶
在本节中,我们将演示如何使用XAD来计算前向模式和伴随模式的一阶导数。
例如,我们选择一个具有4个输入和1个输出变量的简单函数,定义如下:
double f(double x0, double x1, double x2, double x3)
{
double a = sin(x0) * cos(x1);
double b = x2 * x3 - tan(x1 - x2);
double c = a + 2* c;
return c*c;
}
我们将在这一点计算该函数的导数:
前提条件:替换活动变量¶
为了使用XAD对这个函数进行自动微分,我们首先需要用XAD提供的活动数据类型替换所有独立数据类型及其依赖值。在上述函数中,所有变量都依赖于输入参数,因此所有出现的double类型都必须被替换。
这可以通过以下两种方式之一实现:
- 变量可以直接替换,根据所需的微分模式。例如,对于前向模式,
double被替换为类型FReal;对于伴随模式,则替换为类型AReal。 - 该函数被设计为模板,因此可以使用任何数据类型调用,包括原始的
double类型。
在本教程中我们选择第二种方法,因此函数变为:
template <class T>
T f(T x0, T x1, T x2, T x3)
{
T a = sin(x0) * cos(x1);
T b = x2 * x3 - tan(x1 - x2);
T c = a + 2* b;
return c*c;
}
这意味着我们可以使用相同的定义来处理正向和伴随模式。
前向模式¶
如算法微分背景:前向模式所示,当应用于单输出函数时,算法微分的前向模式每次只能计算一个导数。为了说明这一点,我们选择对输入变量x0求导。
要启动前向模式,我们首先需要用适当的类型声明活跃变量。XAD提供了便捷的类型定义来选择微分模式,详情请参阅AD Mode Interface。对于前向模式,我们可以按如下方式声明所需类型:
然后我们可以使用AD类型定义来处理我们的变量。
下一步是初始化因变量,只需将输入值赋给AD类型的新变量即可:
对于前向模式,我们现在需要用值1来为感兴趣的变量设定初始导数(如算法微分背景:前向模式中所述),如下所示:
全局函数derivative是一个便捷函数,适用于任何活跃数据类型。或者,我们也可以使用成员函数FReal::setDerivative。
此时我们可以调用函数,它将计算函数值以及我们感兴趣的导数:
现在我们可以通过输出结果上的value和derivative函数(或成员函数FReal::getDerivative和FReal::getValue)来访问这些结果。例如,以下代码将它们输出到控制台:
另请参阅
这个示例包含在XAD中(fwd_1st)。
伴随模式¶
对于当前这个具有单一输出和多个输入的函数来说,自动微分的伴随模式是自然选择。我们可以在一次执行中就获得全部四个导数。
伴随模式需要一个磁带(tape)来记录计算过程中的操作及其值。在设置输出的伴随项后,可以回放该磁带来计算输入的伴随项。
活动数据类型和磁带类型都可以从接口结构adj中获取:
计算伴随变量的第一步是初始化tape::
这会调用默认构造函数Tape::Tape,该函数会创建并激活记录带。
接下来,我们创建输入变量并将其注册到tape中:
AD x0_ad = x0;
AD x1_ad = x1;
AD x2_ad = x2;
AD x3_ad = x3;
tape.registerInput(x0);
tape.registerInput(x1);
tape.registerInput(x2);
tape.registerInput(x3);
请注意,只有被注册为磁带输入变量的变量及其依赖变量会被记录。还需注意,在注册活动变量之前,当前线程需要拥有一个活动磁带。为确保线程安全,应用程序的每个线程都可以拥有自己的活动磁带。
一旦设置好自变量,我们就可以开始在记录带上记录导数并运行算法:
在此阶段,我们已经记录了所有运算并完成了数值计算。现在我们需要将输出值也注册到记录带中,然后才能按照算法微分基础:伴随模式中所述,用1来初始化输出的伴随导数:
这里使用全局函数derivative,它会返回给定参数的存储导数(或伴随)的引用。或者也可以使用成员函数AReal::setAdjoint或AReal::setDerivative来实现相同目的。
剩下的工作就是解读磁带以计算自变量的伴随导数:
现在我们可以通过全局函数derivative或成员函数AReal::getDerivative来访问输入的伴随变量,这些正是我们感兴趣的导数:
std::cout << "y = " << value(y) << "\n"
<< "dy/dx0 = " << derivative(x0_ad) << "\n"
<< "dy/dx1 = " << derivative(x1_ad) << "\n"
<< "dy/dx2 = " << derivative(x2_ad) << "\n"
<< "dy/dx3 = " << derivative(x3_ad) << "\n";
另请参阅
此示例包含在XAD的adj_1st中。
最佳实践¶
当被评估算法的输出数量少于输入数量时,应优先选择伴随模式。然而,当只需要少量导数时(例如少于5个),通过使用前向模式可以避免为计算图分配内存。建议通过实验来确定给定算法的最佳模式。
外部函数¶
通常,需要求导的算法部分可能没有源代码可用。例如,可能会调用外部数学库中的例程。重新实现这些部分可能并不理想(出于性能或开发工作量的考虑),在这种情况下,需要以某种形式手动实现该函数的导数。
这可以通过以下任一方式实现:
- 对库函数应用有限差分法(bumping)
- 手动实现函数的伴随代码,或者
- 通过解析方式计算导数,可能会使用其他库函数。
在这些情况下,可以使用XAD的外部函数接口来集成手动导数,如下所述。通过同样的技术,可以对应用程序中性能或内存关键的部分进行手动调优。
示例算法¶
我们选取一个计算多维向量长度的示例算法。其定义如下:
目标是通过伴随模式计算\(y\)相对于所有输入向量元素的导数。
该算法可以用C++代码实现如下:
std::vector<double> xsqr(n);
for (int i = 0; i < n; ++i)
xsqr[i] = x[i] * x[i];
double y = sqrt(sum_elements(x, n));
在本示例中,我们假设sum_elements是一个外部函数,由我们没有源代码的库实现。其函数原型为:
伴随模式的外部函数¶
要使用外部函数,我们遵循以下步骤:
- 在调用时,将输入的活动变量值转换为底层基本数据类型 (
double) - 被动调用外部函数
- 将结果值赋给活跃的输出变量,以便继续录制磁带
- 使用检查点回调对象存储输入和输出的磁带槽位,并将其注册到磁带中。
- 在计算伴随导数时,该回调函数需要加载输出的伴随导数,手动将其传播到输入,并通过这些值增加输入伴随导数。
我们将所有功能封装到一个回调对象中。我们继承自CheckpointCallback基类,并至少实现虚方法CheckpointCallback::computeAdjoint。该方法在磁带回滚过程中被调用。我们还将正向计算放在同一个对象内(这也可以在回调类外部完成)。我们的回调类声明如下:
template <class Tape>
class ExternalSumElementsCallback : public xad::CheckpointCallback<Tape>
{
public:
typedef typename Tape::slot_type slot_type; // type for slot in the tape
typedef typename Tape::value_type value_type; // double
typedef typename Tape::active_type active_type; // AReal<double>
active_type computeExternal(const active_type* x, int n); // forward compute
void computeAdjoint(Tape* tape) override; // adjoint compute
private:
std::vector<slot_type> inputSlots_; // slots of inputs in tape
slot_type outputSlot_; // slot of output in tape
};
我们将其声明为适用于任意磁带类型的模板,这是一种良好的实践,因为它允许在更高阶导数中也能复用此实现。
computeExternal 方法¶
在computeExternal方法中,我们首先将输入变量的槽位存储在磁带中,因为在伴随计算期间我们需要这些信息来增加相应的伴随量。我们使用inputSlots_成员向量来保存这些信息:
然后我们创建活动输入的副本,并将它们存储在一个被动值向量中,通过这个向量我们可以调用外部函数:
std::vector<value_type> x_p(n);
for (int i = 0; i < n; ++i)
x_p[i] = value(x[i]);
value_type y = sum_elements(&x_p[0], n);
我们现在需要将这个结果存储在一个活动变量中,将其注册为外部函数的输出(以便磁带能继续记录相关变量),并保留其在磁带中的位置以供后续的伴随计算使用:
最后我们需要将回调函数插入到磁带中,从而要求在磁带的反向传播过程中调用它,并返回:
computeAdjoint 方法¶
computeAdjoint方法在XAD执行磁带回滚时被调用。我们需要重写这个方法并实现手动伴随代码。对于一个简单的求和操作来说,这很直接:由于所有偏导数都是1,所有输入伴随值都等于输出伴随值。因此我们需要获取输出伴随值,并将所有输入伴随值增加这个值:
value_type output_adj = tape->getAndResetOutputAdjoint(outputSlot_);
for (int i = 0; i < inputSlots_.size(); ++i)
tape->incrementAdjoint(inputSlots_[i], output_adj);
函数 Tape::getAndResetOutputAdjoint 获取与给定槽位对应的伴随值并将其重置为零。通常需要这种重置操作,因为在正向计算中输出变量可能覆盖了其他值。函数 Tape::incrementAdjoint 则简单地通过给定值递增指定槽位的伴随值。
包装函数¶
通过设置检查点回调类,我们可以为AReal实现一个sum_elements重载函数,该函数封装了这个回调的创建过程:
template <class T>
xad::AReal<T> sum_elements(const xad::AReal<T>* x, int n)
{
typedef typename xad::AReal<T>::tape_type tape_type;
tape_type* tape = tape_type::getActive();
ExternalSumElementsCallback<tape_type>* ckp =
new ExternalSumElementsCallback<tape_type>;
tape->pushCallback(ckp);
return ckp->computeExternal(x, n);
}
该函数动态分配检查点回调对象,并通过Tape::pushCallback函数让磁带管理其销毁过程。此调用确保当磁带被销毁时回调对象也会被销毁,从而避免内存泄漏。如果回调对象由其他方式管理,则无需此调用。随后它将计算重定向到检查点回调类的computeExternal函数。通过使用这个包装类,sum_elements函数可以像原始外部函数sum_elements处理double类型那样用于活动类型。将其定义为模板允许我们在未来需要时复用此函数来计算高阶导数。
调用站点¶
调用站点可以按如下方式实现(假设x_ad是保存已注册在磁带上的独立变量的向量):
tape.newRecording();
std::vector<AD> xsqr(n);
for (int i = 0; i < n; ++i)
xsqr[i] = x_ad[i] * x_ad[i];
AD y = sqrt(sum_elements(xsqr.data(), n)); // calls external function wrapper
tape.registerOutput(y);
derivative(y) = 1.0;
tape.computeAdjoints();
std::cout << "y = " << value(y) << "\n";
for (int i = 0; i < n; ++i)
std::cout << "dy/dx" << i << " = " << derivative(x[i]) << "\n";
这与基本用法中给出的流程完全一致。
另请参阅
此示例随XAD (external_function) 提供。
前向模式的外部函数¶
由于前向模式不涉及磁带,需要手动实现导数计算并与值计算同时进行。可以使用derivative函数直接将手动计算的导数更新到输出值中。
在我们的示例中,我们可以用前向模式实现外部函数:
template <class T>
xad::FReal<T> sum_elements(const xad::FReal<T>* x, int n)
{
typedef xad::FReal<T> active_type;
std::vector<T> x_p(n);
for (int i = 0; i < n; ++i)
x_p[i] = value(x[i]);
T y_p = sum_elements(&x_p[0], n);
active_type y = y_p;
for (int i = 0; i < n; ++i)
derivative(y) += derivative(x[i]);
return y;
}
我们首先从x向量中提取被动值,并调用外部库函数获取被动输出值y_p。然后将该值赋给主动输出变量y,同时将其导数初始化为0。
由于本例中是一个简单的求和运算,输出的导数是输入导数的总和,这是通过最后的循环计算得出的。
另请参阅
此示例随XAD (external_function) 提供。
检查点¶
检查点技术是一种用于减少伴随模式算法微分中磁带内存占用的方法。不同于将整个算法记录在磁带上(这在大型计算中可能迅速导致内存占用达到GB级别),该方法选择性地记录算法的特定阶段,每次只记录一个阶段。如下图所示:
该算法分为多个阶段,每个阶段的输入数据存储在检查点中,输出结果被动计算(不记录在磁带上)。一旦计算出算法的最终输出,就会初始化输出的伴随值,并在磁带回滚期间的每个检查点进行:
- 检查点的输入被加载,
- 此阶段的操作仅记录在磁带上,
- 此阶段的输出伴随变量被初始化,
- 在此阶段回滚磁带,计算阶段输入的伴随变量,
- 输入的伴随变量通过这些值递增,并且
- 在进行上一阶段操作前,磁带会被清空。
使用此方法,磁带内存仅需记录一个算法阶段而非完整算法,从而限制了内存占用。然而,每次前向计算需要执行两次,因此检查点机制是以计算量换取内存空间。
在实践中,由于使用较少内存会带来更高的缓存效率,尽管需要执行更多计算,检查点技术的整体速度可能仍比记录完整算法更快。
示例算法¶
为了演示检查点方法,我们选择对一个单一输入重复应用正弦函数的简单示例:
我们将for循环划分为等距阶段,并在每个阶段插入一个检查点。
检查点回调¶
要创建检查点,我们需要将阶段的输入以及磁带中用于输入和输出的槽位存储在一个继承自CheckpointCallback的回调对象中。需要重写虚方法CheckpointCallback::computeAdjoint来执行每个阶段的伴随计算。由于所有阶段都是相同的,我们选择在单个回调对象中实现所有检查点的功能,并将所需的输入存储在堆栈数据结构中。或者,我们也可以在每次检查点时创建一个新的检查点回调对象。我们的回调原型是:
template <class Tape>
class SinCheckpointCallback : public xad::CheckpointCallback<Tape>
{
public:
typedef typename Tape::slot_type slot_type; // type for slot in the tape
typedef typename Tape::value_type value_type; // double
typedef typename Tape::active_type active_type; // AReal<double>
active_type computeStage(int n, active_type& x); // forward computation
void computeAdjoint(Tape* tape) override; // adjoint computation
private:
std::stack<int> n_; // number of iterations in this stage
std::stack<value_type> x_; // input values for this stage
std::stack<slot_type> slots_; // tape slots for input and output
};
为了方便实现,我们在computeStage方法中将单阶段的前向计算添加到了同一个类中,这个计算也可以在对象外部执行。
computeStage 方法¶
在computeStage方法中,我们首先将输入值、迭代次数和输入的槽位存储在检查点对象中:
然后我们计算包含被动变量(不记录在磁带上的)的阶段:
输出活动变量的值需要用结果更新,并且我们还需要在检查点中存储输出变量的槽位:
注意,我们不需要像之前外部函数示例那样将x注册为磁带的输出,因为该变量已经在磁带上注册(它既是输入也是输出)。
剩下的工作是将这个回调对象注册到磁带中,以便在磁带回滚时调用其computeAdjoint方法:
computeAdjoint 方法¶
computeAdjoint方法会在XAD在磁带检查点时自动调用。我们首先需要加载这个计算阶段的输入并获取输出的伴随导数:
slot_type outputidx = slots_.top(); slots_.pop();
slot_type inputidx = slots_.top(); slots_.pop();
int n = n_.top(); n_.pop();
value_type outputadj = tape->getAndResetOutputAdjoint(outputslot);
函数Tape::getAndResetOutputAdjoint会读取给定槽位对应的伴随值并将其重置为0。这种重置通常是必需的,因为算法中对应槽位的变量可能会被重复使用(覆盖),正如repeated_sin函数中的情况所示。
我们现在想使用XAD来计算仅针对此计算阶段的伴随导数。这是通过在全局磁带中创建一个嵌套记录来实现的,该记录可以单独回滚:
active_type x = x_.top(); // local independent variable
x_.pop();
tape->registerInput(x); // need to register to record
xad::ScopedNestedRecording<Tape> nested(tape); // nested recording
repeated_sin(n, x_ad); // run actively
tape->registerOutput(x); // register x as an output
derivative(x) = output_adj; // set output adjoint
nested.computeAdjoints(); // rollback nested tape
nested.incrementAdjoint(inputslot, derivative(x)); // incr. input adjoint
与简单伴随模式类似(参见基本用法),我们首先将局部自变量初始化为活跃数据类型并开始嵌套记录。这是通过创建一个类型为ScopedNestedRecording的局部对象nested来实现的,该对象在其构造函数中封装了对Tape::newNestedRecording的调用,并在析构函数中封装了对Tape::endNestedRecording的调用。建议使用ScopedNestedRecording来确保在离开作用域时始终完成嵌套记录。
接下来我们通过主动运行算法来记录这一阶段的操作。然后设置输出的伴随变量并计算输入的伴随变量。这样就能递增这一阶段输入的伴随变量。
请注意,当nested对象超出作用域时(即调用其析构函数时),该计算阶段的嵌套磁带将被清除,内存可被前一阶段重用。这有助于节省总体内存。
调用站点¶
带有检查点的完整算法可以按如下方式启动:
tape_type tape;
AD x_ad = x; // initialized indepedent variables
tape.registerInput(x_ad); // register with the tape
tape.newRecording(); // start recording derivatives
SinCheckpointCallback<tape_type> chkpt; // setup checkpointing object
int checkpt_distance = 4; // we checkpoint every 4 iterations
for (int i = 0; i < n; i += checkpt_distance)
{
int m = min(checkpt_distance, n-i);
chkpt.computeStage(m, x_ad); // one computation stage
}
tape.registerOutput(x_ad);
derivative(x_ad) = 1.0;
tape.computeAdjoints();
std::cout << "xout = " << value(x_ad) << "\n"
<< "dxout/dxin = " << derivative(x_ad) << "\n";
这个过程基本遵循基本用法中的相同步骤,但需要设置检查点对象,并在算法的每个阶段(本例中为4次迭代)调用其computeStage成员函数。
注意
重要的是当调用Tape::computeAdjoints时,检查点回调对象必须有效。在此之前不应销毁该对象。
请参阅Checkpoint Callback Memory Management了解如何将基于tape的销毁机制与动态分配检查点回调结合使用。
另请参阅
此示例随XAD(checkpointing)一起提供。
其他使用模式¶
可以使用其他方法在检查点的CheckpointCallback::computeAdjoint方法中更新伴随变量,例如:
- 在外层伴随模式中使用前向模式算法微分
- 有限差分法(扰动法)
- 解析导数
- 外部库函数(参见External Functions)
检查点也可以递归使用,即在检查点的嵌套磁带中创建新的检查点。
这些方法的优势高度依赖于具体应用场景。
高阶导数¶
如Algorithmic Differentiation Background: Higher Orders所述,高阶导数可以通过嵌套一阶算法微分技术来计算。例如,可以通过在伴随模式上计算前向模式来获得二阶导数。利用XAD,可以直接应用该技术来计算高阶导数。
XAD的自动微分接口结构(参见AD Mode Interface)定义了二阶模式数据类型以便于访问。三阶或更高阶的类型需要从基本的一阶类型手动定义。
接下来我们将使用前向-伴随模式来演示二阶导数。
示例算法¶
出于演示目的,我们使用与基本用法相同的算法:
template <class T>
T f(T x0, T x1, T x2, T x3)
{
T a = sin(x0) * cos(x1);
T b = x2 * x3 - tan(x1 - x2);
T c = a + 2* b;
return c*c;
}
我们对这一点的导数感兴趣:
前向模式优于伴随模式¶
在此模式下,我们可以计算所有一阶导数(作为通过伴随法导出的单输出函数可给出所有一阶导数),以及二阶导数海森矩阵的第一行。完整海森矩阵定义为:
请注意,Hessian矩阵通常是对称的,这一特性可用于减少计算完整Hessian矩阵所需的工作量。
第一步是设置本次计算所需的磁带和活动数据类型:
typedef xad::fwd_adj<double> mode;
typedef mode::tape_type tape_type;
typedef mode::active_type AD;
tape_type tape;
请注意,此模式的活动类型实际上是AReal<FReal<double>>。
现在我们需要设置自变量并注册它们:
AD x0_ad = x0;
AD x1_ad = x1;
AD x2_ad = x2;
AD x3_ad = x3;
tape.registerInput(x0_ad);
tape.registerInput(x1_ad);
tape.registerInput(x2_ad);
tape.registerInput(x3_ad);
当我们使用前向模式计算二阶导数时,在运行算法之前需要为二阶导数设定初始导数值:
内部对value的调用会获取外部类型的值,即返回FReal<double>类型的值,我们将其导数设置为1。
现在我们可以开始在记录带上记录导数并运行算法:
对于内部伴随模式,我们需要注册输出并用1初始化初始伴随:
在这里,内部对derivative的调用给出了外部类型的导数,即伴随模式活动类型的导数。其类型为FReal<double>,我们将其值设为1。
接下来我们计算伴随变量,这将同时计算一阶和二阶导数:
我们现在可以输出结果:
以及一阶导数:
std::cout << "dy/dx0 = " << value(derivative(x0_ad)) << "\n"
<< "dy/dx1 = " << value(derivative(x1_ad)) << "\n"
<< "dy/dx2 = " << value(derivative(x2_ad)) << "\n"
<< "dy/dx3 = " << value(derivative(x3_ad)) << "\n";
再次注意,内部调用derivative会获取外部活动数据类型的导数,因此它返回一个表示一阶伴随值的FReal<double>引用。我们可以通过value调用将该值作为double类型获取。
关于 x0 的二阶导数可以这样获得:
std::cout << "d2y/dx0dx0 = " << derivative(derivative(x0_ad)) << "\n"
<< "d2y/dx0dx1 = " << derivative(derivative(x1_ad)) << "\n"
<< "d2y/dx0dx2 = " << derivative(derivative(x2_ad)) << "\n"
<< "d2y/dx0dx3 = " << derivative(derivative(x3_ad)) << "\n";
该功能用于"解包"一阶和二阶活动类型的导数。
给定输入运行应用程序的结果是:
y = 7.69565
dy/dx0 = 0.21205
dy/dx1 = -16.2093
dy/dx2 = 24.8681
dy/dx3 = 14.4253
d2y/dx0dx0 = -0.327326
d2y/dx0dx1 = -3.21352
d2y/dx0dx2 = 0.342613
d2y/dx0dx3 = 0.198741
正向伴随模式是计算二阶导数的推荐方法。
另请参阅
此示例随XAD (fwd_adj_2nd)一同提供。
其他二阶模式¶
其他二阶模式的工作方式类似。以下是对它们的简要描述。
前向覆盖前向¶
在前向-前向模式下,不需要使用计算磁带,两种阶数的导数都需要在运行算法前进行初始化。如果函数只有一个输出,这种方法可以计算一个海森矩阵元素和一个一阶导数。该模式下的导数初始化顺序通常是:
value(derivative(x)) = 1.0; // initialize the first-order derivative
derivative(value(x)) = 1.0; // initialize the second-order derivative
计算完成后,一阶导数可以这样获取:
以及二阶导数为:
通过不同的初始种子设定,可以获得Hessian矩阵的不同元素。
伴随模式优于前向模式¶
这里内部模式采用前向传播,以无磁带方式计算一个导数;外部模式采用伴随模式,需要磁带记录。使用此模式时,我们需要用以下方式初始化前向传播导数:
由于输出的导数对应一阶结果,我们需要在运行算法后设置其导数(即伴随变量):
在完成磁带解析后,我们现在可以获取一阶导数:
由于这种操作模式的对称性,相同的一阶导数也可以通过以下方式获得:
在此模式下可以获取所有输入的一阶导数,类似于前向-伴随模式。
二阶导数可以通过以下方式获得:
伴随的伴随¶
由于两种嵌套模式都是伴随模式,这种模式需要为两种顺序准备两条记录带。因此接口结构adj_adj中定义的类型需要包含内部和外部记录带类型:
typedef xad::adj_adj<double> mode;
typedef mode::inner_tape_type inner_tape_type;
typedef mode::outer_tape_type outer_tape_type;
typedef mode::active_type AD;
在此模式下,无需设置初始导数,但重要的是在运行算法前需要初始化两个记录带,并在两者上启动新的记录。
执行完成后,外部导数需要初始化为:
然后外部磁带需要计算伴随导数。这会计算value(derivative(x))作为输出,在解释内部磁带之前需要先设置这个导数:
在内层磁带调用computeAdjoints()后,我们可以读取一阶导数如下:
以及二阶导数为: