|
- import mindspore.nn as nn
- import math
- import mindspore.ops as ops
- import mindspore
- from mindspore import dtype
- from mindspore.dataset.transforms import py_transforms
- from config import convmixer_cfg as cfg
-
-
- class Residual(nn.Cell):
- #残差连接
- def __init__(self,fn):
- super(Residual, self).__init__()
- self.fn=fn
- def construct(self,x):
- return self.fn(x)+x
-
- def convmixer(dim=cfg.dim, depth = cfg.depth, kernel_size=cfg.kernel_size,
- patch_size=cfg.patch_size, n_class=cfg.num_classes):
- temp = ceil((cfg.image_height-patch_size+1)/patch_size)
- return nn.SequentialCell(nn.Conv2d(3,dim,patch_size,patch_size, pad_mode='valid'),nn.GELU(),nn.BatchNorm2d(dim),
- *[nn.SequentialCell(Residual(nn.SequentialCell(nn.Conv2d(dim,dim,kernel_size,group=dim),
- nn.GELU(),
- nn.BatchNorm2d(dim)
- )),
- nn.Conv2d(dim,dim,kernel_size=1),
- nn.GELU(),
- nn.BatchNorm2d(dim)
- ) for i in range(depth)],
- nn.AvgPool2d(kernel_size=(temp,temp)),
- nn.Flatten(),
- nn.Dense(dim,n_class)
- )
-
- class AdaptiveAvgpool(nn.Cell):
- def __init__(self):
- super(AdaptiveAvgpool, self).__init__()
- def construct(self, input):
- return AdaptiveAvgpool_function(input)
-
- def floor(x):
- return int(x)
-
- def ceil(x):
- out = int(x)
- if out == x:
- return x
- else:
- return out + 1
-
- def AdaptiveAvgpool_function(input,target_size=(1,1)):
- h=target_size[0]
- w=target_size[1]
- s1=[]
- e1=[]
- s2=[]
- e2=[]
- shape=input.shape
-
- return ops.ReduceMean(keep_dims=True)(input,(-2,-1))
- '''
- for i in range(w):
- s1.append(floor(i*shape[-1]/w))
- e1.append(ceil((i+1)*shape[-1]/w))
- for i in range(h):
- s2.append(floor(i*shape[-2]/h))
- e2.append(ceil((i+1)*shape[-2]/h))
- pool2=[]
- for i_h in range(h):
- pool=[]
- for i_w in range(w):
- pool.append(ops.ReduceMean(keep_dims=True)(input[:,:,s2[i_h]:e2[i_h],s1[i_w]:e1[i_w]],(-2,-1)))
- pool1=pool[0]
- for i_w in range(w-1):
- pool1=ops.Concat(axis=-1)((pool1,pool[i_w+1]))
- pool2.append(pool1)
- out=pool2[0]
- for i_h in range(h-1):
- out=ops.Concat(axis=-2)((out,pool2[i_h+1]))
- return out
- '''
-
-
-
-
- class convmixer_net(nn.Cell):
- def __init__(self,):
- super(convmixer_net,self).__init__()
- self.convmixer_layers = convmixer(dim=cfg.dim, depth=cfg.depth)
- def construct(self,x):
- x = self.convmixer_layers(x)
- return x
-
- if __name__ == '__main__':
- my_net = convmixer_net()
- print('net.py OK')
|