(x, src, k, transpose=False)
| 135 | torch.cuda.set_device(torch.device("cuda:0")) |
| 136 | |
| 137 | def knn(x, src, k, transpose=False): |
| 138 | if transpose: |
| 139 | x = x.transpose(1, 2).contiguous() |
| 140 | src = src.transpose(1, 2).contiguous() |
| 141 | b, n, _ = x.shape |
| 142 | m = src.shape[1] |
| 143 | x = x.view(-1, 3) |
| 144 | src = src.view(-1, 3) |
| 145 | x_offset = torch.full((b,), n, dtype=torch.long, device=x.device) |
| 146 | src_offset = torch.full((b,), m, dtype=torch.long, device=x.device) |
| 147 | x_offset = torch.cumsum(x_offset, dim=0).int() |
| 148 | src_offset = torch.cumsum(src_offset, dim=0).int() |
| 149 | idx, dists = knnquery(k, src, x, src_offset, x_offset) |
| 150 | idx = idx.view(b, n, k) - (src_offset - m)[:, None, None] |
| 151 | return idx.long(), dists.view(b, n, k) |
| 152 | |
| 153 | def fps(x, k): |
| 154 | b, n, _ = x.shape |
nothing calls this directly
no outgoing calls
no test coverage detected