facebookresearch / pytorch3d

PyTorch3D is FAIR's library of reusable components for deep learning with 3D data
https://pytorch3d.org/
Other
8.53k stars 1.28k forks source link

Difference between SoftRas and Pytorch3d textured output #19

Closed aluo-x closed 4 years ago

aluo-x commented 4 years ago

To preface, it is very likely that I am doing something wrong, as I didn't closely look at the blending code in Pytorch3d. Would appreciate any advice on this particular topic.

Was comparing the silhouette and color output between SoftRas and Pytorch3d. The silhouette is very very similar, with a few pixels difference, perhaps due to some kind of rounding or offset issue.

But the textured (color) output is very different with black artifacts.

Links to mask & color output: Imgur link

Minimum replication code for SoftRas:

import matplotlib.pyplot as plt
import os
import numpy as np
import imageio
import soft_renderer as sr
current_dir = "/home/aluo/tools/SoftRas/examples"
data_dir = os.path.join(current_dir, '../data')
import matplotlib.pyplot as plt
class args_obj():
    filename_input = os.path.join(data_dir, 'obj/spot/spot_triangulated.obj')
    output_dir = os.path.join(data_dir, 'results/output_render')
def main():
    args = args_obj()
    mesh = sr.Mesh.from_obj(args.filename_input, load_texture=True, texture_res=5, 
texture_type='surface')
    renderer = sr.SoftRenderer(camera_mode='look_at')
    os.makedirs(args.output_dir, exist_ok=True)
    renderer.transform.set_eyes_from_angles(2.7, 10, 20)
    mesh.reset_()
    gamma_pow = -3
    renderer.set_gamma(10**gamma_pow)
    renderer.set_sigma(10**(gamma_pow-1))
    images = renderer.render_mesh(mesh)
    image = images.detach().cpu().numpy()[0].transpose((1, 2, 0))
    plt.figure(figsize=(10, 10))
    plt.grid("off")
    plt.axis("off")
    plt.imshow(image[:,:,:3])
    plt.show()
    plt.figure(figsize=(10, 10))
    plt.grid("off")
    plt.axis("off")
    plt.imshow(image[:,:,-1])
    plt.show()
main()

For Pytorch3d (following the render_textured_meshes notebook)

blend_params = BlendParams(sigma=1e-4, gamma=1e-4)
raster_settings = RasterizationSettings(
    image_size=256, 
    blur_radius=np.log(1. / 1e-4 - 1.) * blend_params.sigma, 
    faces_per_pixel=100, 
    bin_size=0
)
renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras, 
        raster_settings=raster_settings
    ),
    shader=TexturedPhongShader(
        device=device, 
        cameras=cameras,
        lights=lights
    )
)

images = renderer(mesh)
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.grid("off")
plt.axis("off")

plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., -1].cpu().numpy())
plt.grid("off")
plt.axis("off")
nikhilaravi commented 4 years ago

@aluo-x I can try to replicate this and get back to you tomorrow. I think it is due to some parameter setting in the blending. Can you try the PyTorch3d version again with the same gamma and sigma values as for SoftRas (it seems they are currently not the same)?

In the meantime you could try to copy this function and play around with the settings and then define a new shader which uses it (e.g. copy the TexturedPhongShader but replace with your modified blending function)

aluo-x commented 4 years ago

Playing around with gamma & sigma doesn't seem to cause the artifacts to disappear in the color image. The artifacts occur only if blur radius is set to greater than 0, even for very small values like 1e-5. Using hard/soft rgb blending doesn't seem to affect the artifacts either.

If this is a bug, then probs somewhere in the raster code?

nikhilaravi commented 4 years ago

@aluo-x There are a couple of things you need to change. First, pass in the blend_params to the TexturedPhongShader as currently the default values are being used so any changes you make will not be used:

    shader=TexturedPhongShader(
        blend_params=blend_params,
        device=device, 
        cameras=cameras,
        lights=lights
    )

Second, you need to enable barycentric clipping to [0, 1] before texture interpolation. I will add this in a pull request soon but in the meantime you can enable this by changing this line to

pixel_uvs = interpolate_face_attributes(fragments, faces_verts_uvs, bary_clip=True)

There are a few more things I am debugging to resolve the artifacts. I will get back to you shortly!

aluo-x commented 4 years ago

Apologies, made an error when copy pasting the code. My full notebook had too much unrelated stuff to fully paste here.

Minor question, by clamping the value to [0,1] wouldn't we be stopping the gradient at the border? Unsure if this would matter, but would something like:

clipped[clipped>1.0] = torch.max(clipped[clipped<1.0])

be necessary? Or do we not need the gradients to be calculated at the edge? Much appreciated for the help debugging this.

nikhilaravi commented 4 years ago

@aluo-x we have a function which does the barycentric clipping see here. To enable this you just need to set bary_clip=True at the line mentioned above. This uses the torch.clamp function which is differentiable.

In addition to the two steps mentioned above, one last step to resolving the issue of the black artifacts is to interpolate the z coordinate after barycentric clipping. To do this, pass in meshes to the softmax_rgb_blend function here i.e.

