mit-han-lab / torchsparse

[MICRO'23, MLSys'22] TorchSparse: Efficient Training and Inference Framework for Sparse Convolution on GPUs.
https://torchsparse.mit.edu
MIT License
1.15k stars 131 forks source link

Can't use the spnn.Conv3d #295

Open qifeng22 opened 4 months ago

qifeng22 commented 4 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)

image

zhijian-liu commented 2 months ago

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.