daniilidis-group / neural_renderer

A PyTorch port of the Neural 3D Mesh Renderer
Other
1.12k stars 248 forks source link

speed up flip when rasterize #103

Open ALLinLLM opened 3 years ago

ALLinLLM commented 3 years ago

line 311 rasterize.py

    if return_rgb:
        rgb = rgb.permute((0, 3, 1, 2))
        # pytorch does not support negative slicing for the moment
        # may need to look at this again because it seems to be very slow
        # rgb = rgb[:, :, ::-1, :]
        rgb = rgb[:, :, list(reversed(range(rgb.shape[2]))), :]

after pytorch 0.4.0, we can use following native api to do reverse, which is really faster

import torch
def flip(x, dim):    
    indices = [slice(None)] * x.dim()
    indices[dim] = torch.arange(x.size(dim) - 1, -1, -1,
                                dtype=torch.long, device=x.device)
    return x[tuple(indices)]

And, in pytorch 1.2

just using tensor.flip() refer to https://pytorch.org/docs/1.2.0/torch.html?highlight=flip#torch.flip