chamfer_distance
chamfer_distance copied to clipboard
Implementation of the Chamfer Distance as a module for PyTorch
Chamfer Distance for pyTorch
This is an installable implementation of the Chamfer Distance as a module for pyTorch from Christian Diller. It is written as a custom C++/CUDA extension.
As it is using pyTorch's JIT compilation, there are no additional prerequisite steps that have to be taken. Simply import the module as shown below; CUDA and C++ code will be compiled on the first run.
Requirements
The only requirement is PyTotch with cuda support:
Installation
- Install PyTorch (>= 1.1.0)
- To install the package simply run the following line:
pip install git+'https://github.com/otaheri/chamfer_distance'
Usage
import torch
from chamfer_distance import ChamferDistance
import time
chamfer_dist = ChamferDistance()
p1 = torch.rand([10,25,3])
p2 = torch.rand([10,15,3])
s = time.time()
dist1, dist2, idx1, idx2 = chamfer_dist(p1,p2)
loss = (torch.mean(dist1)) + (torch.mean(dist2))
torch.cuda.synchronize()
print(f"Time: {time.time() - s} seconds")
print(f"Loss: {loss}")
#...
Integration
This code has been integrated into the Kaolin library for 3D Deep Learning by NVIDIAGameWorks. You should probably take a look at it if you are working on anything 3D :)