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.
The only requirement is PyTotch with cuda support:
- Install PyTorch (>= 1.1.0)
- To install the package simply run the following line:
pip install git+'https://github.com/otaheri/chamfer_distance'
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}")
#...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 :)