Theseus-教程1

Theseus 最小二乘优化

本教程演示了如何使用 Theseus 解决曲线拟合问题。
本教程中的示例灵感来源于 Ceres 的 教程，并按照 Ceres 中的 曲线拟合示例 和 鲁棒曲线拟合示例 的结构进行设计。

二次曲线拟合

在本教程中，我们将展示如何拟合一个二次函数：

$$y = ax^2 + b$$

步骤 0：生成数据

我们首先从二次函数 $x^2 + 0.5$ 中采样点来生成数据。
在这些数据上，我们添加高斯噪声，标准差 $\sigma = 0.01$。

import torch

torch.manual_seed(0)

def generate_data(num_points=100, a=1, b=0.5, noise_factor=0.01):
    # 生成数据：从上面提到的二次曲线中采样 100 个点
    data_x = torch.rand((1, num_points))
    noise = torch.randn((1, num_points)) * noise_factor
    data_y = a * data_x.square() + b + noise
    return data_x, data_y

data_x, data_y = generate_data()

# 绘制数据
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.scatter(data_x, data_y);
ax.set_xlabel('x');
ax.set_ylabel('y');

我们演示如何使用 Theseus 解决这个曲线拟合问题，共分三个步骤：

• 步骤 1：表示数据和变量
• 步骤 2：设置优化问题
• 步骤 3：运行优化

步骤 1：在 Theseus 中表示数据和变量

如我们在教程 0 中描述的，Theseus 变量在语义上主要分为两类：

• 优化变量（optimization variables）：将被非线性最小二乘优化器修改，以最小化总代价函数。
• 辅助变量（auxiliary variables）：代价函数在优化过程中需要使用，但不会被非线性最小二乘优化器优化，例如本例中的应用数据（我们后面会看到更多示例）。

我们的第一步是将数据 (x, y) 以及优化变量 (a 和 b) 用 Theseus 数据结构表示。
优化变量必须是 Manifold 类型。对于本例，我们选择其 Vector 子类来表示 a 和 b。
因为它们是一维量，所以在初始化这些 Vector 对象时只需要 1 个自由度。
（或者，我们也可以将两个变量表示为一个二维的 Vector 对象；不过这样会改变误差函数的写法。）

辅助变量（数据变量）可以是任意 Variable 类型的实例。对于本例，直接使用 Variable 类型即可。

import theseus as th

# 数据为 Variable 类型
x = th.Variable(data_x, name="x")
y = th.Variable(data_y, name="y")

# 优化变量为 Vector 类型，具有 1 个自由度（dof）
a = th.Vector(1, name="a")
b = th.Vector(1, name="b")

步骤 2：设置优化问题

最小二乘拟合的残差误差由 CostFunction 捕获。在本例中，我们将使用 Theseus 提供的 AutoDiffCostFunction，它提供了一种简便方式来表示任意代价函数。
AutoDiffCostFunction 只需要我们定义优化变量和辅助变量，并提供一个计算残差误差的误差函数。之后，它会利用 PyTorch 的 autograd 自动微分来计算优化变量的雅可比矩阵。

在下面的示例中，quad_error_fn 捕获了使用两个一维 Vector 对象 a 和 b 拟合二次函数的最小二乘误差。

总代价误差可以通过两种方式捕获：

1. 一个 100 维的 **AutoDiffCostFunction**（每个维度表示一个数据点的误差）；
2. 100 个一维 AutoDiffCostFunction（每个代价函数表示一个数据点的误差）。

本例中我们使用前者（即一个 100 维的 AutoDiffCostFunction），在教程 4 和 5 中将看到后者的示例。

最后，我们将代价函数组合成 Theseus 的优化问题：

• 优化准则由 Objective 表示，通过将所有代价函数添加到 Objective 中来构建。
• 然后可以选择优化器并设置一些默认配置（例如下面示例中使用 GaussNewton，max_iterations=15）。
• Objective 及其关联的优化器随后用于构建 TheseusLayer，它表示一层优化。

