simple_example.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# A minimal example using Theseus that is fitting a curve to a dataset of observations.

import torch

import theseus as th


def y_model(x, c):
    return c * torch.exp(x)


def generate_data(num_points=10, c=0.5):
    data_x = torch.linspace(-1, 1, num_points).view(1, -1)
    data_y = y_model(data_x, c)
    return data_x, data_y


def read_data():
    num_points = 20
    data_x, data_y_clean = generate_data(num_points)
    return data_x, data_y_clean, 0.5 * torch.ones(1, 1)


x_true, y_true, v_true = read_data()  # shapes (1, N), (1, N), (1, 1)
x = th.Variable(torch.randn_like(x_true), name="x")
y = th.Variable(y_true, name="y")
v = th.Vector(1, name="v")  # a manifold subclass of Variable for optim_vars


def error_fn(optim_vars, aux_vars):  # returns y - v * exp(x)
    x, y = aux_vars
    return y.tensor - optim_vars[0].tensor * torch.exp(x.tensor)


objective = th.Objective()
cost_function = th.AutoDiffCostFunction(
    [v], error_fn, y_true.shape[1], aux_vars=[x, y], cost_weight=th.ScaleCostWeight(1.0)
)
objective.add(cost_function)
layer = th.TheseusLayer(th.GaussNewton(objective, max_iterations=10))

phi = torch.nn.Parameter(x_true + 0.1 * torch.ones_like(x_true))
outer_optimizer = torch.optim.Adam([phi], lr=0.001)
for epoch in range(20):
    solution, info = layer.forward(
        input_tensors={"x": phi.clone(), "v": torch.ones(1, 1)},
        optimizer_kwargs={"backward_mode": "implicit"},
    )
    outer_loss = torch.nn.functional.mse_loss(solution["v"], v_true)
    outer_loss.backward()
    outer_optimizer.step()
    print("Outer loss: ", outer_loss.item())