• Georgia Gkioxari's avatar
    knn autograd · b2b0c5a4
    Georgia Gkioxari authored
    Summary:
    Adds knn backward to return `grad_pts1` and `grad_pts2`. Adds `knn_gather` to return the nearest neighbors in pts2.
    
    The BM tests include backward pass and are ran on an M40.
    ```
    Benchmark                               Avg Time(μs)      Peak Time(μs) Iterations
    --------------------------------------------------------------------------------
    KNN_SQUARE_32_256_128_3_24_cpu              39558           43485             13
    KNN_SQUARE_32_256_128_3_24_cuda:0            1080            1404            463
    KNN_SQUARE_32_256_512_3_24_cpu              81950           85781              7
    KNN_SQUARE_32_256_512_3_24_cuda:0            1519            1641            330
    --------------------------------------------------------------------------------
    
    Benchmark                               Avg Time(μs)      Peak Time(μs) Iterations
    --------------------------------------------------------------------------------
    KNN_RAGGED_32_256_128_3_24_cpu              13798           14650             37
    KNN_RAGGED_32_256_128_3_24_cuda:0            1576            1713            318
    KNN_RAGGED_32_256_512_3_24_cpu              31255           32210             16
    KNN_RAGGED_32_256_512_3_24_cuda:0            2024            2162            248
    --------------------------------------------------------------------------------
    ```
    
    Reviewed By: jcjohnson
    
    Differential Revision: D20945556
    
    fbshipit-source-id: a16f616029c6b5f8c2afceb5f2bc12c5c20d2f3c
    b2b0c5a4
bm_knn.py 698 Bytes