We briefly show how the example from the earlier section on differentiable rendering can be made to work when combining differentiable rendering with an optimization expressed using PyTorch. The ability to combine these frameworks enables sandwiching Mitsuba 2 between neural layers and differentiating the combination end-to-end.
Note that communication and synchronization between Enoki and PyTorch along with the complexity of traversing two separate computation graph data structures causes an overhead of 10-20% compared to optimization implemented using only Enoki. We generally recommend sticking with Enoki unless the problem requires neural network building blocks like fully connected layers or convolutions, where PyTorch provides a clear advantage.
As before, we specify the relevant variant, load the scene, and retain relevant differentiable parameters.
import mitsuba
mitsuba.set_variant('gpu_autodiff_rgb')
import enoki as ek
from mitsuba.core import Thread, Vector3f
from mitsuba.core.xml import load_file
from mitsuba.python.util import traverse
from mitsuba.python.autodiff import render_torch, write_bitmap
import torch
import time
Thread.thread().file_resolver().append('cbox')
scene = load_file('cbox/cbox.xml')
# Find differentiable scene parameters
params = traverse(scene)
# Discard all parameters except for one we want to differentiate
params.keep(['red.reflectance.value'])
The .torch()
method can be used to convert any Enoki CUDA type
into a corresponding PyTorch tensor.
# Print the current value and keep a backup copy
param_ref = params['red.reflectance.value'].torch()
print(param_ref)
The render_torch()
function works
analogously to render()
except that it
returns a PyTorch tensor.
# Render a reference image (no derivatives used yet)
image_ref = render_torch(scene, spp=8)
crop_size = scene.sensors()[0].film().crop_size()
write_bitmap('out.png', image_ref, crop_size)
As before, we change one of the input parameters and initialize an optimizer.
# Change the left wall into a bright white surface
params['red.reflectance.value'] = [.9, .9, .9]
params.update()
# Which parameters should be exposed to the PyTorch optimizer?
params_torch = params.torch()
# Construct a PyTorch Adam optimizer that will adjust 'params_torch'
opt = torch.optim.Adam(params_torch.values(), lr=.2)
objective = torch.nn.MSELoss()
Note that the scene parameters are effectively duplicated: we represent them
once using Enoki arrays (params
), and once using PyTorch arrays
(params_torch
). To perform a differentiable rendering, the function
render_torch()
requires that both are given
as arguments. Due to technical specifics of how PyTorch detects differentiable
parameters, it is furthermore necessary that params_torch
is expanded into
a list of keyword arguments (**params_torch
). The function then keeps both
representation in sync and creates an interface between the underlying
computation graphs.
The main optimization loop looks as follows:
for it in range(100):
# Zero out gradients before each iteration
opt.zero_grad()
# Perform a differentiable rendering of the scene
image = render_torch(scene, params=params, unbiased=True,
spp=1, **params_torch)
write_bitmap('out_%03i.png' % it, image, crop_size)
# Objective: MSE between 'image' and 'image_ref'
ob_val = objective(image, image_ref)
# Back-propagate errors to input parameters
ob_val.backward()
# Optimizer: take a gradient step
opt.step()
# Compare iterate against ground-truth value
err_ref = objective(params_torch['red.reflectance.value'], param_ref)
print('Iteration %03i: error=%g' % (it, err_ref * 3))
Warning
Memory caching: When a GPU array in Enoki or PyTorch is destroyed, its memory is not immediately released back to the GPU. The reason for this is that allocating and releasing GPU memory are both extremely expensive operations, and any unused memory is therefore instead placed into a cache for later re-use.
The fact that this happens is normally irrelevant when only using Enoki or only using PyTorch, but it can be a problem when using both at the same time, as the cache of one system may grow sufficiently large that allocations by the other system fail, despite plenty of free memory technically being available.
If you notice that your programs crash with out-of-memory errors, try
passing malloc_trim=True
to the render_torch
function. This
flushes PyTorch’s memory cache before executing any Enoki code, and vice
versa. This is something of a last resort—generally, it’s better to
reduce memory requirements by lowering the number of samples per pixel,
as flushing the cache causes severe performance penalty.
Note
The full Python script of this tutorial can be found in the file:
docs/examples/10_diff_render/invert_cbox_torch.py
.