Pytorch3d: Batch lighting position is not working for point light

[Jiayi-Wang-MPI]

The diffuse and specular function of point light gives a runtime error about mismatched tensor size when the input locations is batched. This is due to tensor broadcasting. Is batch light location supported?

[nikhilaravi]

@Jiayi-Wang-MPI we support one light per mesh in the batch. Can you send a code snippet of what you are trying to do so I can reproduce the error?

[Jiayi-Wang-MPI]

@nikhilaravi Using the teapot example as a starting point and duplicating the the inputs:

verts, faces_idx, _ = load_obj("./data/teapot.obj")
faces = faces_idx.verts_idx

verts = torch.stack(4*[verts])
faces = torch.stack(4*[faces])

verts_rgb = torch.ones_like(verts) 

location= (torch.rand(4,3)-0.5)*2
location[:,2] = -torch.abs(location[:,2])

distance = 3   
elevation = 40.0
azimuth = 10.0

R, T = look_at_view_transform(distance, elevation, azimuth, device=device)
R = torch.cat(4*[R])
T = torch.cat(4*[T])

textures = Textures(verts_rgb=verts_rgb.to(device))

teapot_mesh = Meshes(
    verts=verts.to(device),   
    faces=faces.to(device), 
    textures=textures
)

cameras = OpenGLPerspectiveCameras(device=device)

raster_settings = RasterizationSettings(
    image_size=256, 
    blur_radius=0.0, 
    faces_per_pixel=1, 
    bin_size=0
)

#input variables
lights = PointLights(device=device, location=location)
phong_renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras, 
        raster_settings=raster_settings
    ),
    shader=PhongShader(device=device, lights=lights)
)

image_ref = phong_renderer(meshes_world=teapot_mesh, R=R, T=T)

This causes the following error:

RuntimeError: The size of tensor a (4) must match the size of tensor b (1048576) at non-singleton dimension 0

This is due to the direction calculation for point light in lighting.py

direction = self.location - points

If location is provided as a (N,3) tensor, and since points is a (N,H,W,K,3) tensor, broadcasting produces the incorrect shape of (N,H,W,N,3) when K is 1. When K is larger, the subtraction itself gives an error.

[nikhilaravi]

@Jiayi-Wang-MPI thanks for the detailed explanation. I think this could be fixed by expanding the dimensions of location (as done in several other places in lighting.py) e.g.

# Reshape direction so it has all the arbitrary intermediate
# dimensions as points. Assume first dim = batch dim and last dim = 3.
points_dims = points.shape[1:-1]
expand_dims = (-1,) + (1,) * len(points_dims) + (3,)
if self.location.shape != points.shape:
        direction =  self.location.view(expand_dims) - points
else:
       direction = self.location - points

Please try the above fix on your local clone in place of direction = self.location - points. I will try to replicate the error based on the code you have provided above and fix this issue in master.

[Jiayi-Wang-MPI]

@nikhilaravi I tried the suggested fix and found that a similar bug is present in shading.py.

In line 82 and 127, dimensions of diffuse and ambient does not match and the dimensions of ambient need to be similarly expanded.

ambient, diffuse, specular = _apply_lighting(
        pixel_coords, pixel_normals, lights, cameras, materials
    )
#Note ambient_color has the same shape as materials.ambient_color.
#but diffuse_color has the same shape as the input points
colors = (ambient + diffuse) * texels + specular

[nikhilaravi]

@Jiayi-Wang-MPI ok thanks for letting me know. I will add a fix for this with more tests to check the cases that you mentioned.

[HectorAnadon]

Is this bug being addressed?

[nikhilaravi]

@HectorAnadon thanks for following up on this and apologies for the delay - we will fix this issue soon!

[HectorAnadon]

I've solved the problem mentioned by @Jiayi-Wang-MPI in shading.py like this:
colors = (ambient.unsqueeze(1).unsqueeze(1).unsqueeze(1) + diffuse) * texels + specular

Not a nice fix though.

[hectorbasevi]

I have also run into this problem. It exists in several places including lighting.py (camera location), blending.py (background), and shading.py (material diffuse_color, material specular_color, ambient). The solution in all cases is to reshape the relevant tensors to add singleton dimensions. I would submit a PR but getting approval for the Facebook CLA will take some time.

Also, this appears to be the same bug as #192.

@HectorAnadon A slightly more general solution is:
colors = (ambient.view([ambient.size(0)]+[1]*(diffuse.dim()-2)+[ambient.size(1)]) + diffuse) * texels + specular
However, this assumes that packed mode isn't being used and I'm not familiar enough with the logic to know whether this assumption is appropriate everywhere.

[HectorAnadon]

@hectorbasevi Do you mean that I missed other "fixes" in different files? Could you submit a MR or a fork so I can see what cases I've missed while Facebook CLA fix these bugs? That would be super helpful, thanks!

[hectorbasevi]

@HectorAnadon You haven't necessarily missed fixes in other files, depending on which parameters you're interested in. According to the Pytorch3D API, it's possible to specify a variety of parameters individually (within a batch tensor) for each data point in a batch, including light location, light colors, material colours, and also the background color to use as a default. Because the coding pattern for the application of each of these parameters is the same the same bugs exist, in different parts of the code for each parameter. I can't guarantee that I've found every case because I haven't been looking.

I mentioned packed mode because there's a view of the mesh data which removes the batch dimension, and would require a different solution involving repeating the the tensor instead of just reshaping it. I don't know whether that's ever used in the current logic. Addressing every possible set of inputs would require an if statement and at least three branches (unbatched, batched and non-packed, batched and packed).

Sorry, I'll need to get permission to put code up publicly, whether as a MR/PR or a fork. Once I get approval I'll put code up and let you know.