Gradient based optimization of scene parameters (Mesh, BSDF, environment map)

[stefanjp]

Hi,

I tried to optimize alpha_u/alpha_v of the roughtconductor BRDF, the environment map illumination and the mesh (separately). I did not manage to optimize the parameters with the Adam optimizer from this repository (utils/adam.py). I did not find an example on how to use the shipped adam optimizer, so I might be using it in a non intended way. After the first iteration, the gradients stay zero. Here is my code, executed from the examples directory and with a scene with reduced sample_count:

### Fail gradient based optimization ###
import enoki as ek
from enoki.cuda_autodiff import Float32 as FloatD, Vector3f as Vector3fD, Matrix4f as Matrix4fD
import psdr_cuda
from utils.adam import Adam as AdamPsdr
import matplotlib.pyplot as plt

scene = psdr_cuda.Scene()

# bunny_env.xml reduced independent sampler sample_count to 10 to avoid out of memory exception
scene.load_file('data/scenes/bunny_env.xml')

# render target image

# Initialize an integrator
integrator = psdr_cuda.DirectIntegrator()

# Start rendering!
sensor_index = 0
target_image = integrator.renderD(scene, sensor_index)

# rotate mesh
P = FloatD(1.)
# Construct a rotation matrix
mat = Matrix4fD.rotate(Vector3fD(0., 0., 1.), P)

# Left-multiply this matrix to the to-world transformation of the first mesh of the scene
scene.param_map["Mesh[0]"].append_transform(mat)

lr=.2
opt = AdamPsdr(None, scene.param_map, [], ["Mesh[0]"], lr=lr)

scene.configure()

image_res = [scene.opts.width, scene.opts.height]

image = integrator.renderD(scene, sensor_index)

target_image_vis = target_image.numpy().reshape(image_res[1],image_res[0],-1)
image_vis = image.numpy().reshape(image_res[1],image_res[0],-1)

plt.imshow(target_image_vis)
plt.show()
plt.imshow(image_vis)
plt.show()

for it in range(20):
    # calc loss
    loss = ek.sqrt(ek.hmean(ek.squared_norm(target_image - image)))
    print(loss)

    # Reverse-mode autodiff
    ek.backward(loss)

    grad = ek.gradient(scene.param_map["Mesh[0]"].vertex_positions)
    print(ek.hmax(ek.hmax(grad)))

    opt.step()

plt.imshow(image_vis)
plt.show()

[andyyankai]

We don't usually use that adam optimizer. However, enoki have pytorch binding and you can use pytorch optimizor like torch.optim.Adam to optimize the scene param_map. You can write your torch forward() and backward() method for rendering and backward. something like this:

    class PSDRRender(torch.autograd.Function):
        @staticmethod 
        def forward(ctx, V):
            _vertex_pos = Vector3fD(V)
            ek.set_requires_gradient(_vertex_pos, V.requires_grad)
            ctx.in1 = _vertex_pos
            sc.param_map['Mesh[0]'].vertex_positions = _vertex_pos;

            sc.configure()
            integrator = psdr_cuda.DirectIntegrator(1, 1)
            loss = 0
            for sensor_id in range(num_sensors):
                img = integrator.renderD(sc, sensor_id)
                tar = Vector3fD(tars[sensor_id].cuda())

                tmp_loss = ek.abs(img - tar)
                loss += ek.hsum(ek.hsum(tmp_loss)) / (ro.height * ro.width * 3) / num_sensors
            ctx.out = loss
            out_torch = ctx.out.torch() 
            return out_torch

        @staticmethod
        def backward(ctx, grad_out):
            ek.set_gradient(ctx.out, FloatC(grad_out))
            FloatD.backward() 
            gV = ek.gradient(ctx.in1)

            nan_mask = ek.isnan(gV)
            gV = ek.select(nan_mask, 0, gV)

            gradV = gV.torch()            
            result = (gradV, None)
            del ctx.out, ctx.in1
            return result

[andyyankai]

You can optimize this method just like other common pytorch example using for example torch.optim.adam

[andyyankai]

Also, we doesn't fully support anisotropic for now(although you might see the code do have these terms). You need to double check in implementation by yourself to separate the alpha into alpha_u and alpha_v or optimize alpha only.

[stefanjp]

Thank you for providing an example code for pytorch integration, that was a big help! I have tried to optimize with pytorch before, but did not manage to implement the integration correctly (did not update the scene properly).

I managed to implement gradient based optimization for the environment map and the vertices. The optimization does not converge at the moment. Not sure yet if it is due to the problem statement or a coding error.

[stefanjp]

Everything is working now. Using not enough SPP in the reference image was the issue for my problem with convergence. If anyone needs example code for optimizing the environment map, I can post it here!

[l1346792580123]

Everything is working now. Using not enough SPP in the reference image was the issue for my problem with convergence. If anyone needs example code for optimizing the environment map, I can post it here!

Sorry to bother you. Could you tell me what is the spp, sppe and sppse you use? I also encounterd optimization problems. Also, have you encountered noise in the generatered image?

[stefanjp]

Everything is working now. Using not enough SPP in the reference image was the issue for my problem with convergence. If anyone needs example code for optimizing the environment map, I can post it here!

