Running out of memory when running differentiable loop with PyTorch

[adithyaknarayan]

I wanted to run a differentiable loop to update the camera pose inside a scene. However, I keep running into the following issue.

For reference, the current settings are

• spp : 256
• Scene: The Bedroom scene under 'Architecture' here: https://mitsuba.readthedocs.io/en/stable/src/gallery.html
• CUDA Version: 12.6
• NVIDIA RTX A4500 20GB

I've also included the code:

class CameraModel(nn.Module):
    def __init__(self, device):
        super().__init__()
        self.device = device
        # define optimizable parameters
        self.translation_param = nn.Parameter(
            torch.from_numpy(np.array(model_config['init_position'], dtype=np.float32)).to(self.device))

        self.rotation_param = nn.Parameter(torch.from_numpy(np.array(model_config['init_rotation_rad'], dtype=np.float32)).to(self.device))
    
    def get_camera_extrinsics(self, R, T):
        # Combine R and T into a 3x4 matrix
        camera_matrix_3x4 = torch.cat([R, T], dim=1)
        
        # Add the bottom row [0, 0, 0, 1] to make it 4x4
        bottom_row = torch.tensor([[0.0, 0.0, 0.0, 1.0]]).to(self.device)
        camera_matrix_4x4 = torch.cat([camera_matrix_3x4, bottom_row], dim=0)
        return camera_matrix_4x4

    def forward(self):
        R = axis_angle_to_matrix(self.rotation_param)
        T = self.translation_param.unsqueeze(1)
        # assemble
        camera_extrinsics = self.get_camera_extrinsics(R, T)
        return camera_extrinsics
    
mi.set_variant("cuda_ad_rgb")
scene = mi.load_file('/bedroom/scene.xml', integrator='aov')
params = mi.traverse(scene)
key = 'PerspectiveCamera.to_world'
dr.enable_grad(params[key])
def denoise(image):
    denoiser = mi.OptixDenoiser(input_size=image.shape[:2], albedo=False, normals=False, temporal=False)
    denoised = denoiser(image[:,:,:3])
    return denoised

@dr.wrap(source='torch', target='drjit')
def render_texture(camera_extrinsics, spp=8, seed=1):
    params[key] = camera_extrinsics
    params.update()
    params.keep(key)
    return mi.render(scene, params=params, spp=spp, seed=seed, seed_grad=seed+1)

if __name__ == '__main__':
    cmodel = CameraModel('cuda')
    cmodel.train(mode=True)
    optimizer = torch.optim.Adam(cmodel.parameters(), lr=0.02)
    num_iter = 10
    loss_fn = nn.MSELoss()
    for i in range(num_iter):
        optimizer.zero_grad()
        camera_extrinsics = cmodel.forward()
        print(camera_extrinsics)
        rendered_image = render_texture(camera_extrinsics)
        image = rendered_image[:,:,:3]
        loss = loss_fn(image, image*0) # some dummy code as a placeholder
        print(loss)
        loss.backward()
        optimizer.step()
        torch.cuda.empty_cache()

Is there a reason this could be happening? So far I've tried

1. Lowering the spp - This only works if I use something like 5
2. Changing the resolution of the image - I'm currently using 256 as the output render size, I've lowered it but nothing changes and I still get the same error.