facebookresearch / theseus

A library for differentiable nonlinear optimization
MIT License
1.78k stars 128 forks source link

Creating custom intrinsic matrices for use in error_fn in AutoDiffCostFunction #644

Open Unturned3 opened 6 months ago

Unturned3 commented 6 months ago

❓ Questions and Help

Thanks for this amazing piece of software, first of all!

I am trying to create an error function for use with th.AutoDiffCostFunction to measure the reprojection error from one camera to another, which involves computing a homography matrix in the form of Kj * Rj * Ri.inverse() * Ki.inverse(), where Ki and Ri are the intrinsic matrix and (rotation) pose of a camera respectively.

Ki is parameterized by a single variable: the camera's focal length f, and take the following form (due to some unfortunate coordinate system conventions):

[[-f, 0, 320],
 [ 0, f, 240],
 [ 0, 0,   1]]

I tried the following code for creating the custom Ki:

...

f0 = th.Vector(tensor=torch.tensor([[640]]).float(), name='f0')
f1 = th.Vector(tensor=torch.tensor([[640]]).float(), name='f1')
...

def error_fn(optim_vars: list[th.Variable],
             aux_vars: list[th.Variable]):

    logR0, f0, logR1, f1 = [v.tensor for v in optim_vars]
    src_pts, dst_pts = [v.tensor for v in aux_vars]

    R0 = kornia.geometry.axis_angle_to_rotation_matrix(logR0)
    K0 = torch.zeros(1, 3, 3)
    K0[:, 0, 0] = -f0[:, 0]
    K0[:, 1, 1] = f0[:, 0]
    K0[:, 0, 2] = 320
    K0[:, 1, 2] = 240
    K0[:, 2, 2] = 1

    R1 = ...
    K1 = ...

    H = K1 @ R1 @ R0.inverse() @ K0.inverse()
    return reprojection_error(H, src_pts, dst_pts)

...

cost = th.AutoDiffCostFunction(
    optim_vars=(logR0, f0, logR1, f1),
    err_fn=error_fn,
    dim=N,
    aux_vars=(src_pts, dst_pts),
)

But when I tried to run the optimizer, Theseus gave the following error for the assignment K0[:, 0, 0] = -f0[:, 0]:

vmap: inplace arithmetic(self, *extra_args) is not possible because there exists a 
Tensor `other` in extra_args that has more elements than `self`. This happened 
due to `other` being vmapped over but `self` not being vmapped over in a vmap. 
Please try to use out-of-place operators instead of inplace arithmetic.

It seems that I cannot do in-place operations like assignments. May I ask what is the proper/recommended way to create these custom matrices for use with an error function intended for th.AutoDiffCostFunction?

Thanks!

luisenp commented 5 months ago

Hi @Unturned3. Thanks for your interest in Theseus, and my sincere apologies for the long response time.

I'm not 100% sure, but it's possible that this error is the results of using torch.zeros(1, 3, 3) when creating K0. Inside functions being vmaped, usually the right thing to use is torch.new_zeros(), which vmap knows how to batch with the right dimensions. Does the error happen if you do K0 = torch.new_zeros(3, 3)?

You could also try without vmap, by setting autograd_mode="dense" when creating the autodiff cost function, although I suspect this would be too slow for your use case.