Can I add a neural network as an auxilary variable to the Theseus Layer? How?

[MickShen7558]

❓ Questions and Help

My question is as titled.

Take examples/homography_estimation.py as an example. In each iteration of the inner optimization, I would like to have the source image transformed by the estimated transformation matrix in that iteration and then update the neural features. Can I do this? How?

[luisenp]

HI @MickShen7558, thanks a lot for the question. One pattern that we are using in some of our applications is to have the TheseusLayer as part of a larger nn.Module class, and then use class membership to access neural features in custom cost functions. I tried to illustrate this idea in the code below which does it for a toy problem; do let me know if you have questions about it.

Alternatively, while you cannot pass a full neural net as an auxiliary variable, it's definitely possible to pass its parameters as individual aux vars and then update the model's state dict. However, this would be extremely clunky IMO.

Finally, the other pattern is that, when your cost function doesn't depend directly on the neural network but only on a set of features that will be constant throughout the optimization (i.e., the features are not a function of optim vars), then you should be able to precompute the features before calling TheseusLayer.forward() and passing them to Theseus as an auxiliary variable.

import theseus as th
import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # some arbitrary NN
        self.nn = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))

        # Add a theseus layer with a single cost function whose error depends on the NN
        objective = th.Objective()
        x = th.Vector(2, name="x")
        y = th.Vector(1, name="y")
        # This cost function computes `err(x) = nn(x) - y`
        objective.add(th.AutoDiffCostFunction([x], self._error_fn, 1, aux_vars=[y]))
        optimizer = th.LevenbergMarquardt(objective, max_iterations=10)
        self.layer = th.TheseusLayer(optimizer)

    def _error_fn(self, optim_vars, aux_vars):
        x = optim_vars[0].tensor
        y = aux_vars[0].tensor
        err = self.nn(x) - y
        return err

    # Run theseus so that NN(x*) is close to y
    def forward(self, y):
        x0 = torch.ones(y.shape[0], 2)
        sol, info = self.layer.forward(
            {"x": x0, "y": y}, optimizer_kwargs={"damping": 0.1}
        )
        print("Optim error: ", info.last_err.item())
        return sol["x"]


# Outer loop will modify NN weights to make x* as small as possible, while
# inner loop guarantees that NN(x*) is close to y
m = Model()
optim = torch.optim.Adam(m.nn.parameters(), lr=0.01)
y = torch.ones(1, 1)
for i in range(5):
    optim.zero_grad()
    xopt = m.forward(y)
    loss = (xopt**2).sum()
    loss.backward()
    optim.step()
    print("Outer loss:", loss.item(), "\n------------------------")

[MickShen7558]

Hi @luisenp,
Thank you for your great example!