Unexpected rendering when using batch lights

🐛 Bugs / Unexpected behaviors

Hey!
I'm trying to render different meshes with different lights, and process them in batches. However, I found the results are not what I was expected.
To simplify the steps as much as possible, let's try a simple case:

`

ambient_color = torch.FloatTensor([[0.0, 0.0, 0.0]]).to(self.device)
diffuse_color = torch.FloatTensor([[1.0, 1.0, 1.0]]).to(self.device)
specular_color = torch.FloatTensor([[0.0, 0.0, 0.0]]).to(self.device)
direction = torch.FloatTensor([[1, 0, 0]]).to(self.device)
lights = DirectionalLights(ambient_color=ambient_color,
                           diffuse_color=diffuse_color,
                           specular_color=specular_color,
                           direction=direction,
                           device=self.device)
shader = TexturedSoftPhongShader(device=self.device, cameras=self.cameras, lights=lights)
renderer = MeshRenderer(rasterizer=self.rasterizer, shader=shader)
render_image = renderer(meshes) # there are 2 meshes to be rendered

`
In this case, the arguments of DirectionalLights are tensor of shape (1, 3), and the results seem OK.

However, if I change the shape of arguments of DirectionalLights to (2, 3), the results are wired:
`

ambient_color = torch.FloatTensor([[0.0, 0.0, 0.0]]).to(self.device).repeat(2, 1)
diffuse_color = torch.FloatTensor([[1.0, 1.0, 1.0]]).to(self.device).repeat(2, 1)
specular_color = torch.FloatTensor([[0.0, 0.0, 0.0]]).to(self.device).repeat(2, 1)
direction = torch.FloatTensor([[1, 0, 0]]).to(self.device).repeat(2, 1)
lights = DirectionalLights(ambient_color=ambient_color,
                           diffuse_color=diffuse_color,
                           specular_color=specular_color,
                           direction=direction,
                           device=self.device)

`

Theoretically, the results should be the same as before, so I looked into the code and found the problem may be caused by broadcasting issue in

pytorch3d/pytorch3d/renderer/mesh/shading.py

Line 82 in 6a365d2

colors = (ambient + diffuse) * texels + specular

In my case, the shape of ambient is (2, 3) and diffuse is (2, 256, 256, 1, 3).
And I believe this line could be fixed as:
colors = (ambient[:, None, None, None] + diffuse) * texels + specular
It produces the right results for me, but I'm not sure if this is the best way.

P.S. I also found similar broadcasting issues in:

pytorch3d/pytorch3d/renderer/lighting.py

Line 251 in 6a365d2

direction = self.location - points

pytorch3d/pytorch3d/renderer/lighting.py

Line 255 in 6a365d2

direction = self.location - points

and I think those should be:
direction = self.location[:, None, None, None] - points

I wish to add to the above that i have had a problem with the lighting too, with shapes not broadcasting correctly with each other and failing with

  File "[...]\lib\site-packages\pytorch3d-0.2.0-py3.7-win-amd64.egg\pytorch3d\renderer\lighting.py", line 67, in diffuse
    angle = F.relu(torch.sum(normals * direction, dim=-1))
RuntimeError: The size of tensor a (4) must match the size of tensor b (16777216) at non-singleton dimension 0

having a batch size of 4 in meshes, lights, and camera T+R.

however with the fixes above (or by removing the explicit batch size from the PointLights), the forward function didn't fail anymore

@Edward-Lin, @michele-arrival thanks for pointing this out. We will look into this issue - feel free to submit a PR with a fix.

This issue should now be fixed by bc8361f.