def quad_error_fn(optim_vars, aux_vars):
    a, b = optim_vars 
    x, y = aux_vars
    est = a.tensor * x.tensor.square() + b.tensor
    err = y.tensor - est
    return err

optim_vars = a, b
aux_vars = x, y
cost_function = th.AutoDiffCostFunction(
    optim_vars, quad_error_fn, 100, aux_vars=aux_vars, name="quadratic_cost_fn"
)
objective = th.Objective()
objective.add(cost_function)
optimizer = th.GaussNewton(
    objective,
    max_iterations=15,
    step_size=0.5,
)
theseus_optim = th.TheseusLayer(optimizer)

步骤 3：运行优化**

现在运行优化问题只需提供输入数据和初始值，然后在 TheseusLayer 上调用 forward 函数即可。

输入以字典的形式提供，其中键表示优化变量（与其初始值对应）或辅助变量（与其数据对应）。下面的 theseus_inputs 字典展示了一个示例。

有了这些输入，我们就可以在 Theseus 中运行最小二乘优化。通过在 TheseusLayer 上调用 forward 函数即可完成。每次调用 forward 后会返回两个量：

1. updated_inputs 对象：包含优化变量的最终值，以及未改变的辅助变量值。这使我们能够将 updated_inputs 用作下游函数或 Theseus 层的输入（例如，对于需要多次 forward 的问题，后续教程 2 中会看到示例）。
2. info 对象：如果需要，可以跟踪最优解，并包含优化过程中的其他有用信息。跟踪最优解很有用，因为优化算法在误差比前一次迭代增大时并不会停止。（在进行反向传播时，最优解的作用不大，因为反向传播会使用整个优化序列；参见教程 2。）

theseus_inputs = {
"x": data_x,
"y": data_y,
"a": 2 * torch.ones((1, 1)),
"b": torch.ones((1, 1))
}
with torch.no_grad():
    updated_inputs, info = theseus_optim.forward(
        theseus_inputs, optimizer_kwargs={"track_best_solution": True, "verbose":True})
print("Best solution:", info.best_solution)

# Plot the optimized function
fig, ax = plt.subplots()
ax.scatter(data_x, data_y);

a = info.best_solution['a'].squeeze()
b = info.best_solution['b'].squeeze()
x = torch.linspace(0., 1., steps=100)
y = a*x*x + b
ax.plot(x, y, color='k', lw=4, linestyle='--',
        label='Optimized quadratic')
ax.legend()

ax.set_xlabel('x');
ax.set_ylabel('y');

得到输出结果为：

Nonlinear optimizer. Iteration: 0. Error: 38.42743682861328
Nonlinear optimizer. Iteration: 1. Error: 9.609884262084961
Nonlinear optimizer. Iteration: 2. Error: 2.405491828918457
Nonlinear optimizer. Iteration: 3. Error: 0.6043925285339355
Nonlinear optimizer. Iteration: 4. Error: 0.15411755442619324
Nonlinear optimizer. Iteration: 5. Error: 0.04154873266816139
Nonlinear optimizer. Iteration: 6. Error: 0.013406438753008842
Nonlinear optimizer. Iteration: 7. Error: 0.006370890885591507
Nonlinear optimizer. Iteration: 8. Error: 0.0046120136976242065
Nonlinear optimizer. Iteration: 9. Error: 0.0041722883470356464
Nonlinear optimizer. Iteration: 10. Error: 0.004062363877892494
Nonlinear optimizer. Iteration: 11. Error: 0.004034876357764006
Nonlinear optimizer. Iteration: 12. Error: 0.004028005059808493
Nonlinear optimizer. Iteration: 13. Error: 0.004026289563626051
Nonlinear optimizer. Iteration: 14. Error: 0.004025860223919153
Nonlinear optimizer. Iteration: 15. Error: 0.00402575358748436
Best solution: {'a': tensor([[0.9945]]), 'b': tensor([[0.5018]])}

