Commit 19e73bbe authored by Yan Yan's avatar Yan Yan
Browse files

format code with clang-format, better c++ code

parent c336139f
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
...@@ -14,14 +14,14 @@ ...@@ -14,14 +14,14 @@
#pragma once #pragma once
#include <chrono> #include <chrono>
#ifdef SPCONV_CUDA #ifdef TV_CUDA
#include <cuda_runtime_api.h> #include <cuda_runtime_api.h>
#endif #endif
#include <iostream> #include <iostream>
namespace spconv { namespace spconv {
#ifdef SPCONV_CUDA #ifdef TV_CUDA
template <typename TimeT = std::chrono::microseconds> struct CudaContextTimer { template <typename TimeT = std::chrono::microseconds> struct CudaContextTimer {
CudaContextTimer() { CudaContextTimer() {
cudaDeviceSynchronize(); cudaDeviceSynchronize();
......
import os import os
import re
import sys
import platform import platform
import re
import subprocess import subprocess
import torch import sys
from setuptools import setup, Extension, find_packages
from setuptools.command.build_ext import build_ext
from distutils.version import LooseVersion from distutils.version import LooseVersion
from pathlib import Path from pathlib import Path
import torch
from setuptools import Extension, find_packages, setup
from setuptools.command.build_ext import build_ext
# if 'LIBTORCH_ROOT' not in os.environ: # if 'LIBTORCH_ROOT' not in os.environ:
# raise ValueError("You must set LIBTORCH_ROOT to your torch c++ library.") # raise ValueError("You must set LIBTORCH_ROOT to your torch c++ library.")
...@@ -100,4 +100,3 @@ setup( ...@@ -100,4 +100,3 @@ setup(
cmdclass=dict(build_ext=CMakeBuild), cmdclass=dict(build_ext=CMakeBuild),
zip_safe=False, zip_safe=False,
) )
...@@ -12,21 +12,20 @@ ...@@ -12,21 +12,20 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import platform
from pathlib import Path from pathlib import Path
import platform
import numpy as np import numpy as np
import torch import torch
from spconv import utils
from spconv.conv import SparseConv2d, SparseConv3d, SubMConv2d, SubMConv3d from spconv import ops, utils
from spconv.conv import SparseConvTranspose2d, SparseConvTranspose3d from spconv.conv import (SparseConv2d, SparseConv3d, SparseConvTranspose2d,
from spconv.conv import SparseInverseConv2d, SparseInverseConv3d SparseConvTranspose3d, SparseInverseConv2d,
SparseInverseConv3d, SubMConv2d, SubMConv3d)
from spconv.identity import Identity
from spconv.modules import SparseModule, SparseSequential from spconv.modules import SparseModule, SparseSequential
from spconv.pool import SparseMaxPool2d, SparseMaxPool3d from spconv.pool import SparseMaxPool2d, SparseMaxPool3d
from spconv.tables import ConcatTable, JoinTable, AddTable from spconv.tables import AddTable, ConcatTable, JoinTable
from spconv.identity import Identity
from spconv import ops
_LIB_FILE_NAME = "libspconv.so" _LIB_FILE_NAME = "libspconv.so"
if platform.system() == "Windows": if platform.system() == "Windows":
...@@ -34,6 +33,7 @@ if platform.system() == "Windows": ...@@ -34,6 +33,7 @@ if platform.system() == "Windows":
_LIB_PATH = str(Path(__file__).parent / _LIB_FILE_NAME) _LIB_PATH = str(Path(__file__).parent / _LIB_FILE_NAME)
torch.ops.load_library(_LIB_PATH) torch.ops.load_library(_LIB_PATH)
def scatter_nd(indices, updates, shape): def scatter_nd(indices, updates, shape):
"""pytorch edition of tensorflow scatter_nd. """pytorch edition of tensorflow scatter_nd.
this function don't contain except handle code. so use this carefully this function don't contain except handle code. so use this carefully
...@@ -49,8 +49,10 @@ def scatter_nd(indices, updates, shape): ...@@ -49,8 +49,10 @@ def scatter_nd(indices, updates, shape):
ret[slices] = updates.view(*output_shape) ret[slices] = updates.view(*output_shape)
return ret return ret
class SparseConvTensor(object): class SparseConvTensor(object):
def __init__(self, features, indices, spatial_shape, batch_size, grid=None): def __init__(self, features, indices, spatial_shape, batch_size,
grid=None):
""" """
Args: Args:
grid: pre-allocated grid tensor. should be used when the volume of spatial shape grid: pre-allocated grid tensor. should be used when the volume of spatial shape
...@@ -77,7 +79,8 @@ class SparseConvTensor(object): ...@@ -77,7 +79,8 @@ class SparseConvTensor(object):
return None return None
def dense(self, channels_first=True): def dense(self, channels_first=True):
output_shape = [self.batch_size] + list(self.spatial_shape) + [self.features.shape[1]] output_shape = [self.batch_size] + list(
self.spatial_shape) + [self.features.shape[1]]
res = scatter_nd(self.indices.long(), self.features, output_shape) res = scatter_nd(self.indices.long(), self.features, output_shape)
if not channels_first: if not channels_first:
return res return res
...@@ -88,7 +91,8 @@ class SparseConvTensor(object): ...@@ -88,7 +91,8 @@ class SparseConvTensor(object):
@property @property
def sparity(self): def sparity(self):
return self.indices.shape[0] / np.prod(self.spatial_shape) / self.batch_size return self.indices.shape[0] / np.prod(
self.spatial_shape) / self.batch_size
class ToDense(SparseModule): class ToDense(SparseModule):
...@@ -97,6 +101,7 @@ class ToDense(SparseModule): ...@@ -97,6 +101,7 @@ class ToDense(SparseModule):
def forward(self, x: SparseConvTensor): def forward(self, x: SparseConvTensor):
return x.dense() return x.dense()
class RemoveGrid(SparseModule): class RemoveGrid(SparseModule):
"""remove pre-allocated grid buffer. """remove pre-allocated grid buffer.
""" """
......
This diff is collapsed.
# Copyright 2019 Yan Yan # Copyright 2019 Yan Yan
# #
# Licensed under the Apache License, Version 2.0 (the "License"); # Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. # you may not use this file except in compliance with the License.
# You may obtain a copy of the License at # You may obtain a copy of the License at
# #
# http://www.apache.org/licenses/LICENSE-2.0 # http://www.apache.org/licenses/LICENSE-2.0
# #
# Unless required by applicable law or agreed to in writing, software # Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, # distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import spconv.ops as ops
import torch import torch
from torch import nn from torch import nn
from torch.autograd import Function from torch.autograd import Function
import spconv.ops as ops
class SparseConvFunction(Function): class SparseConvFunction(Function):
@staticmethod @staticmethod
def forward( def forward(ctx, features, filters, indice_pairs, indice_pair_num,
ctx, num_activate_out):
features, ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
filters, return ops.indice_conv(features, filters, indice_pairs,
indice_pairs, indice_pair_num, num_activate_out, False)
indice_pair_num,
num_activate_out):
ctx.save_for_backward(
indice_pairs,
indice_pair_num,
features,
filters)
return ops.indice_conv(features, filters, indice_pairs, indice_pair_num, num_activate_out, False)
@staticmethod @staticmethod
def backward(ctx, grad_output): def backward(ctx, grad_output):
indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
input_bp, filters_bp = ops.indice_conv_backward(features, filters, grad_output, indice_pairs, indice_pair_num, False) input_bp, filters_bp = ops.indice_conv_backward(
features, filters, grad_output, indice_pairs, indice_pair_num,
False)
return input_bp, filters_bp, None, None, None return input_bp, filters_bp, None, None, None
class SparseInverseConvFunction(Function): class SparseInverseConvFunction(Function):
@staticmethod @staticmethod
def forward( def forward(ctx, features, filters, indice_pairs, indice_pair_num,
ctx, num_activate_out):
features, ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
filters, return ops.indice_conv(features, filters, indice_pairs,
indice_pairs, indice_pair_num, num_activate_out, True, False)
indice_pair_num,
num_activate_out):
ctx.save_for_backward(
indice_pairs,
indice_pair_num,
features,
filters)
return ops.indice_conv(features, filters, indice_pairs, indice_pair_num, num_activate_out, True, False)
@staticmethod @staticmethod
def backward(ctx, grad_output): def backward(ctx, grad_output):
indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
input_bp, filters_bp = ops.indice_conv_backward(features, filters, grad_output, indice_pairs, indice_pair_num, True, False) input_bp, filters_bp = ops.indice_conv_backward(
features, filters, grad_output, indice_pairs, indice_pair_num,
True, False)
return input_bp, filters_bp, None, None, None return input_bp, filters_bp, None, None, None
class SubMConvFunction(Function): class SubMConvFunction(Function):
@staticmethod @staticmethod
def forward( def forward(ctx, features, filters, indice_pairs, indice_pair_num,
ctx, num_activate_out):
features, ctx.save_for_backward(indice_pairs, indice_pair_num, features, filters)
filters, return ops.indice_conv(features, filters, indice_pairs,
indice_pairs, indice_pair_num, num_activate_out, False, True)
indice_pair_num,
num_activate_out):
ctx.save_for_backward(
indice_pairs,
indice_pair_num,
features,
filters)
return ops.indice_conv(features, filters, indice_pairs, indice_pair_num, num_activate_out, False, True)
@staticmethod @staticmethod
def backward(ctx, grad_output): def backward(ctx, grad_output):
indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
input_bp, filters_bp = ops.indice_conv_backward(features, filters, grad_output, indice_pairs, indice_pair_num, False, True) input_bp, filters_bp = ops.indice_conv_backward(
features, filters, grad_output, indice_pairs, indice_pair_num,
False, True)
return input_bp, filters_bp, None, None, None return input_bp, filters_bp, None, None, None
class SparseMaxPoolFunction(Function): class SparseMaxPoolFunction(Function):
@staticmethod @staticmethod
def forward( def forward(ctx, features, indice_pairs, indice_pair_num,
ctx, num_activate_out):
features, out = ops.indice_maxpool(features, indice_pairs, indice_pair_num,
indice_pairs, num_activate_out)
indice_pair_num, ctx.save_for_backward(indice_pairs, indice_pair_num, features, out)
num_activate_out):
out = ops.indice_maxpool(features, indice_pairs, indice_pair_num, num_activate_out)
ctx.save_for_backward(
indice_pairs,
indice_pair_num,
features,
out)
return out return out
@staticmethod @staticmethod
def backward(ctx, grad_output): def backward(ctx, grad_output):
indice_pairs, indice_pair_num, features, out = ctx.saved_tensors indice_pairs, indice_pair_num, features, out = ctx.saved_tensors
input_bp = ops.indice_maxpool_backward(features, out, grad_output, indice_pairs, indice_pair_num) input_bp = ops.indice_maxpool_backward(features, out, grad_output,
indice_pairs, indice_pair_num)
return input_bp, None, None, None return input_bp, None, None, None
indice_conv = SparseConvFunction.apply indice_conv = SparseConvFunction.apply
indice_inverse_conv = SparseInverseConvFunction.apply indice_inverse_conv = SparseInverseConvFunction.apply
indice_subm_conv = SubMConvFunction.apply indice_subm_conv = SubMConvFunction.apply
......
...@@ -12,12 +12,13 @@ ...@@ -12,12 +12,13 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import time
from collections import OrderedDict from collections import OrderedDict
import spconv
import torch import torch
from torch import nn from torch import nn
import time
import spconv
def is_spconv_module(module): def is_spconv_module(module):
...@@ -81,7 +82,6 @@ class SparseSequential(SparseModule): ...@@ -81,7 +82,6 @@ class SparseSequential(SparseModule):
relu2=nn.ReLU() relu2=nn.ReLU()
) )
""" """
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super(SparseSequential, self).__init__() super(SparseSequential, self).__init__()
if len(args) == 1 and isinstance(args[0], OrderedDict): if len(args) == 1 and isinstance(args[0], OrderedDict):
...@@ -148,7 +148,8 @@ class SparseSequential(SparseModule): ...@@ -148,7 +148,8 @@ class SparseSequential(SparseModule):
idx = 0 idx = 0
while idx < len(mods): while idx < len(mods):
if is_sparse_conv(mods[idx]): if is_sparse_conv(mods[idx]):
if idx < len(mods) - 1 and isinstance(mods[idx + 1], nn.BatchNorm1d): if idx < len(mods) - 1 and isinstance(mods[idx + 1],
nn.BatchNorm1d):
new_module = SparseConvolution( new_module = SparseConvolution(
ndim=mods[idx].ndim, ndim=mods[idx].ndim,
in_channels=mods[idx].in_channels, in_channels=mods[idx].in_channels,
......
This diff is collapsed.
This diff is collapsed.
import torch
from torch.autograd import Function from torch.autograd import Function
import spconv
#from torch.nn import Module #from torch.nn import Module
from spconv.modules import SparseModule from spconv.modules import SparseModule
import spconv
import torch
class JoinTable(SparseModule):# Module):
class JoinTable(SparseModule): # Module):
def forward(self, input): def forward(self, input):
output = spconv.SparseConvTensor( output = spconv.SparseConvTensor(
torch.cat([i.features for i in input],1), input[1].indices, torch.cat([i.features for i in input], 1), input[1].indices,
input[1].spatial_shape, input[0].batch_size ) input[1].spatial_shape, input[0].batch_size)
output.indice_dict = input[1].indice_dict output.indice_dict = input[1].indice_dict
output.grid = input[1].grid output.grid = input[1].grid
return output return output
...@@ -18,11 +19,12 @@ class JoinTable(SparseModule):# Module): ...@@ -18,11 +19,12 @@ class JoinTable(SparseModule):# Module):
return out_size return out_size
class AddTable(SparseModule): # Module): class AddTable(SparseModule): # Module):
def forward(self, input): def forward(self, input):
output = spconv.SparseConvTensor( output = spconv.SparseConvTensor(sum([i.features for i in input]),
sum([i.features for i in input]), input[1].indices, input[1].indices,
input[1].spatial_shape, input[1].batch_size ) input[1].spatial_shape,
input[1].batch_size)
output.indice_dict = input[1].indice_dict output.indice_dict = input[1].indice_dict
output.grid = input[1].grid output.grid = input[1].grid
...@@ -32,7 +34,7 @@ class AddTable(SparseModule): # Module): ...@@ -32,7 +34,7 @@ class AddTable(SparseModule): # Module):
return out_size return out_size
class ConcatTable(SparseModule): # Module): class ConcatTable(SparseModule): # Module):
def forward(self, input): def forward(self, input):
return [module(input) for module in self._modules.values()] return [module(input) for module in self._modules.values()]
......
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment