Skip to content

Implementation of the Chamfer Distance as a module for PyTorch

License

Notifications You must be signed in to change notification settings

otaheri/chamfer_distance

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

32 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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.

Update

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.

Requirements

The only requirements are PyTotch and Pytorch3D with cuda support:

Installation

  1. Install PyTorch (>= 1.1.0)
  2. Install PyTorch3d gpu version.
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
  1. 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 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}")

#...