def test_cpugpu(self): b = 20 n = 1000 known = torch.randn((b, 2 * n, 3)) unknown = torch.randn((b, n, 3)) dist_cpu, cpu_idx = three_nn(unknown, known) dist_cuda, cuda_idx = three_nn(unknown.cuda(), known.cuda()) torch.testing.assert_allclose(dist_cpu, dist_cuda.cpu())
def test_gpu(self): pos = torch.randn([16, 100, 3]).cuda() pos_skip = torch.randn([16, 500, 3]).cuda() x = torch.randn([16, 30, 100], requires_grad=True).cuda() dist, idx = three_nn(pos_skip, pos) dist_recip = 1.0 / (dist + 1e-8) norm = torch.sum(dist_recip, dim=2, keepdim=True) weight = dist_recip / norm interpolated_feats = three_interpolate(x, idx, weight) dist, idx = three_nn(pos_skip.cpu(), pos.cpu()) dist_recip = 1.0 / (dist + 1e-8) norm = torch.sum(dist_recip, dim=2, keepdim=True) weight = dist_recip / norm interpolated_feats_cpu = three_interpolate(x.cpu(), idx, weight) torch.testing.assert_allclose(interpolated_feats_cpu, interpolated_feats.cpu())
def test_grad(self): b, n, k = (2, 10, 3) pos = torch.randn([b, n, k]).double() pos_skip = torch.randn([b, 2 * n, k]).double() x = torch.randn([b, 30, n], requires_grad=True).double() dist, idx = three_nn(pos_skip, pos) dist_recip = 1.0 / (dist + 1e-8) norm = torch.sum(dist_recip, dim=2, keepdim=True) weight = dist_recip / norm input = (x, idx, weight.double()) test = gradcheck(three_interpolate, input, eps=1e-6, atol=1e-4)
def conv(self, pos, pos_skip, x): assert pos_skip.shape[2] == 3 if pos is not None: dist, idx = tp.three_nn(pos_skip, pos) dist_recip = 1.0 / (dist + 1e-8) norm = torch.sum(dist_recip, dim=2, keepdim=True) weight = dist_recip / norm interpolated_feats = tp.three_interpolate(x, idx, weight) else: interpolated_feats = x.expand(*(x.size()[0:2] + (pos_skip.size(1), ))) return interpolated_feats