Failure to fit complex data.

[zhangrentu]

❓ Questions and Help

When attempting to fit some simple regression problems, such as y = ax + b, where both x and y are complex numbers, I encountered errors. Could you please advise on methods or modifications to resolve this issue?

[luisenp]

Hi @zhangrentu. We have never tested Theseus with complex numbers, so I'm not really sure what parts of it would need to be modified. Certainly our custom linear solvers will not work with this type of data, so in the best possible case you are limited to use DenseLinearization and either CholeskyDenseSolver or LUDenseSolver, but even with this there is no guarantee that it will work.

If you send me a code snippet of code I can try to run and perhaps suggest what modifications would be necessary.

[luisenp]

Hi @zhangrentu. I made a number of tweaks and removed some of our dtype constraints (add torch.complex32 here), and eventually failed with a torch error when trying to run your error_fn(). You can see my changes below

import torch
import theseus as th


def y_model(x, a, b, c):
   return a * torch.exp((-1j * b + c) * x)  # y = a * exp((-1ja + b) * x)


def generate_data(num_points=512, a=1, b=2, c=3):
   data_x = torch.linspace(0, 1, num_points).view(1, -1)
   data_y = y_model(data_x, a, b, c)
   return data_x, data_y


def read_data():
   data_x, data_y_clean = generate_data()
   return (
       data_x,
       data_y_clean,
       1 * torch.ones(1, 1),
       2 * torch.ones(1, 1),
       3 * torch.ones(1, 1),
   )


x_true, y_true, a_true, b_true, c_true = read_data()
x = th.Variable(torch.randn_like(x_true), name="x")
y = th.Variable(y_true, name="y")
a = th.Vector(1, name="a")  # a manifold subclass of Variable for optim_vars
b = th.Vector(1, name="b")  # a manifold subclass of Variable for optim_vars
c = th.Vector(1, name="c")  # a manifold subclass of Variable for optim_vars
for v in [x, y, a, b, c]:
   v.to(dtype=torch.complex32)


def error_fn(optim_vars, aux_vars):  # returns y - a * exp((-1j*a + b) * x)
   x, y = aux_vars
   return y.tensor - optim_vars[0].tensor * torch.exp(
       (-1j * optim_vars[1].tensor + optim_vars[2].tensor) * x.tensor
   )


objective = th.Objective(dtype=torch.complex32)
w = th.ScaleCostWeight(1.0)
w.to(dtype=torch.complex32)
cost_function = th.AutoDiffCostFunction(
   [a, b, c], error_fn, y_true.shape[1], aux_vars=[x, y], cost_weight=w
)
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)
input_tensors = {
   "x": phi.clone(),
   "a": 0.5 * torch.ones(1, 1),
   "b": torch.ones(1, 1),
   "c": torch.ones(1, 1),
}
input_tensors = {k: t.to(dtype=torch.complex32) for k, t in input_tensors.items()}
for epoch in range(20):
   solution, info = layer.forward(
       input_tensors=input_tensors,
       optimizer_kwargs={"backward_mode": "implicit"},
   )
   outer_loss1 = torch.nn.functional.mse_loss(solution["a"], a_true)
   outer_loss2 = torch.nn.functional.mse_loss(solution["b"], b_true)
   outer_loss3 = torch.nn.functional.mse_loss(solution["c"], c_true)
   outer_loss = outer_loss1 + outer_loss2 + outer_loss3
   outer_loss.backward()
   outer_optimizer.step()
   print("Outer loss: ", outer_loss.item())

[zhangrentu]

Thank you very much for your response. I added the data attribute (torch.complex64) and performed calculations with PyTorch, but I found that the results are not converging. Do you have any suggestions?

[luisenp]

Can you share the new version of your script?

[zhangrentu]

Thanks. The following is the new version, with the added dtype constraints (torch.complex64)

import torch
import theseus as th


def y_model(x, a, b, c):
   return a * torch.exp((-1j * b + c) * x)  # y = a * exp((-1ja + b) * x)


def generate_data(num_points=512, a=1, b=2, c=3):
   data_x = torch.linspace(0, 1, num_points).view(1, -1)
   data_y = y_model(data_x, a, b, c)
   return data_x, data_y