Sorry to bother you. Could you tell me what is the spp, sppe and sppse you use? I also encounterd optimization problems. Also, have you encountered noise in the generatered image?

In the end I've only done benchmark experiments optimizing the illumination of a scene. I used 10 spp. My goal was to replace mitsuba 2 for performance reasons. In the benchmark I did not care about the result of the optimization. A problem I could not resolve was the huge memory consumption of my PSDR installation (I gave up investigating the cause of the problem).

My benchmark code (data from the psdr cuda repository, bunny_env.xml with changed sample_count of 10):

### PSDR-CUDA Benchmark + Pytorch ###
import enoki as ek
from enoki.cuda_autodiff import Float32 as FloatD, Vector3f as Vector3fD, Matrix4f as Matrix4fD
from enoki.cuda import Float32 as FloatC, Vector3f as Vector3fC
import psdr_cuda
import matplotlib.pyplot as plt
import torch
import time

scene = psdr_cuda.Scene()

# bunny_env.xml reduced independent sampler sample_count to 10 to avoid out of memory exception
scene.load_file('data/scenes/bunny_env.xml')
ro = scene.opts

# render target image

# Initialize an integrator
integrator = psdr_cuda.DirectIntegrator(1, 1)

sensor_index = 0
target_image_torch = integrator.renderD(scene, sensor_index).torch().detach()
target_image = Vector3fC(target_image_torch)

class PSDRRender(torch.autograd.Function):
    @staticmethod 
    def forward(ctx, intorch):
        inenoki = Vector3fD(intorch)
        ek.set_requires_gradient(inenoki, intorch.requires_grad)
        ctx.in1 = inenoki
        scene.param_map["Emitter[0]"].radiance.data = inenoki
        scene.configure()
        ctx.img = integrator.renderD(scene, sensor_index)
        return ctx.img.torch() 

    @staticmethod
    def backward(ctx, grad_out):
        ek.set_gradient(ctx.img, ek.detach(Vector3fD(grad_out)))
        FloatD.backward() 
        gV = ek.gradient(ctx.in1)

        nan_mask = ek.isnan(gV)
        gV = ek.select(nan_mask, 0, gV)

        gradV = gV.torch()            
        result = gradV
        del ctx.in1, ctx.img
        return result

# reset illumination that well be optimized in this test
emitter = scene.param_map["Emitter[0]"]
emitter_resolution = emitter.radiance.resolution
emitter_data_gt = Vector3fD(emitter.radiance.data)
emitter_data = ek.full(Vector3fD, .5, len(emitter_data_gt.x))
scene.param_map["Emitter[0]"].radiance.data = emitter_data
ek.set_requires_gradient(scene.param_map["Emitter[0]"].radiance.data)
scene.configure() # TODO just in case, might be not required

opt_params = [scene.param_map["Emitter[0]"].radiance.data.torch()]
opt_params[0].requires_grad = True
lr=1e-2
opt = torch.optim.Adam([
    {      
        'params': opt_params
        #, 'lr': 1e-2
    }
], lr=lr)

start = time.time()
elapsed = []
rendertime = []

pltiter = 200
objective = torch.nn.MSELoss()
for it in range(100):
    # Start rendering!
    renderstart = time.time()
    image = PSDRRender.apply(opt_params[0])
    rendertime.append(time.time() - renderstart)

    loss = objective(image, target_image_torch)
    # Reverse-mode autodiff
    loss.backward()

    opt.step()

    # Compare iterate against ground-truth value
    err_ref = ek.hsum(ek.sqr(emitter_data_gt - Vector3fD(opt_params[0].detach())))

    print(f'Iteration {it:=3d}: error={ torch.mean(err_ref.torch())}, loss= {loss.item()}')
    elapsed.append(time.time() - start)
    if it % pltiter == 0:
        image_vis = image.detach().cpu().numpy().reshape(ro.height,ro.width,-1)
        plt.imshow(image_vis)
        plt.show()

end = time.time()

print(f'elapsed {end - start} seconds for {it + 1} iterations.')
plt.figure(figsize=(8, 10), dpi=160)
plt.suptitle('PSDR-CUDA benchmark optimize envmap', fontsize="x-large")
plotindex = 1
plt.subplot(2,2,plotindex)
plt.imshow(emitter_data_gt.numpy().reshape(emitter_resolution[1], emitter_resolution[0], -1))
plt.title(f'gt envmap')

plotindex += 1
plt.subplot(2,2,plotindex)
plt.imshow(scene.param_map["Emitter[0]"].radiance.data.numpy().reshape(emitter_resolution[1], emitter_resolution[0], -1))
plt.title(f'result envmap after {it + 1} iterations')

plotindex += 1
plt.subplot(2,2,plotindex)
plt.plot(elapsed)
plt.xlabel('iterations')
plt.ylabel('time [s]')
plt.title(f'time per iteration accumulated')

plotindex += 1
plt.subplot(2,2,plotindex)
plt.plot(rendertime)
plt.xlabel('iteration')
plt.ylabel('time [s]')
plt.title(f'render only per iteration')

plt.tight_layout()
plt.savefig("PSDR-CUDA-bench.png", facecolor='w')
plt.show()

With following output: