From 52beb6acf10826b3b60fa9135c9c93f549f55bc2 Mon Sep 17 00:00:00 2001 From: jkay Date: Mon, 25 Jul 2022 22:32:30 +0800 Subject: [PATCH] add activation, pool and norm --- python/__init__.py | 1 + python/__pycache__/__init__.cpython-39.pyc | Bin 0 -> 172 bytes python/__pycache__/sparse.cpython-39.pyc | Bin 0 -> 2175 bytes python/__pycache__/utils.cpython-39.pyc | Bin 0 -> 1229 bytes python/indice_manager.py | 34 ++++++++++++++ python/nn/__init__.py | 1 + python/nn/__pycache__/__init__.cpython-39.pyc | Bin 0 -> 176 bytes python/nn/functional/__init__.py | 4 +- .../__pycache__/__init__.cpython-39.pyc | Bin 239 -> 280 bytes .../__pycache__/activation.cpython-39.pyc | Bin 0 -> 679 bytes .../__pycache__/conv.cpython-39.pyc | Bin 7847 -> 7858 bytes .../__pycache__/pooling.cpython-39.pyc | Bin 0 -> 928 bytes python/nn/functional/activation.py | 24 ++++++++++ python/nn/functional/conv.py | 3 ++ python/nn/functional/pooling.py | 26 +++++++++++ python/nn/modules/__init__.py | 6 +++ .../__pycache__/__init__.cpython-39.pyc | Bin 0 -> 238 bytes .../__pycache__/activation.cpython-39.pyc | Bin 0 -> 442 bytes .../modules/__pycache__/conv.cpython-39.pyc | Bin 0 -> 1952 bytes .../modules/__pycache__/norm.cpython-39.pyc | Bin 0 -> 1436 bytes .../__pycache__/pooling.cpython-39.pyc | Bin 0 -> 465 bytes python/nn/modules/activation.py | 21 +++++++++ python/nn/modules/{conv3d.py => conv.py} | 2 +- python/nn/modules/norm.py | 34 ++++++++++++++ python/nn/modules/pooling.py | 11 +++++ python/nn/utils/__init__.py | 1 + .../utils/__pycache__/__init__.cpython-39.pyc | Bin 181 -> 198 bytes .../nn/utils/__pycache__/apply.cpython-39.pyc | Bin 0 -> 601 bytes python/nn/utils/apply.py | 15 +++++++ python/sparse.py | 4 +- python/sparse_dense_function.py | 42 ++++++++++++++++++ 31 files changed, 225 insertions(+), 4 deletions(-) create mode 100644 python/__init__.py create mode 100644 python/__pycache__/__init__.cpython-39.pyc create mode 100644 python/__pycache__/sparse.cpython-39.pyc create mode 100644 python/__pycache__/utils.cpython-39.pyc create mode 100644 python/indice_manager.py create mode 100644 python/nn/__init__.py create mode 100644 python/nn/__pycache__/__init__.cpython-39.pyc create mode 100644 python/nn/functional/__pycache__/activation.cpython-39.pyc create mode 100644 python/nn/functional/__pycache__/pooling.cpython-39.pyc create mode 100644 python/nn/functional/activation.py create mode 100644 python/nn/functional/pooling.py create mode 100644 python/nn/modules/__init__.py create mode 100644 python/nn/modules/__pycache__/__init__.cpython-39.pyc create mode 100644 python/nn/modules/__pycache__/activation.cpython-39.pyc create mode 100644 python/nn/modules/__pycache__/conv.cpython-39.pyc create mode 100644 python/nn/modules/__pycache__/norm.cpython-39.pyc create mode 100644 python/nn/modules/__pycache__/pooling.cpython-39.pyc create mode 100644 python/nn/modules/activation.py rename python/nn/modules/{conv3d.py => conv.py} (99%) create mode 100644 python/nn/modules/norm.py create mode 100644 python/nn/modules/pooling.py create mode 100644 python/nn/utils/__pycache__/apply.cpython-39.pyc create mode 100644 python/nn/utils/apply.py create mode 100644 python/sparse_dense_function.py 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 0000000000000000000000000000000000000000..51c3946201cc127c2531c3de548e6a5f52ef66d0 GIT binary patch literal 172 zcmYe~<>g`kf>$@LC5Zy*#~=OI3Wi_kGi_+(N)#;=t=wiE`h=OzuvnU z5%L=%=Z6ExEto0>zzL@XDQQe8?peWN$F6Sd!Ri!#8N@*u#$nlsJ7qWSS~<7ql~Ej( z{kUHa;sGUJ5bkmRG2#A-#Y4{Sl5ubX`pK9YFRPc8o{nAPd|DkC_jaCXW` z+_lpaz6gZR-TQ2d;N;$q4)=NR*o#9R^3G!tcX*fgVD0jV_hIdE2mBwJ=&p5RS5&H& zClm}AM=i_gHuxLl70)xFjK7~2OI($fITu+6#LmV5Ht68zFx4-hK$C`QI%g{;N3dt@ zUR}OJzNZbFJDlz@^|r;_h9qu7=N?C=*@yt{Xuok90&>nV13YLvfMMeU?BpoXwXnyL z8|l6tfON2;D|(xJcm1f-5SPrOrt<+|z(2)6qUmgrk1XwBqlah#X0$>rYx+V%0b7F^ za@T&|vHvYvz6f*usC(2qikcqmH^3g7dn?M>9+Q8;iFNSbhrrXLes%pR7jF3jmDir} z@e~g*K}UU{c6ewdW-ph3sBJcVICF`2j@kSb?t^uPNj%1=jjFlT5m_BS{XZQxHlrrG zf>9Y&Bdf7*HU8W{T;}Ur81E4$P%4Az34;{N9snqK7fPetg(8Wh&`Vh*P%Gk*89c@I z1ypx@#$L(I{&UU-?s3}NZNLT}MYthd zF|~$IAGTC$o5@#!Z#qem6{%85GIs36jaTcLka0J;mshDgNb1FlzFm06f_~BPli3AT z$^}3m5r?lo@8?rJr@~{`HnZ);>2w@n4OD3g1BF-=dZai0_s6#c|X*H*z`hEqS}${4L*2kSNwu!C*7Bin4($Mg{t z{R=w?Sq@gzs0%#h7M_D6I4-@sdxdK}^n=0_%Gg}ZrDui0f&0mH=+g5lTmdbuXW!=9 z^Xo$X;3L)5?b^q->AkFyYD{&8@j4*<8`1CJ>+eUCWmC(^dNE&0HThDSm2KL|`N!wy zrmJhICl^oJTpKxQ*LK;cNoR}7oV9Bmh>+qo4^R*~fCn(ZX9UOeedxV&2%h~4p$YE( zI{!i?j9B`au|bk zxQiFC&rC5pa|}Toh-eqO9XiYs@hSFvMR%}CFdkvvcb5Fu*$H+MYB-LNCYN{NT0DQ& zG!+s#5_a9nr+S1&9X#WaJ`MqqS>~rZS;Q1*(|TB*-Re>e`vg?QD5@4h5~JM=*p*$>qc~y{4Cs?30Wg}i>CNI Jgwdg`k0%3=HN#a2IF^Gc(44TX@8G%BYjJFuI z{D35TZhlH>PHOQ=h9VZAIGFfltDljdo2p-#mYk8Ar|*(loL!P%pl@VgWK>+5o10ps u?-g8-SX7*Zza8-$ScdJGErMah$)yslldhhP+XJo7NeG*CdGp_ld)Aw?nfvb%Mh$aatxdh1^KgZ=;) zKFEhkhy&smIJC?-Me0a1e&ZR>%o{gHM{@#le{%Qnlo9d^ou|UVc?)VkfDuGcMJgIj z)Q+x(Dr<-`o@KpPsYF%8=S>wue)B1N}k|<*%x3knzI}ll`6vzGS#0l-Sg=0X44y-33f^K4{5E6%w7`hZOoP;9x?B3-raZ<;rT9liB zk#E4zj{Qm=SeW`1n0U?sIw~QArT264{QT_q{uWzX8wAJR%dbZd7$M*M&-y?gd5%_( z5Ryp9$rx~^IW5?j0g+4wQ#uY5Q!o!@_=QR?*~eha*F2F7&=W~GB!l=GGm(Pl_2>c&4ky){Yxx8>1)SQ~ewq=Me1$T78DF z*{jge!52*}E0m(nE3;92kGGYY!}w$*dUMa&>!&+RFX_-w_?Sjmz z!i7fYsdD^ON$#7UVMWIacde#4=I=8b@Ct4@Rjh@)7)FMy3 z1eu;buFj2r0})z%8$sY91ki&$*hX%{-!*74_`fvR|EC7 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 0000000000000000000000000000000000000000..bd3d11672fa53bdc266bbda20bf29de13c231328 GIT binary patch literal 238 zcmYe~<>g`kg2%r1l3GACh=UB6fE))PE>-~&DGVu$ISe5Xnh{DfL1|_n%@oX_$^4QL zs8o~j7NeG*Cd)0Z#N?99vc!_i{JdK%$@zI@w^;J>i*j$V7v$&XWag!>WGG?>ss|Il zT=g^Zb5r#z(~>h%^YmR(i?d7e3-pZ)jEss)b8}OR^u2-$5{rsc^$RLXGV=5E^YZj_ n^HWN5Qj7KD<1_OzOXB183My}L*g!<>KrSg}0}>oeK*$IHA^|g_S(snYDLk`E;5h7av#OuQvptA5Z>A=*YD@HFHp+5(~tzyNN)OFoL}nh(Lxc z?DJ5JWW*vFv)J)b8NCsf0R9|9bigcyluddTm+^tH3^Fz?@e|6Eqw9|za%_oG4zs6< zn;M=$)!O95h6Y4yQvrP66@4B)!t%wAmtOJd(b-4&iH&yKrqRZdU1PjtqpE8*TA9F; zvw_T&ny*{6YPMQ&Vf>b6{lMOWiJkC;%&GtF z;rEViv;^9%R~u0D7V3SYE4nyeEb3NDF!X*{iw-WRQgoP5Q|`RHbMI$fbvC&e^h@m8 T)-wkQ&Pax5L6W98#2NkrGpTlk literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..fe3e23840cd897fc1f87f4f8d6852673d3c7f5a2 GIT binary patch literal 1952 zcma(S%WfP+uzTJ!yX(h#6D7Q40U}vIUgVsRP(&1DqzHUTkkDL4eKJm-Lo{S} z>>BTjLTTeaXq%E7|0_|{VoZ%A5PO9v^w=?e*5rjWVWwpPW}xZcZ*r;l11?l8flp?w zm})qimnOJdi$}NE6g^yCy(Oa<5S*lxk_^GKRNmIH}CpSy@Cpq;JVacp{p$fgIm5WG~TU$^># z>m70l;!G9C+yn~JW!yxpD09rs1pBgX8)dw`B3CBSG8d|;6=x=pT(?q8WAJjm(02eE z-K%ZzugM8HrDueaj&>w-ItQ?;z0TGCi6h_bToH9H!psFt02I`3cYsdsPFDUubt5EDR?`V{Q@%A$|A5%bWitB2OXIQXY>Rk zTLC(ezXEoJ2N>dP!kdh&%{KOG_mew}@n`CZZ@U2+Cz zvkft{TGTkqKf|0iy6|TzZ(6ytPyb^-5-xZMZa1mSNXtnOp$_xr$Z^xDiU{ zVhSUoaZ=ZoDej2gXwmrIy$mCTvA9 zt7T<^Z+NkPpyddp;~tpsa&{5IOZ}ng(&;Q0##NdX_+gPWzWdpcP0mQ#V1Fz z=>Zp$&$v3&bu-!iXnR|=RmJ7x-eQm^&5=H+#YBinRkOC_YBGf>-fE8Ib?|~Ys`mhq zIH3Uzyw@l~Kt}+9GwfZXD^5b+a)uzg<{Ymso`9|MyTmu75KXJ)Rrto#!7nWr()fbj zg{^J^(qv9?K5W;)9L%u;eb~_1^AL8#=uZ*!hO;{kx zywJLq#+hsBqnT~%-a@)n6>2I&6i3#0(PlPw6EB!8AtW|{{17d|EXy$q6MSj=&_;eK gC7YbBMfAzy%5LKWp|JW99XSz=XdhtB0qh3OKk|6Z!vFvP literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..334f604af2efdfd6534708f61f9d0ab874ce3de4 GIT binary patch literal 1436 zcmb_cJCEEp5GJX&R_l8pK;Xo2m1dh;<7y+Y0|N$9T;sq2gam|6;(E0!Q8GydcL%9j z2d?rDTp0Gk80BNaYLN7|AjDWn%7i1Fg1%#K$$}QFaNZ=T`MR zm@(^WQ7u#io43?Poaa@hHRreF@2B%+Q;YfaB3}wQeRPDz5BpZk z+iSCIKHdhfpx1B6xc8X)p%=2|CdJyFmMRslg3ZMfwy7v^&cR5 zxg5dzlB>#Wt*Tcesy6+XZS-p}3`U*b%1%*iIQt9ic6q!g7-L&*=B z)3tTV;wlXtjcJtn29Mz6SPzniaYC3`e?M;fejM!=XHuSZrmXaS n$^*``aXJv2YJ~28b8=Ua4Il5P?K^zh@xKv583{uU*{S<4Zg)<< literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..9d8fc1ef6603f1528203c49db7ed910156793e24 GIT binary patch literal 465 zcmYjOO-sWt7*5iz?Y3_2;GgJak)A~q2ZEpo^Av=Vty%YxkCJp_`!{;@FXiCHlYhaJ zNh{(DdGdUO_vIyYG#UU`+uM&t4gvV|;(rBBT)Sgl9F(Ah90*7t^c-B#4CQBcl{!VU%9u|lYF&zX*W zo;Y|1dg8f 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 0b03769cb3973484b04158620f20ced8da864c08..6d461f471aaa307d63462d0fe985c2537e58df44 100644 GIT binary patch delta 96 zcmdnWc#M%Zk(ZZ?0SLD6-AmG%$SZ3v0_3DHq%h_%gg|IUAk7rapvnA_5h$g}c#BcX oPm}2uYhpn`PUS7O?9`&X)SQ(JMXW&CBG!pfVeG|BKqdzx00LGOmjD0& delta 79 zcmX@cxRsGNk(ZZ?0SHv9t|y63LKsq*f*CZKUorxPG#PI(YWV?4 bw(QiRywseP3`HzJ;UboanPIG-fh-09lMN9u 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 0000000000000000000000000000000000000000..722228ef3c95703596dba12957004aa040215cb0 GIT binary patch literal 601 zcmYjPy^ho{5cYp^uy+y^2p%B~Mq49v0z^S^NZg4cO_W=EceCp_k?nwjavd!^0TrU* z0eBe~*Iae4Km+3}#Z5Hd=W&uR<8M+QAKya|R`gIZi- zoYusUk^)o1T3&M4V;fuw%{5kB(_OYEPmtzHe5Wc?>xiD!W*x*@>%ej)b& zp7?Vrw_U5{xL$3wm9Mqmhpv~)XUnA@+EzRH_TZE3$FS|Jv{nwGFPMGGmRj@k zWXWUv2 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):