def read_data():
   data_x, data_y_clean = generate_data()
   return (
       data_x,
       data_y_clean,
       1 * torch.ones(1, 1),
       2 * torch.ones(1, 1),
       3 * torch.ones(1, 1),
   )


x_true, y_true, a_true, b_true, c_true = read_data()
x = th.Variable(torch.randn_like(x_true), name="x")
y = th.Variable(y_true, name="y")
a = th.Vector(1, name="a")  # a manifold subclass of Variable for optim_vars
b = th.Vector(1, name="b")  # a manifold subclass of Variable for optim_vars
c = th.Vector(1, name="c")  # a manifold subclass of Variable for optim_vars
for v in [x, y, a, b, c]:
   v.to(dtype=torch.complex64)


def error_fn(optim_vars, aux_vars):  # returns y - a * exp((-1j*a + b) * x)
   x, y = aux_vars
   return y.tensor - optim_vars[0].tensor * torch.exp(
       (-1j * optim_vars[1].tensor + optim_vars[2].tensor) * x.tensor
   )


objective = th.Objective(dtype=torch.complex64)
w = th.ScaleCostWeight(1.0)
w.to(dtype=torch.complex64)
cost_function = th.AutoDiffCostFunction(
   [a, b, c], error_fn, y_true.shape[1], aux_vars=[x, y], cost_weight=w
)
objective.add(cost_function)
optimizer = th.LevenbergMarquardt(
    objective,
    th.CholeskyDenseSolver,
    max_iterations=10,
    step_size=0.001,
)
layer = th.TheseusLayer(optimizer)

phi = torch.nn.Parameter(x_true + 0.1 * torch.ones_like(x_true))
outer_optimizer = torch.optim.Adam([phi], lr=0.001)
input_tensors = {
   "x": phi.clone(),
   "a": 0.5 * torch.ones(1, 1),
   "b": torch.ones(1, 1),
   "c": torch.ones(1, 1),
}
input_tensors = {k: t.to(dtype=torch.complex64) for k, t in input_tensors.items()}
for epoch in range(20):
   solution, info = layer.forward(
       input_tensors=input_tensors,
       optimizer_kwargs={"backward_mode": "implicit"},
   )
   outer_loss1 = torch.nn.functional.mse_loss(solution["a"], a_true)
   outer_loss2 = torch.nn.functional.mse_loss(solution["b"], b_true)
   outer_loss3 = torch.nn.functional.mse_loss(solution["c"], c_true)
   outer_loss = outer_loss1 + outer_loss2 + outer_loss3
   outer_loss.backward()
   outer_optimizer.step()
   print("Outer loss: ", outer_loss.item())

[zhangrentu]

In the process of attempting parallel network（torch.nn.DataParallel） computations, the following error was encountered. A model composed of th.TheseusLayer and the rest of the network structures is placed on the same device. Do you know what problem is causing this?

[luisenp]

Hi @zhangrentu. The error in your last comment should now be fixed after #623 is merged.

[luisenp]

I took a quick look at your script. One change I had to make was to set autograd_mode="dense" in AutoDiffCostFunction (looks like vmap doesn't support complex data). I did see problems in the linear system used to compute gradients for our optimizers, but I don't have a good intuition on how to solve this because I'm not too familiar with complex data. Perhaps you need to change your error function to return some real valued error vector that depends on your complex inputs, but I'm not sure.

[luisenp]

I'm able to run the optimizer if I use a very high damping, set the error to return .abs() of the previous error, and some slight changes to the loss (although I don't see the other loss changing at all).

You can see my changes here

import torch
import theseus as th


def y_model(x, a, b, c):
    return a * torch.exp((-1j * b + c) * x)  # y = a * exp((-1ja + b) * x)


def generate_data(num_points=4, a=1, b=2, c=3):
    data_x = torch.linspace(0, 1, num_points).view(1, -1)
    data_y = y_model(data_x, a, b, c)
    return data_x, data_y


def read_data():
    data_x, data_y_clean = generate_data()
    return (
        data_x,
        data_y_clean,
        1 * torch.ones(1, 1),
        2 * torch.ones(1, 1),
        3 * torch.ones(1, 1),
    )


