diff --git a/python/__init__.py b/python/__init__.py new file mode 100644 index 0000000..4228645 --- /dev/null +++ b/python/__init__.py @@ -0,0 +1 @@ +from .sparse import * \ No newline at end of file diff --git a/python/__pycache__/__init__.cpython-39.pyc b/python/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..51c3946 Binary files /dev/null and b/python/__pycache__/__init__.cpython-39.pyc differ diff --git a/python/__pycache__/sparse.cpython-39.pyc b/python/__pycache__/sparse.cpython-39.pyc new file mode 100644 index 0000000..1a0ba05 Binary files /dev/null and b/python/__pycache__/sparse.cpython-39.pyc differ diff --git a/python/__pycache__/utils.cpython-39.pyc b/python/__pycache__/utils.cpython-39.pyc new file mode 100644 index 0000000..29217de Binary files /dev/null and b/python/__pycache__/utils.cpython-39.pyc differ diff --git a/python/indice_manager.py b/python/indice_manager.py new file mode 100644 index 0000000..18e16a2 --- /dev/null +++ b/python/indice_manager.py @@ -0,0 +1,34 @@ +import os +import numpy as np +from typing import Union, List, Tuple + +import jittor as jt +from jittor import Function + +class IndiceManager: + def __init__( + self, + ndim, + # indice_map_type, + # sparse_alorithm, # set m_hashtable_occupancy for concurrent_unordered_map + ): + # if indice_map_type == 'GPU': + # assert(jt.has_cuda) + + self.ndim = ndim + # self.indice_map_type = indice_map_type + # self.sparse_algorithm = sparse_alorithm + self.stride_key_manager = {} + self.indice_map_manager = {} + self.kernel_map_manager = {} + + def insert(self, stride, indice_key, indice_hash): + self.stride_key_manager[stride] = indice_key + self.indice_map_manager[indice_key] = indice_hash + + +# class IndiceMapManager + + + + \ No newline at end of file diff --git a/python/nn/__init__.py b/python/nn/__init__.py new file mode 100644 index 0000000..270dceb --- /dev/null +++ b/python/nn/__init__.py @@ -0,0 +1 @@ +from .modules import * diff --git a/python/nn/__pycache__/__init__.cpython-39.pyc b/python/nn/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..3d661fc Binary files /dev/null and b/python/nn/__pycache__/__init__.cpython-39.pyc differ diff --git a/python/nn/functional/__init__.py b/python/nn/functional/__init__.py index de6e4d9..240b7d0 100644 --- a/python/nn/functional/__init__.py +++ b/python/nn/functional/__init__.py @@ -1,10 +1,10 @@ -# from .activation import * +from .activation import * from .conv import * # from .count import * # from .crop import * # from .devoxelize import * from .downsample import * from .hash import * -# from .pooling import * +from .pooling import * from .query import * # from .voxelize import * diff --git a/python/nn/functional/__pycache__/__init__.cpython-39.pyc b/python/nn/functional/__pycache__/__init__.cpython-39.pyc index f98f401..dafd0b9 100644 Binary files a/python/nn/functional/__pycache__/__init__.cpython-39.pyc and b/python/nn/functional/__pycache__/__init__.cpython-39.pyc differ diff --git a/python/nn/functional/__pycache__/activation.cpython-39.pyc b/python/nn/functional/__pycache__/activation.cpython-39.pyc new file mode 100644 index 0000000..6b15650 Binary files /dev/null and b/python/nn/functional/__pycache__/activation.cpython-39.pyc differ diff --git a/python/nn/functional/__pycache__/conv.cpython-39.pyc b/python/nn/functional/__pycache__/conv.cpython-39.pyc index 43dad00..0b42695 100644 Binary files a/python/nn/functional/__pycache__/conv.cpython-39.pyc and b/python/nn/functional/__pycache__/conv.cpython-39.pyc differ diff --git a/python/nn/functional/__pycache__/pooling.cpython-39.pyc b/python/nn/functional/__pycache__/pooling.cpython-39.pyc new file mode 100644 index 0000000..2e8e8e8 Binary files /dev/null and b/python/nn/functional/__pycache__/pooling.cpython-39.pyc differ diff --git a/python/nn/functional/activation.py b/python/nn/functional/activation.py new file mode 100644 index 0000000..52b3158 --- /dev/null +++ b/python/nn/functional/activation.py @@ -0,0 +1,24 @@ +import jittor as jt +import jittor.nn as nn + +from python import SparseTensor +from python.nn.utils import fapply + +__all__ = ['relu', 'leaky_relu'] +# __all__ = ['relu', 'leaky_relu', 'ReLU', 'LeakyReLU'] + +def relu(input: SparseTensor) -> SparseTensor: + return fapply(input, nn.relu) + + +def leaky_relu(input: SparseTensor, + scale: float = 0.01) -> SparseTensor: + return fapply(input, + nn.leaky_relu, + scale=scale) + +# Relu = jt.make_module(relu) +# ReLU = Relu +# Leaky_relu = jt.make_module(leaky_relu, 2) +# LeakyReLU = Leaky_relu + diff --git a/python/nn/functional/conv.py b/python/nn/functional/conv.py index 5812dc8..096210c 100644 --- a/python/nn/functional/conv.py +++ b/python/nn/functional/conv.py @@ -270,6 +270,7 @@ def conv3d( bias: Optional[jt.Var] = None, stride: Union[int, Tuple[int, ...]] = 1, dilation: Union[int, Tuple[int, ...]] = 1, + group: int = 1, transposed: bool = False, ) -> SparseTensor: # kernel_size = make_ntuple(kernel_size, ndim=3) @@ -337,10 +338,12 @@ def conv3d( if bias is not None: output_values += bias + # size have to be set output = SparseTensor( indices=output_indices, values=output_values, stride=output_stride, + size=input.size ) output.cmaps = input.cmaps output.cmaps.setdefault(output_stride, output_indices) diff --git a/python/nn/functional/pooling.py b/python/nn/functional/pooling.py new file mode 100644 index 0000000..3d84fce --- /dev/null +++ b/python/nn/functional/pooling.py @@ -0,0 +1,26 @@ +import jittor as jt + +from python import SparseTensor + +__all__ = ['global_avg_pool', 'global_max_pool'] + +def global_avg_pool(inputs: SparseTensor) -> jt.Var: + batch_size = jt.max(inputs.indices[:, 0]).item() + 1 + outputs = [] + for k in range(batch_size): + input = inputs.values[inputs.indices[:, 0] == k] + output = jt.mean(input, dim=0) + outputs.append(output) + outputs = jt.stack(outputs, dim=0) + return outputs + + +def global_max_pool(inputs: SparseTensor) -> jt.Var: + batch_size = jt.max(inputs.indices[:, 0]).item() + 1 + outputs = [] + for k in range(batch_size): + input = inputs.values[inputs.indices[:, 0] == k] + output = jt.max(input, dim=0)[0] + outputs.append(output) + outputs = jt.stack(outputs, dim=0) + return outputs \ No newline at end of file diff --git a/python/nn/modules/__init__.py b/python/nn/modules/__init__.py new file mode 100644 index 0000000..82665d0 --- /dev/null +++ b/python/nn/modules/__init__.py @@ -0,0 +1,6 @@ +from .activation import * +# from .bev import * +from .conv import * +# from .crop import * +from .norm import * +from .pooling import * diff --git a/python/nn/modules/__pycache__/__init__.cpython-39.pyc b/python/nn/modules/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..bd3d116 Binary files /dev/null and b/python/nn/modules/__pycache__/__init__.cpython-39.pyc differ diff --git a/python/nn/modules/__pycache__/activation.cpython-39.pyc b/python/nn/modules/__pycache__/activation.cpython-39.pyc new file mode 100644 index 0000000..4666e13 Binary files /dev/null and b/python/nn/modules/__pycache__/activation.cpython-39.pyc differ diff --git a/python/nn/modules/__pycache__/conv.cpython-39.pyc b/python/nn/modules/__pycache__/conv.cpython-39.pyc new file mode 100644 index 0000000..fe3e238 Binary files /dev/null and b/python/nn/modules/__pycache__/conv.cpython-39.pyc differ diff --git a/python/nn/modules/__pycache__/norm.cpython-39.pyc b/python/nn/modules/__pycache__/norm.cpython-39.pyc new file mode 100644 index 0000000..334f604 Binary files /dev/null and b/python/nn/modules/__pycache__/norm.cpython-39.pyc differ diff --git a/python/nn/modules/__pycache__/pooling.cpython-39.pyc b/python/nn/modules/__pycache__/pooling.cpython-39.pyc new file mode 100644 index 0000000..9d8fc1e Binary files /dev/null and b/python/nn/modules/__pycache__/pooling.cpython-39.pyc differ diff --git a/python/nn/modules/activation.py b/python/nn/modules/activation.py new file mode 100644 index 0000000..e751d58 --- /dev/null +++ b/python/nn/modules/activation.py @@ -0,0 +1,21 @@ +import jittor as jt +from jittor import nn + +from python import SparseTensor +from python.nn.functional import relu, leaky_relu +# from python.nn.utils import fapply + +__all__ = ['ReLU', 'LeakyReLU'] + +# class ReLU(nn.ReLU): +# def execute(self, input: SparseTensor) -> SparseTensor: +# return fapply(input, super().execute) + +# class LeakyReLU(nn.LeakyReLU): +# def execute(self, input: SparseTensor) -> SparseTensor: +# return fapply(input, super().execute) + +Relu = jt.make_module(relu) +ReLU = Relu +Leaky_relu = jt.make_module(leaky_relu, 2) +LeakyReLU = Leaky_relu \ No newline at end of file diff --git a/python/nn/modules/conv3d.py b/python/nn/modules/conv.py similarity index 99% rename from python/nn/modules/conv3d.py rename to python/nn/modules/conv.py index d96ec97..5790979 100644 --- a/python/nn/modules/conv3d.py +++ b/python/nn/modules/conv.py @@ -9,7 +9,7 @@ from jittor.misc import _pair, _triple from python import SparseTensor from python.nn import functional as F -from utils import make_ntuple +# from utils import make_ntuple __all__ = ['Conv3d'] diff --git a/python/nn/modules/norm.py b/python/nn/modules/norm.py new file mode 100644 index 0000000..de483f4 --- /dev/null +++ b/python/nn/modules/norm.py @@ -0,0 +1,34 @@ +import jittor as jt +from jittor import nn +from numpy import kaiser + +from python import SparseTensor +from python.nn.utils import fapply + +__all__ = ['BatchNorm', 'GroupNorm'] + +class BatchNorm(nn.BatchNorm): + def execute(self, input: SparseTensor) -> SparseTensor: + return fapply(input, super().execute) + +class GroupNorm(nn.GroupNorm): + def execute(self, input: SparseTensor) -> SparseTensor: + indices, values, stride, size = input.indices, input.values, input.stride, input.size + + batch_size = jt.max(indices[:, 0]).item() + 1 + num_channels = values.shape[1] + + n_values = jt.zeros_like(values) + for k in range(batch_size): + idx = indices[:, 0] == k + b_values = values[idx] + b_values = b_values.t().reshape(1, num_channels, -1) + b_values = super().execute(b_values) + b_values = b_values.reshape(num_channels, -1).t() + n_values[idx] = b_values + + output = SparseTensor(indices=indices, values=n_values, stride=stride, size=size) + output.cmaps = input.cmaps + output.kmaps = input.kmaps + return output + diff --git a/python/nn/modules/pooling.py b/python/nn/modules/pooling.py new file mode 100644 index 0000000..871ea00 --- /dev/null +++ b/python/nn/modules/pooling.py @@ -0,0 +1,11 @@ +from ast import Global +import jittor as jt +from jittor import nn + +from python import SparseTensor +from python.nn.functional import global_avg_pool, global_max_pool + +__all__ = ['GlobalAvgPool', 'GlobalMaxPool'] + +GlobalAvgPool = jt.make_module(global_avg_pool) +GlobalMaxPool = jt.make_module(global_max_pool) \ No newline at end of file diff --git a/python/nn/utils/__init__.py b/python/nn/utils/__init__.py index 662c212..919f5b0 100644 --- a/python/nn/utils/__init__.py +++ b/python/nn/utils/__init__.py @@ -1 +1,2 @@ +from .apply import * from .kernel import * \ No newline at end of file diff --git a/python/nn/utils/__pycache__/__init__.cpython-39.pyc b/python/nn/utils/__pycache__/__init__.cpython-39.pyc index 0b03769..6d461f4 100644 Binary files a/python/nn/utils/__pycache__/__init__.cpython-39.pyc and b/python/nn/utils/__pycache__/__init__.cpython-39.pyc differ diff --git a/python/nn/utils/__pycache__/apply.cpython-39.pyc b/python/nn/utils/__pycache__/apply.cpython-39.pyc new file mode 100644 index 0000000..722228e Binary files /dev/null and b/python/nn/utils/__pycache__/apply.cpython-39.pyc differ diff --git a/python/nn/utils/apply.py b/python/nn/utils/apply.py new file mode 100644 index 0000000..2ca3432 --- /dev/null +++ b/python/nn/utils/apply.py @@ -0,0 +1,15 @@ +from typing import Callable + +import jittor as jt + +from python import SparseTensor + +__all__ = ['fapply'] + +def fapply(input: SparseTensor, fn: Callable[..., jt.Var], *args, + **kwargs) -> SparseTensor: + values = fn(input.values, *args, **kwargs) + output = SparseTensor(indices=input.indices, values=values, stride=input.stride, size=input.size) + output.cmaps = input.cmaps + output.kmaps = input.kmaps + return output \ No newline at end of file diff --git a/python/sparse.py b/python/sparse.py index 9baee94..3484ab7 100644 --- a/python/sparse.py +++ b/python/sparse.py @@ -24,11 +24,13 @@ class SparseTensor: assert (values.ndim == 2) # self.indices = indices # self.values = values - self.shape = size + self.size = size self.ndim = indices.shape[1] - 1 self.stride = make_ntuple(stride, ndim=self.ndim) self.voxel_size = voxel_size self.coalesce_mode = coalesce_mode + self.cmaps = {} + self.kmaps = {} ########################## # Setup CoordsManager diff --git a/python/sparse_dense_function.py b/python/sparse_dense_function.py new file mode 100644 index 0000000..f8bb855 --- /dev/null +++ b/python/sparse_dense_function.py @@ -0,0 +1,42 @@ +import jittor as jt +from jittor import Function + + + + +def spmm( + rows: jt.Var, + cols: jt.Var, + vals: jt.Var, + size: jt.NanoVector, + mat: jt.Var, + spmm_mode='scatter', + is_sorted: bool = False, + cuda_spmm_alg: int = 1, +) -> jt.Var: + assert len(rows) == len(cols), "Invalid length" + assert len(rows) == len(vals), "Invalid length" + assert vals.dtype == mat.dtype, "dtype mismatch" + + if jt.flags.use_cuda > 1: + assert jt.has_cuda == 1, "No GPUs available" + rows = rows.int() + cols = cols.int() + ''' + TODO: Using the coo_spmm of cuSPARSE on GPU + result = coo_spmm_int32( + rows, cols, vals, size[0], size[1], mat, cuda_spmm_alg, is_sorted + ) + ''' + else: + if (spmm_mode == 'scatter'): + + + + + + + + + +class SPMM(Function):