我们可以看到，我们几乎完全恢复了用于采样二次函数的原始 a 和 b 值。

鲁棒二次曲线拟合

这个示例也可以应用于加权误差的问题，例如使用 Cauchy 损失函数，可以降低极端误差数据点的权重。这类似于 Ceres 求解器中的鲁棒曲线拟合示例。

在本教程中，我们做了一个简单修改，为误差函数添加 Cauchy 损失加权：
我们在 AutoDiffCostFunction 中将上面的 quad_error_fn 替换为 cauchy_loss_quad_error_fn，从而对误差进行加权。

def cauchy_fn(x):
    return torch.sqrt(0.5 * torch.log(1 + x ** 2))

def cauchy_loss_quad_error_fn(optim_vars, aux_vars):
    err = quad_error_fn(optim_vars, aux_vars)
    return cauchy_fn(err)

wt_cost_function = th.AutoDiffCostFunction(
    optim_vars, cauchy_loss_quad_error_fn, 100, aux_vars=aux_vars, name="cauchy_quad_cost_fn"
)

与上面的示例类似，我们现在可以使用这个加权代价函数构建 Theseus 优化问题：创建 Objective、选择一个优化器、构建 TheseusLayer，然后运行优化。

objective = th.Objective()
objective.add(wt_cost_function)
optimizer = th.GaussNewton(
    objective,
    max_iterations=20,
    step_size=0.3,
)
theseus_optim = th.TheseusLayer(optimizer)
theseus_inputs = {
"x": data_x,
"y": data_y,
"a": 2 * torch.ones((1, 1)),
"b": torch.ones((1, 1))
}

# 在此次优化调用中我们忽略警告，因为我们观察到使用该数据时，
# Cauchy 损失在接近最优值时经常会导致数值计算出现奇异系统。
# 注意：如果启用了 torch 的梯度追踪，在 forward 优化过程中出现奇异系统会抛出错误。
import warnings
warnings.simplefilter("ignore")   

with torch.no_grad():
    _, info = theseus_optim.forward(
        theseus_inputs, optimizer_kwargs={"track_best_solution": True, "verbose":True})
print("最优解:", info.best_solution)

# 绘制优化后的函数
fig, ax = plt.subplots()
ax.scatter(data_x, data_y);

a = info.best_solution['a'].squeeze()
b = info.best_solution['b'].squeeze()
x = torch.linspace(0., 1., steps=100)
y = a*x*x + b
ax.plot(x, y, color='k', lw=4, linestyle='--',
        label='优化后的二次函数')
ax.legend()

ax.set_xlabel('x');
ax.set_ylabel('y');

得到如下结果：

Nonlinear optimizer. Iteration: 0. Error: 13.170896530151367
Nonlinear optimizer. Iteration: 1. Error: 5.595989227294922
Nonlinear optimizer. Iteration: 2. Error: 2.6045525074005127
Nonlinear optimizer. Iteration: 3. Error: 1.248531460762024
Nonlinear optimizer. Iteration: 4. Error: 0.6063637733459473
Nonlinear optimizer. Iteration: 5. Error: 0.2966576814651489
Nonlinear optimizer. Iteration: 6. Error: 0.14604422450065613
Nonlinear optimizer. Iteration: 7. Error: 0.0725097507238388
Nonlinear optimizer. Iteration: 8. Error: 0.03653936833143234
Nonlinear optimizer. Iteration: 9. Error: 0.018927372992038727
Nonlinear optimizer. Iteration: 10. Error: 0.01030135527253151
Nonlinear optimizer. Iteration: 11. Error: 0.0060737887397408485
Best solution: {'a': tensor([[1.0039]]), 'b': tensor([[0.5112]])}

为了更高效地解决曲线拟合问题，也可以编写一个自定义的 CostFunction 并提供闭式雅可比矩阵，而不是使用 AutoDiffCostFunction 依赖 PyTorch 的数值计算雅可比矩阵。
我们将在教程 3 中展示一个示例。