x_true, y_true, a_true, b_true, c_true = read_data()
x = th.Variable(torch.randn_like(x_true), name="x")
y = th.Variable(y_true, name="y")
a = th.Vector(1, name="a")  # a manifold subclass of Variable for optim_vars
b = th.Vector(1, name="b")  # a manifold subclass of Variable for optim_vars
c = th.Vector(1, name="c")  # a manifold subclass of Variable for optim_vars
for v in [x, y, a, b, c]:
    v.to(dtype=torch.complex64)


def error_fn(optim_vars, aux_vars):  # returns y - a * exp((-1j*a + b) * x)
    x, y = aux_vars
    return (
        y.tensor
        - optim_vars[0].tensor
        * torch.exp((-1j * optim_vars[1].tensor + optim_vars[2].tensor) * x.tensor)
    ).abs()


objective = th.Objective(dtype=torch.complex64)
w = th.ScaleCostWeight(1.0)
w.to(dtype=torch.complex64)
cost_function = th.AutoDiffCostFunction(
    [a, b, c],
    error_fn,
    y_true.shape[1],
    aux_vars=[x, y],
    cost_weight=w,
    autograd_mode="dense",
)
objective.add(cost_function)
optimizer = th.LevenbergMarquardt(
    objective,
    th.CholeskyDenseSolver,
    max_iterations=5,
    step_size=0.1,
)
layer = th.TheseusLayer(optimizer)

phi = torch.nn.Parameter(x_true + 0.1 * torch.ones_like(x_true))
outer_optimizer = torch.optim.Adam([phi], lr=1.0)
input_tensors = {
    "x": phi.clone(),
    "a": 0.5 * torch.ones(1, 1),
    "b": torch.ones(1, 1),
    "c": torch.ones(1, 1),
}
input_tensors = {k: t.to(dtype=torch.complex64) for k, t in input_tensors.items()}
for epoch in range(20):
    solution, info = layer.forward(
        input_tensors=input_tensors,
        optimizer_kwargs={
            "backward_mode": "unroll",
            "verbose": True,
            "damping": 100.0,
        },
    )
    outer_loss1 = torch.nn.functional.mse_loss(solution["a"].real, a_true)
    outer_loss2 = torch.nn.functional.mse_loss(solution["b"].real, b_true)
    outer_loss3 = torch.nn.functional.mse_loss(solution["c"].real, c_true)
    outer_loss = outer_loss1 + outer_loss2 + outer_loss3
    outer_loss.backward()
    outer_optimizer.step()
    print("Outer loss: ", outer_loss.item())

[zhangrentu]

Thank you very much for your response. However, when we use TheseusLayer in the model, it still returns NoneType. Additionally, optim_vars and aux_vars data are automatically placed on Cuda:0, but the objection is on cuda, which causes an error when calling objection.update, making it impossible to use data parallel training (nn.DataParallel).

[zhangrentu]

We tried changing the error function to the real part, for example, real(exp(ia)) = cos a, but the error did not converge after the modification, or the matrix is non-positive definite, etc. We also attempted to reduce the step_size, which showed a slight improvement. Are there any other adjustment methods? Currently, we are using the following configuration:
optimizer = th.LevenbergMarquardt(
objective,
th.CholeskyDenseSolver,
th.DenseLinearization,
max_iterations=20,
step_size=0.02,
)

[luisenp]

@zhangrentu This code has not been merged to main yet. Are you using the code directly from that branch? If you are, then please share a short snippet of code that results in the device error, because I'm not sure how that can happen in the code from that branch. Thanks!

[zhangrentu]

Yes, I called it from a branch. The main process of parallel computing is as follows: data is placed on the primary GPU (cuda:0), and the model is distributed to the GPUs used for parallel computing (e.g., cuda:0, cuda:1, cuda:2). Since the TheseusLayer is treated as a layer within the network, it belongs to the model part. However, during parameter updates, the objective of the cost function was not distributed, resulting in the error. The specific error is as shown in the image below:

[luisenp]

Ah, I see. We have never tested this inside a DataParallel model, so I don't have a lot of insight yet. Could you share a short repro script?

[luisenp]

@zhangrentu Regarding the convergence, in the script I shared above one thing I did was to increase the damping to a really large value (I used 100.0), and the error does seem to decrease when doing this.