Closed qifeng22 closed 3 months ago
import random from typing import Any, Dict import numpy as np import torch import torch.utils.data from torch import nn from torch.cuda import amp import torchsparse from torchsparse import SparseTensor from torchsparse import nn as spnn from torchsparse.nn import functional as F from torchsparse.utils.collate import sparse_collate_fn from torchsparse.utils.quantize import sparse_quantize inputs = np.random.uniform(-10, 10, size=(10, 4)) coords, feats = inputs[:, :3], inputs coords -= np.min(coords, axis=0, keepdims=True) coords, indices = sparse_quantize(coords, 0.01, return_index=True) coords = torch.tensor(coords, dtype=torch.int) feats = torch.tensor(feats[indices], dtype=torch.float) input = SparseTensor(coords=coords, feats=feats) tt = spnn.Conv3d(4, 3, 1,stride=2).cuda()(input) print(tt.feats)
For SparseTensor, the 'coords' need to be a four-channel vector, where the first three dimensions represent the voxelized coordinates, and the last one should indicate the batch index.