images = softmax_rgb_blend(colors, fragments, self.blend_params, meshes)

Then modify the softmax_rgb_blend function to take meshes as an input argument and then replace line 167 with:

# Reinterpolate the z values using clipped barycentrics
verts_z = meshes_world.verts_packed()
faces = meshes_world.faces_packed()
faces_verts_z = verts_z[faces][..., 2][..., None]
pixel_z = interpolate_face_attributes(fragments, faces_verts_z, bary_clip=True)
pixel_z = pixel_z.squeeze()[None, ...]  # (1, H, W, K)

Bear in mind that the lighting/texturing approach in SoftRas and Pytorch3d are different so may give slightly different results.

We have plans to support the texture atlas method soon.

Here are some example outputs with different settings for sigma and gamma:

sigma = 5e-4, gamma = 1e-4 sigma_5_e4_gamma_1e4

sigma = 1e-3, gamma = 1e-3 sigma_3_gamma_3

sigma = 1e-4, gamma = 1e-3 sigma_4_gamma_3

Gamma controls the opacity of the image so lower gamma means the image is more transparent - in the image below you can see the part of the back of the cow. sigma = 1e-6, gamma = 1e-2 sigma_6_gamma_2

aluo-x commented 4 years ago

Much appreciated! This does look more right.

Minor nitpick, I was under the impression that clamp only passed gradients to points within or on a given bound, and a quick test seems to confirm this.

import torch
a = torch.randn(20, requires_grad=True)
b = a.clamp(1)
b.sum().backward()
print(a.grad)

I have no idea if this is enough to be an issue. If not please feel free to close the issue.

nikhilaravi commented 4 years ago

@aluo-x If you want to still pass gradients in the backward pass but clamp in the forward pass then you can try using a 'fakeclamp' function. For example a simple version could look like this:

class fakeclamp(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        input_tensor,
        min=0.0,
        max=1.0,
    ):
        ctx.save_for_backward(input_tensor)
        return torch.clamp(input_tensor, min=min, max=max)

    @staticmethod
    def backward(
        ctx, grad_output_tensor
    ):
        input_tensor = ctx.saved_tensors[0]
        grad_input_tensor = input_tensor.new_ones(input_tensor.shape) * grad_output_tensor
        return grad_input_tensor, None, None
>> clamp = fakeclamp.apply
>> a = torch.randn(20, requires_grad=True)
>> b = clamp(a, max=0.6)
>> b.sum().backward()
>> print(a.grad)
>> tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1.])

For reference, in SoftRasterizer, the clamping is done in the CUDA kernel not in PyTorch. In the backward pass CUDA kernel, there is some gradient flow to the clipped inputs (from what I understand w > 1 clipped to 1.0 still gets a gradient but w < 0 clipped to 0 doesn't get a gradient). This isn't the complete implementation of clamp + normalize backward so it wouldn't pass any gradient tests in comparison to the PyTorch autograd version. The fakeclamp method above would give you gradients to all the inputs even if they are clipped.

We plan to add more support for blurry blending of textured meshes so stay tuned! :)

nikhilaravi commented 4 years ago

@aluo-x Let me know if this answers your question. Also what is the task you are working on? It would be easier to identify what the effect of clamping on the gradients would be depending on the use case e.g. are you use the RGB images for computing a loss?

aluo-x commented 4 years ago

Much appreciated! You really went above and beyond. It really has answered all of my questions. Hopefully the z interpolation bit can be merged into master.

The project that I am working involves many small, potentially very irregular meshes in a single scene that can warp and change depending on several factors, and the loss currently is RGB & mask with the usual regularization applied, in this respect having texture is less meaningful than having per face or per vertex colorization options (neural mesh renderer does really well with their per face option), since the bits of the scene can change depending on warp.

But there are also a few other losses that are currently being experimented with that depend on outputs that other diff renders do not provide, or do not provide in sufficiently high quality. Hopefully once it is a little bit mature I can talk about it more.

Closing this issue now. Thanks again!

nikhilaravi commented 4 years ago

@aluo-x Great, glad to help! :) We'll get the z interpolation with clipping added soon and also support for the per face texturing.

Regarding the "outputs that other diff renderers do not provide or do not provide in sufficiently high quality" could you share what these are? No worries if you are not ready to share this yet.

It's great for us to know how people are using the PyTorch3d renderer so we can prioritize improvements and features!

nikhilaravi commented 4 years ago

@aluo-x did you end up needing to use the fakeclamp function instead of torch.clamp?

aluo-x commented 4 years ago

Much appreciated for the followup.

I'm currently using the torch.clamp function without any bad effect. It should be easy to change to the fakeclamp function down the road I think if it ever becomes necessary since pytorch3d is so modular.

Minor (unrelated question), would it be possible to have documentation on how to canonically do multi GPU training? I could open another issue if it is better that way.

nikhilaravi commented 4 years ago

@aluo-x the barycentric clipping fix has now been landed in master.

Yes please raise a separate issue about multi GPU training.