Antialiasing memory leak?
odusseys opened this issue · comments
I am using the following code snippet to render a scene:
def render_untextured(mesh, cam=None, light_dir=None):
if cam is None or light_dir is None:
cam, light_dir = random_camera_and_light()
# Rasterize the mesh.
vertices_camera = cam.extrinsics.transform(mesh.vertices)
proj = cam.projection_matrix()[None]
homogeneous_vecs = kal.render.camera.up_to_homogeneous(
vertices_camera
)[..., None]
vertices_clip = (proj @ homogeneous_vecs).squeeze(-1).contiguous()
faces = mesh.faces.int().contiguous()
rast_out = nvdiffrast.torch.rasterize(
glctx, vertices_clip, faces, (cam.height, cam.width)
)[0]
# Compute rasterized normals and apply material with light direction
normals = mesh.vertex_normals
vertex_shade = torch.sum(normals * light_dir, dim=2, keepdim=True)
mask=torch.full((1, normals.shape[1], 1), 1).to("cuda")
features = torch.cat((vertex_shade, mask), dim=2)
rendered = nvdiffrast.torch.interpolate(features, rast_out, faces)[0]
shade = rendered[:,:,:, :1]
hard_mask = rendered[:,:,:, 1:]
res = hard_mask * material_color * shade + (1 - hard_mask) * background_color
res = torch.clamp(res, 0, 1)
# REMOVING THIS CAUSES MEMORY LEAK
# res = nvdiffrast.torch.antialias(res.contiguous(), rast_out, vertices_clip, faces, pos_gradient_boost=3)
return torch.squeeze(res).to("cuda")
Here, cam
is a Kaolin camera, but you can assume that everything from proj=...
is standard
If I remove the commented line which performs the antialiasing step, I get an instant memory leak. After rendering a couple images, CUDA memory is full, and resetting its cache / performing gc only clears a fraction of the used memory. Removing the antialiasing leads to no noticeable GPU memory usage whatsoever (T4 card).
For reference - I am using
- Google Colab
- CUDA 118
- pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
- nvdiffrast@latest
- using RasterizeCudaContext!
I don't see a reason why this would leak memory, so there may be something going on outside this function. The antialias op stores a couple of temporary objects for the gradient pass, and if you accidentally hold onto these across iterations, that might explain increasing memory consumption — see the first item here.
However, these objects shouldn't be big enough to immediately fill the GPU memory, unless your mesh has a massive number of triangles (tens or hundreds of millions).
The function above looks fine to me, so hunting this down would require a stand-along reproducer. If this is a bug in the antialiasing op itself, it has to be some sort of a corner case, or the inputs have to be extreme in some way that hasn't come up in other projects that have used this op.