updates source code in nn example

Author: wjakob

docs/nn.rst

@@ -43,48 +43,86 @@ Example
 Below is a fully worked out example demonstrating how to use it to declare and
 optimize a small `multilayer perceptron
 <https://en.wikipedia.org/wiki/Multilayer_perceptron>`__ (MLP). This network
-implements a 2D neural field that we fit to an image.
+implements a 2D neural field (right) that we then fit to a low-resolution image of `The
+Great Wave off Kanagawa
+<https://en.wikipedia.org/wiki/The_Great_Wave_off_Kanagawa>`__ (left).
+
+.. image:: https://rgl.s3.eu-central-1.amazonaws.com/media/uploads/wjakob/2024/06/coopvec-screenshot.png
+  :width: 300
+  :align: center
 
 .. code-block:: python
 
+    from tqdm.auto import tqdm
+    import imageio.v3 as iio
     import drjit as dr
     import drjit.nn as nn
+    from drjit.opt import Adam, GradScaler
+    from drjit.auto.ad import Texture2f, TensorXf, TensorXf16, Float16, Float32, Array2f, Array3f
+
+    # Load a test image and construct a texture object
+    ref = TensorXf(iio.imread("https://rgl.s3.eu-central-1.amazonaws.com/media/uploads/wjakob/2024/06/wave-128.png") / 256)
+    tex = Texture2f(ref)
 
-    from drjit.llvm.ad import TensorXf16
-    from drjit.opt import Adam
+     # Ensure consistent results when re-running the following
+    dr.seed(0)
 
     # Establish the network structure
     net = nn.Sequential(
-        nn.Linear(-1, 32, bias=False),
-        nn.ReLU(),
-        nn.Linear(-1, -1),
-        nn.ReLU(),
+        nn.TriEncode(16, 0.2),
+        nn.Cast(Float16),
+        nn.Linear(-1, -1, bias=False),
+        nn.LeakyReLU(),
+        nn.Linear(-1, -1, bias=False),
+        nn.LeakyReLU(),
+        nn.Linear(-1, -1, bias=False),
+        nn.LeakyReLU(),
         nn.Linear(-1, 3, bias=False),
-        nn.Tanh()
+        nn.Exp()
     )
 
     # Instantiate the network for a specific backend + input size
     net = net.alloc(TensorXf16, 2)
 
-    # Pack coefficients into a training-optimal layout
+    # Convert to training-optimal layout
     coeffs, net = nn.pack(net, layout='training')
+    print(net)
 
-    # Optimize a float32 version of the packed coefficients
+    # Optimize a single precision copy of the parameters
     opt = Adam(lr=1e-3, params={'coeffs': Float32(coeffs)})
 
-    # Update network state from optimizer
-    for i in range(1000):
-        # Update neural network state
-        coeffs[:] = Float16(opt['coeffs'])
+    # This is an adaptive mixed-precision (AMP) optimization, where a half
+    # precision computation runs within a larger single precision program.
+    # Gradient scaling is required to make this numerically well-behaved.
+    scaler = GradScaler()
 
-        # Create input
-        out = net(nn.CoopVec(...))
+    res = 256
 
-        # Unpack
-        out = Array3f16(result)
-
-        # Backpropagate
-        dr.backward(dr.square(reference-out))
+    for i in tqdm(range(40000)):
+        # Update network state from optimizer
+        coeffs[:] = Float16(opt['coeffs'])
 
-        # Take a gradient step
-        opt.step()
+        # Generate jittered positions on [0, 1]^2
+        t = dr.arange(Float32, res)
+        p = (Array2f(dr.meshgrid(t, t)) + dr.rand(Array2f, (2, res*res))) / res
+
+        # Evaluate neural net + L2 loss
+        img = Array3f(net(nn.CoopVec(p)))
+        loss = dr.squared_norm(tex.eval(p)-img)
+
+        # Mixed-precision training: take suitably scaled steps
+        dr.backward(scaler.scale(loss))
+        scaler.step(opt)
+
+    # Done optimizing, now let's plot the result
+    t = dr.linspace(Float32, 0, 1, res)
+    p= Array2f(dr.meshgrid(t, t))
+    img = Array3f(net(nn.CoopVec(p)))
+    img = dr.reshape(TensorXf(img, flip_axes=True), (res, res, 3))
+
+    import matplotlib.pyplot as plt
+    fig, ax = plt.subplots(1, 2, figsize=(10,5))
+    ax[0].imshow(ref)
+    ax[1].imshow(dr.clip(img, 0, 1))
+    fig.tight_layout()
+    plt.show()