Metadata-Version: 2.4
Name: torch-fps
Version: 0.3
Summary: Native PyTorch farthest point sampling for point cloud workloads
Author: Felix Yu
License-Expression: MIT
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Information Analysis
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Dynamic: license-file

# torch-fps

Optimized standard farthest point sampling (FPS) for PyTorch written in C++.

```bash
pip install torch-fps
```
**Note**: Ensure gcc > 9 and < 14. Install might take a while since its building from source (to be fixed in future). 

## Usage

```python
import torch
from torch_fps import farthest_point_sampling, farthest_point_sampling_with_knn

# Create example inputs
points = torch.randn(4, 1000, 3)     # [B, N, D] - batch of point clouds
mask = torch.ones(4, 1000, dtype=torch.bool)  # [B, N] - valid point mask
K = 512  # Number of samples per batch (must be <= number of valid points)

# Perform farthest point sampling
idx = farthest_point_sampling(points, mask, K)  # [B, K] - selected point indices

# Use indices to gather sampled points
sampled_points = points.gather(1, idx.unsqueeze(-1).expand(-1, -1, 3))  # [B, K, D]

# Fused FPS + kNN: get centroids and their k nearest neighbors in one pass
centroid_idx, neighbor_idx = farthest_point_sampling_with_knn(
    points, mask, K=512, k_neighbors=32
)  # centroid_idx: [B, K], neighbor_idx: [B, K, k_neighbors]
```

## Performance

Benchmarked on AMD Threadripper 7970X and NVIDIA RTX 5090. Values show CPU / CUDA measurements. By default uses float32; override with `precision=` parameter.
Numbers below come from the in-repo benchmark script (`python tests/profile.py`) against the local extension build.

**FPS:**

| B  | N    | K   | Baseline (ms)   | Optimized (ms) | Speedup        |
|---:|-----:|----:|----------------:|---------------:|---------------:|
| 4  | 100  | 20  | 0.45 / 1.38     | 0.05 / 0.10    | 9.50x / 13.70x |
| 8  | 512  | 64  | 2.88 / 4.05     | 0.11 / 0.17    | 25.36x / 23.89x |
| 16 | 1024 | 128 | 29.92 / 7.81    | 0.39 / 0.31    | 77.57x / 25.42x |
| 32 | 2048 | 256 | 154.44 / 15.59  | 1.52 / 0.74    | 101.92x / 21.16x |

**FPS+kNN:**

| B  | N    | K   | k  | Baseline (ms)   | Optimized (ms) | Speedup        |
|---:|-----:|----:|---:|----------------:|---------------:|---------------:|
| 4  | 100  | 16  | 8  | 0.50 / 1.21     | 0.05 / 0.21    | 9.98x / 5.80x  |
| 8  | 512  | 64  | 16 | 4.90 / 4.11     | 0.21 / 1.13    | 23.56x / 3.65x |
| 16 | 1024 | 128 | 16 | 37.60 / 8.08    | 0.80 / 2.24    | 46.88x / 3.60x |
| 32 | 2048 | 256 | 16 | 180.33 / 16.86  | 2.57 / 4.84    | 70.24x / 3.48x |
