def forward(ctx,
features,
kernel,
neighbor_map,
neighbor_offset,
sizes,
transpose=False):
r"""
features : torch.FloatTensor (N, c_in) Features of the input point cloud.
kernel : torch.FloatTensor (K, c_in, c_out) Kernel. with K to be kernel volume.
neighbor_map: torch.IntTensor (N_nonzero, 2) from -> to
return: (N', c_out).
"""
features = features.contiguous()
kernel = kernel.contiguous()
if not transpose:
out = torch.zeros(sizes[1], kernel.size(-1), device=features.device)
else:
out = torch.zeros(sizes[0], kernel.size(-1), device=features.device)
if 'cuda' in str(features.device):
torchsparse_cuda.sparseconv_forward(features, out, kernel, neighbor_map, neighbor_offset, transpose)
else:
raise NotImplementedError
ctx.for_backwards = (features, kernel, neighbor_map, neighbor_offset, transpose)
return out
def forward(ctx,
features,
kernel,
neighbor_map,
neighbor_offset,
sizes,
transpose=False):
features = features.contiguous()
kernel = kernel.contiguous()
if not transpose:
out = torch.zeros(sizes[1],
kernel.size(-1),
device=features.device)
else:
# tbd: ensure the original, upsampled size to be the same.
out = torch.zeros(sizes[0],
kernel.size(-1),
device=features.device)
if 'cuda' in str(features.device):
torchsparse_cuda.sparseconv_forward(features, out, kernel,
neighbor_map, neighbor_offset,
transpose)
#elif 'cpu' in str(features.device):
# torchsparse_cuda.sparseconv_cpu_forward(features, out, kernel, neighbor_map, neighbor_offset.cpu(), transpose)
else:
# use the native pytorch XLA APIs for the TPU.
cur_st = 0
for kernel_idx in range(kernel.shape[0]):
cur_ed = cur_st + neighbor_offset[kernel_idx]
in_map = neighbor_map[cur_st:cur_ed, 0].long()
out_map = neighbor_map[cur_st:cur_ed, 1].long()
cur_st += neighbor_offset[kernel_idx]
if transpose:
in_map, out_map = out_map, in_map
# gather
cur_feat = features[in_map]
# gemm
cur_feat = torch.mm(cur_feat, kernel[kernel_idx])
# scatter
out[out_map] += cur_feat
ctx.for_backwards = (features, kernel, neighbor_map, neighbor_offset,
transpose)
return out
def forward(ctx,
features,
kernel,
neighbor_map,
neighbor_offset,
sizes,
transpose=False):
# type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor
r"""
Parameters
----------
features : torch.FloatTensor
(N, c_in) Features of the input point cloud.
kernel : torch.FloatTensor
(K, c_in, c_out) Kernel. with K to be kernel volume.
neighbor_map: torch.IntTensor
(K, N) K-volumetric neighborhood of each point.
For entries with value -1, it means that this neighbor is inexistent.
Returns
-------
torch.FloatTensor
(N, c_out) The output tensor.
"""
features = features.contiguous()
kernel = kernel.contiguous()
if not transpose:
out = torch.zeros(sizes[1],
kernel.size(-1),
device=features.device)
else:
# tbd: ensure the original, upsampled size to be the same.
out = torch.zeros(sizes[0],
kernel.size(-1),
device=features.device)
if 'cuda' in str(features.device):
torchsparse_cuda.sparseconv_forward(features, out, kernel,
neighbor_map, neighbor_offset,
transpose)
#elif 'cpu' in str(features.device):
# torchsparse_cuda.sparseconv_cpu_forward(features, out, kernel, neighbor_map, neighbor_offset.cpu(), transpose)
else:
# use the native pytorch XLA APIs for the TPU.
cur_st = 0
for kernel_idx in range(kernel.shape[0]):
cur_ed = cur_st + neighbor_offset[kernel_idx]
in_map = neighbor_map[cur_st:cur_ed, 0].long()
out_map = neighbor_map[cur_st:cur_ed, 1].long()
cur_st += neighbor_offset[kernel_idx]
if transpose:
in_map, out_map = out_map, in_map
# gather
cur_feat = features[in_map]
# gemm
cur_feat = torch.mm(cur_feat, kernel[kernel_idx])
# scatter
out[out_map] += cur_feat
ctx.for_backwards = (features, kernel, neighbor_map, neighbor_offset,
transpose)
return out