| 168 | torch.cuda.set_device(torch.device("cuda:0")) |
| 169 | |
| 170 | def knn(x, src, k, transpose=False): |
| 171 | if transpose: |
| 172 | x = x.transpose(1, 2).contiguous() |
| 173 | src = src.transpose(1, 2).contiguous() |
| 174 | b, n, _ = x.shape |
| 175 | m = src.shape[1] |
| 176 | x = x.view(-1, 3) |
| 177 | src = src.view(-1, 3) |
| 178 | x_offset = torch.full((b,), n, dtype=torch.long, device=x.device) |
| 179 | src_offset = torch.full((b,), m, dtype=torch.long, device=x.device) |
| 180 | x_offset = torch.cumsum(x_offset, dim=0).int() |
| 181 | src_offset = torch.cumsum(src_offset, dim=0).int() |
| 182 | idx, dists = knnquery(k, src, x, src_offset, x_offset) |
| 183 | idx = idx.view(b, n, k) - (src_offset - m)[:, None, None] |
| 184 | return idx.long(), dists.view(b, n, k) |
| 185 | |
| 186 | def fps(x, k): |
| 187 | b, n, _ = x.shape |