{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Gradient-based optimization" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "## Overview\n", "\n", "Mitsuba 3 can be used to solve inverse problems involving light using a technique known as *differentiable rendering*. It interprets the rendering algorithm as a function $f(\\mathbf{x})$ that converts an input $\\mathbf{x}$ (the scene description) into an output $\\mathbf{y}$ (the rendering). This function $f$ is then mathematically differentiated to obtain $\\frac{d\\mathbf{y}}{d\\mathbf{x}}$, providing a first-order approximation of how a desired change in the output $\\mathbf{y}$ (the rendering) can be achieved by changing the inputs $\\mathbf{x}$ (the scene description). Together with a differentiable *objective function* $g(\\mathbf{y})$ that quantifies the suitability of tentative scene parameters, a gradient-based optimization algorithm such as stochastic gradient descent or Adam can then be used to find a sequence of scene parameters $\\mathbf{x_0}$, $\\mathbf{x_1}$, $\\mathbf{x_2}$, etc., that successively improve the objective function. In pictures:\n", "\n", "\n", "\n", "\n", "In this tutorial, we will build a simple example application that showcases differentiation and optimization through a light transport simulation:\n", "\n", "1. We will first render a reference image of the Cornell Box scene.\n", "2. Then, we will perturb the color of one of the walls, e.g. changing it to blue.\n", "3. Finally, we will try to recover the original color of the wall using differentiation along with the reference image generated in step 1.\n", "\n", "Mitsuba’s ability to automatically differentiate entire rendering algorithms builds on differentiable JIT array types provided by the Dr.Jit library. Those are explained in the [Dr.Jit documentation][1]. The linked document also discusses key differences compared to related frameworks like PyTorch and TensorFlow. For *automatic differentiation* (AD), Dr.Jit records and simplifies computation graphs and uses them to propagate derivatives in forward or reverse mode. Before getting further into this tutorial, we recommend that you familiarize yourself Dr.Jit.\n", "\n", "
Optimizer
classesprb
][2]) introduced by Vicini et al. (2021). It is essentially a path tracer, augmented with a specialized algorithm to efficiently compute the gradients in a separate adjoint pass.\n",
"\n",
"[1]: https://mitsuba.readthedocs.io/en/latest/src/key_topics/scene_format.html\n",
"[2]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.integrators.prb"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"ExecuteTime": {
"end_time": "2021-06-15T14:13:40.914750Z",
"start_time": "2021-06-15T14:13:40.807334Z"
}
},
"outputs": [],
"source": [
"scene = mi.load_file('../scenes/cbox.xml', res=128, integrator='prb')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Reference image\n",
"\n",
"We render a reference image of the original scene that will later be used in the objective function for the optimization. Ideally, this reference image should expose very little noise as it will pertube optimization process otherwise. For best results, we should render it with an even larger sample count."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"SGD
][1]) with and without momentum, as well as [Adam
][2] [KB14]. We will instantiate the latter and optimize our scene parameter with a learning rate of `0.05`. \n",
"\n",
"We then set the color to optimize on the optimizer, which will now hold a copy of this parameter and enable gradient tracking on it. During the optimization process, the optimizer will always perfom gradient steps on those variables. To propagate those changes to the scene, we need to call the `update()` method which will copy the values back into the `params` data structure. As always this method also notifies all objects in the scene whose parameters have changed, in case they need to update their internal state.\n",
"\n",
"This first call to `params.update()` ensures that gradient tracking with respect to our wall color parameter is propagated to the scene internal state. For more detailed explanation on how-to-use the optimizer classes, please refer to the dedicated [how-to-guide][3].\n",
"\n",
"[1]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.SGD\n",
"[2]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.Adam\n",
"[3]: https://mitsuba.readthedocs.io/en/latest/src/how_to_guides/use_optimizers.html"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"ExecuteTime": {
"end_time": "2021-06-15T14:13:51.368393Z",
"start_time": "2021-06-15T14:13:51.362695Z"
}
},
"outputs": [],
"source": [
"opt = mi.ad.Adam(lr=0.05)\n",
"opt[key] = params[key]\n",
"params.update(opt);"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"At every iteration of the gradient descent, we will compute the derivatives of the scene parameters with respect to the objective function. In this simple experiment, we use the [*mean square error*][1], or $L_2$ error, between the current image and the reference created above.\n",
"\n",
"[1]: https://en.wikipedia.org/wiki/Mean_squared_error"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"def mse(image):\n",
" return dr.mean(dr.sqr(image - image_ref))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the following cell we define the hyper parameters controlling our optimization loop, such as the number of iterations:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"iteration_count = 50"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"nbsphinx": "hidden",
"tags": []
},
"outputs": [],
"source": [
"# IGNORE THIS: When running under pytest, adjust parameters to reduce computation time\n",
"import os\n",
"if 'PYTEST_CURRENT_TEST' in os.environ:\n",
" iteration_count = 2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"It is now time to actually perform the gradient-descent loop that executes 50 differentiable rendering iterations."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"ExecuteTime": {
"end_time": "2021-06-15T14:14:04.664564Z",
"start_time": "2021-06-15T14:13:54.209445Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Iteration 49: parameter error = 0.001432\n",
"Optimization complete.\n"
]
}
],
"source": [
"errors = []\n",
"for it in range(iteration_count):\n",
" # Perform a (noisy) differentiable rendering of the scene\n",
" image = mi.render(scene, params, spp=4)\n",
" \n",
" # Evaluate the objective function from the current rendered image\n",
" loss = mse(image)\n",
"\n",
" # Backpropagate through the rendering process\n",
" dr.backward(loss)\n",
"\n",
" # Optimizer: take a gradient descent step\n",
" opt.step()\n",
"\n",
" # Post-process the optimized parameters to ensure legal color values.\n",
" opt[key] = dr.clamp(opt[key], 0.0, 1.0)\n",
"\n",
" # Update the scene state to the new optimized values\n",
" params.update(opt)\n",
" \n",
" # Track the difference between the current color and the true value\n",
" err_ref = dr.sum(dr.sqr(param_ref - params[key]))\n",
" print(f\"Iteration {it:02d}: parameter error = {err_ref[0]:6f}\", end='\\r')\n",
" errors.append(err_ref)\n",
"print('\\nOptimization complete.')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Results\n",
"\n",
"We can now render the scene again to check whether the optimization process successfully recovered the color of the red wall."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"ExecuteTime": {
"end_time": "2021-06-15T14:14:07.765382Z",
"start_time": "2021-06-15T14:14:06.895491Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"Optimizer
](https://mitsuba.readthedocs.io/en/latest/src/how_to_guides/use_optimizers.html)\n",
"- API reference:\n",
" - [mitsuba.ad.Optimizer
](https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.Optimizer)\n",
" - [prb
plugin](https://mitsuba.readthedocs.io/en/latest/src/generated/plugins_integrators.html#path-replay-backpropagation-prb)\n",
" - [mitsuba.ad.SGD
](https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.SGD)\n",
" - [mitsuba.ad.Adam
](https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.ad.Adam)\n",
" - [drjit.backward
](https://drjit.readthedocs.io/en/latest/reference.html#drjit.backward)"
]
}
],
"metadata": {
"file_extension": ".py",
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.12"
},
"metadata": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
},
"mimetype": "text/x-python",
"name": "python",
"npconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": 3
},
"nbformat": 4,
"nbformat_minor": 4
}