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.
I updated the package to use a wrapper around the Pytorch3D package chamfer distance due to some gradients bugs in the original code. Please update to the new version if you face any issues.
The only requirements are PyTotch and Pytorch3D with cuda support:
- Install PyTorch (>= 1.1.0)
- Install PyTorch3d gpu version.
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
- 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 as chamfer_dist
import time
p1 = torch.rand([10,25,3])
p2 = torch.rand([10,15,3])
s = time.time()
chd = chamfer_dist()
dist1, dist2, idx1, idx2 = chd(p1,p2)
loss = (torch.mean(dist1)) + (torch.mean(dist2))
torch.cuda.synchronize()
print(f"Time: {time.time() - s} seconds")
print(f"Loss: {loss}")
#...