Waabi OA Interview Problem: Speed Up NeRF (PyTorch, NeRF, Voxel Feature Volume)

Speed Up NeRF

Here we give you a vanilla NeRF implementation (~100 lines of code).

Your task: Make it faster! Goal is PSNR > 20 in 5 mins (original takes 2 hours).

You’ll need to understand the code, implement a PyTorch model, debug, and potentially do some parameter tuning.

Instructions:

• Briefly go over the code, take care of the forward method in NerfModel.
• Reimplement it in the forward method in FastNerfModel.

Some extra info:

• Volume AABB is -1.5 to 1.5.
• Suggested approach: use volume representation (CxDxHxW) plus a small MLP.
• You are welcome to explore Hashtable, Triplane, or other ideas, but they might be trickier to do in an hour.

This task asks you to accelerate a vanilla NeRF implementation by understanding the original forward pass and then replacing it with a faster model in PyTorch. The suggested direction is a voxel feature volume with a small MLP, which maps 3D world coordinates and viewing directions to density and color more efficiently than a fully shared NeRF MLP. The key challenges are implementing the forward logic correctly, keeping the model numerically stable, and tuning it enough to reach the PSNR target within the time limit.