def check_finite_difference(test_name, make_scene, get_diff_params, diff_integrator=diff_integrator_default, diff_spp=4, diff_passes=10, ref_integrator=ref_integrator_default, ref_spp=128, ref_passes=10, ref_eps=0.001): from mitsuba.core import Bitmap, Struct from mitsuba.python.autodiff import render # Render groundtruth image and gradients (using finite difference) img_ref, grad_ref = compute_groundtruth(make_scene, ref_integrator, ref_spp, ref_passes, ref_eps) ek.cuda_malloc_trim() scene = make_scene(diff_integrator, diff_spp, 0.0) fsize = scene.sensors()[0].film().size() img, grad = render_gradient(scene, diff_passes, get_diff_params(scene)) error_img = np.abs(img_ref - img).mean() error_grad = np.abs(grad_ref - grad).mean() if error_img > 0.1: print("error_img:", error_img) Bitmap(img_ref).write('%s_img_ref.exr' % test_name) Bitmap(img).write('%s_img.exr' % test_name) assert False if error_grad > 0.1: print("error_grad:", error_grad) scale = np.abs(grad_ref).max() write_gradient_image(grad_ref / scale, '%s_grad_ref' % test_name) write_gradient_image(grad / scale, '%s_grad' % test_name) Bitmap(img_ref).write('%s_img_ref.exr' % test_name) Bitmap(img).write('%s_img.exr' % test_name) assert False
def backward(ctx, grad_out): ek.set_gradient(ctx.out, ek.FloatC(grad_out)) ek.FloatD.backward() result = (ek.gradient(ctx.in1).torch() if ek.requires_gradient(ctx.in1) else None, ek.gradient(ctx.in2).torch() if ek.requires_gradient(ctx.in2) else None) del ctx.out, ctx.in1, ctx.in2 ek.cuda_malloc_trim() return result
def forward(ctx, arg1, arg2): ctx.in1 = ek.FloatD(arg1) ctx.in2 = ek.FloatD(arg2) ek.set_requires_gradient(ctx.in1, arg1.requires_grad) ek.set_requires_gradient(ctx.in2, arg2.requires_grad) ctx.out = ek.atan2(ctx.in1, ctx.in2) out_torch = ctx.out.torch() ek.cuda_malloc_trim() return out_torch
def forward(ctx, scene, params, *args): try: assert len(args) % 2 == 0 args = dict(zip(args[0::2], args[1::2])) spp = None sensor_index = 0 unbiased = True ctx.inputs = [None, None] for k, v in args.items(): if k == 'spp': spp = v elif k == 'sensor_index': sensor_index = v elif k == 'unbiased': unbiased = v elif params is not None: params[k] = type(params[k])(v) ctx.inputs.append(None) ctx.inputs.append( params[k] if v.requires_grad else None) ctx.inputs.append(None) ctx.inputs.append(None) if type(spp) is not tuple: spp = (spp, spp) result = None if params is not None: params.update() if unbiased: result = render(scene, spp=spp[0], sensor_index=sensor_index).torch() for v in ctx.inputs: if v is not None: ek.set_requires_gradient(v) ctx.output = render(scene, spp=spp[1], sensor_index=sensor_index) if result is None: result = ctx.output.torch() ek.cuda_malloc_trim() return result except Exception as e: print("render_torch(): critical exception during " "forward pass: %s" % str(e)) raise e
def backward(ctx, grad_output): try: ek.set_gradient(ctx.output, ek.detach(Float(grad_output))) Float.backward() result = tuple(ek.gradient(i).torch() if i is not None else None for i in ctx.inputs) del ctx.output del ctx.inputs ek.cuda_malloc_trim() return result except Exception as e: print("render_torch(): critical exception during " "backward pass: %s" % str(e)) raise e