|
- import mindspore.nn as nn
- import math
- import numpy as np
- import mindspore.ops as ops
- 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, depth, kernel_size=9, patch_size=7, n_class=1000):
- 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)],
- # AdaptiveAvgpool((1,1)),
- nn.AvgPool2d(kernel_size=temp),
- nn.Flatten(),
- nn.Dense(dim, n_class)
- )
-
-
- class AdaptiveAvgpool(nn.Cell):
- def __init__(self, target_size):
- super(AdaptiveAvgpool, self).__init__()
- self.target_size = target_size
-
- def construct(self, input):
- return AdaptiveAvgpool_function(input, self.target_size)
-
-
- 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
-
- for i in range(w):
- s1.append(np.floor(i * shape[-1] / w))
- e1.append(np.ceil((i + 1) * shape[-1] / w))
- for i in range(h):
- s2.append(np.floor(i * shape[-2] / h))
- e2.append(np.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, dim, depth, kernel_size, patch_size):
- super(convmixer_net, self).__init__()
- self.convmixer_layers = convmixer(dim, depth, kernel_size, patch_size)
-
- def construct(self, x):
- x = self.convmixer_layers(x)
- return x
-
-
- if __name__ == '__main__':
- my_net = convmixer_net(cfg.dim, cfg.depth, cfg.kernel_size, cfg.patch_size)
- print('net.py OK')
|