Integrate Theseus with a differentiable render.

[neanea04]

Hello, I am trying to build an image based optimization using theseus. The problem is we need a differentiable render. I was trying to use pytorch3d for this task. The issue is pytorch3d is not vmap friendly. See this issue in their repo for more details. Are there any suggestions for using Theseus with differentiable rendering or any guideline to use pytorch3d with theseus. Any help would be great. Thanks!

[luisenp]

Hi @neanea04, one work around would be to set autograd_mode="dense" when you construct your AutoDiffCostFunction. The downside is that if you have a large batch size or many cost functions of the same type it's probably going to be much slower.

[neanea04]

Thanks for the reply. Is there a way to solve without compromising processing time? Your suggestion seems to be working but it is really slow for my application. What is the difference between dense autograd and normal one?

[luisenp]

dense autograd uses torch.autograd.functional.jacobians to compute the jacobians, instead of vmap and jacrev. The problem is that if you have a batch size of B, this function computes jacobians across batch elements (i.e., the jacobian of the j-th batch output wrt to the i-th batch input); hence the term "dense". Since we only need per-batch jacobians, most of these are unnecessary.

Note that even if your problem batch size is 1, if you have many cost functions of the same type, these are aggregated together into a single batch, so there can still be some overhead of using "dense" vs. "vmap".

Unfortunately, I don't think there are any easy solutions other than finding a vmap-friendly differentiable renderer (or writing your own).