diff --git a/model_zoo/research/cv/TNT/eval.py b/model_zoo/research/cv/TNT/eval.py new file mode 100644 index 0000000000..2d09684cff --- /dev/null +++ b/model_zoo/research/cv/TNT/eval.py @@ -0,0 +1,75 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +""" +eval. +""" +import os +import argparse +from mindspore import context +from mindspore import nn +from mindspore.train.model import Model +from mindspore.train.serialization import load_checkpoint, load_param_into_net +from mindspore.common import dtype as mstype +from src.pet_dataset import create_dataset +from src.config import config_ascend, config_gpu +from src.tnt import tnt_b + + +parser = argparse.ArgumentParser(description='Image classification') +parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path') +parser.add_argument('--dataset_path', type=str, default=None, help='Dataset path') +parser.add_argument('--platform', type=str, default=None, help='run platform') +args_opt = parser.parse_args() + + +if __name__ == '__main__': + config_platform = None + if args_opt.platform == "Ascend": + config_platform = config_ascend + device_id = int(os.getenv('DEVICE_ID')) + context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", + device_id=device_id, save_graphs=False) + elif args_opt.platform == "GPU": + config_platform = config_gpu + context.set_context(mode=context.PYNATIVE_MODE, + device_target="GPU", save_graphs=False) + else: + raise ValueError("Unsupported platform.") + + loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') + + net = tnt_b(num_class=config_platform.num_classes) + + if args_opt.checkpoint_path: + param_dict = load_checkpoint(args_opt.checkpoint_path) + load_param_into_net(net, param_dict) + net.set_train(False) + + if args_opt.platform == "Ascend": + net.to_float(mstype.float16) + for _, cell in net.cells_and_names(): + if isinstance(cell, nn.Dense): + cell.to_float(mstype.float32) + + dataset = create_dataset(dataset_path=args_opt.dataset_path, + do_train=False, + config=config_platform, + platform=args_opt.platform, + batch_size=config_platform.batch_size) + step_size = dataset.get_dataset_size() + + model = Model(net, loss_fn=loss, metrics={'acc'}) + res = model.eval(dataset) + print("result:", res, "ckpt=", args_opt.checkpoint_path) diff --git a/model_zoo/research/cv/TNT/fig/tnt.PNG b/model_zoo/research/cv/TNT/fig/tnt.PNG new file mode 100644 index 0000000000..9cdac80e82 Binary files /dev/null and b/model_zoo/research/cv/TNT/fig/tnt.PNG differ diff --git a/model_zoo/research/cv/TNT/mindpsore_hub_conf.py b/model_zoo/research/cv/TNT/mindpsore_hub_conf.py new file mode 100644 index 0000000000..7ea7538c4f --- /dev/null +++ b/model_zoo/research/cv/TNT/mindpsore_hub_conf.py @@ -0,0 +1,22 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""hub config.""" +from src.tnt import tnt_b + + +def create_network(name, *args, **kwargs): + if name == 'TNT-B': + return tnt_b(*args, **kwargs) + raise NotImplementedError(f"{name} is not implemented in the repo") diff --git a/model_zoo/research/cv/TNT/readme.md b/model_zoo/research/cv/TNT/readme.md new file mode 100644 index 0000000000..16b3e5d549 --- /dev/null +++ b/model_zoo/research/cv/TNT/readme.md @@ -0,0 +1,128 @@ +# Contents + +- [TNT Description](#tnt-description) +- [Model Architecture](#model-architecture) +- [Dataset](#dataset) +- [Environment Requirements](#environment-requirements) +- [Script Description](#script-description) + - [Script and Sample Code](#script-and-sample-code) + - [Training Process](#training-process) + - [Evaluation Process](#evaluation-process) + - [Evaluation](#evaluation) +- [Model Description](#model-description) + - [Performance](#performance) + - [Training Performance](#evaluation-performance) + - [Inference Performance](#evaluation-performance) +- [Description of Random Situation](#description-of-random-situation) +- [ModelZoo Homepage](#modelzoo-homepage) + +## [TNT Description](#contents) + +The TNT (Transformer in Transformer) network is a pure transformer model for visual recognition. TNT treats an image as a sequence of patches and treats a patch as a sequence of pixels. TNT block utilizes a outer transformer block to process the sequence of patches and an inner transformer block to process the sequence of pixels. + +[Paper](https://arxiv.org/abs/2103.00112): Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, Yunhe Wang. Transformer in Transformer. preprint 2021. + +## [Model architecture](#contents) + +The overall network architecture of TNT is show below: +![](./fig/tnt.PNG) + +## [Dataset](#contents) + +Dataset used: [Oxford-IIIT Pet](https://www.robots.ox.ac.uk/~vgg/data/pets/) + +- Dataset size: 7049 colorful images in 1000 classes + - Train: 3680 images + - Test: 3369 images +- Data format: RGB images. + - Note: Data will be processed in src/dataset.py + +## [Environment Requirements](#contents) + +- Hardware(Ascend/GPU) + - Prepare hardware environment with Ascend or GPU. If you want to try Ascend, please send the [application form](https://obs-9be7.obs.cn-east-2.myhuaweicloud.com/file/other/Ascend%20Model%20Zoo%E4%BD%93%E9%AA%8C%E8%B5%84%E6%BA%90%E7%94%B3%E8%AF%B7%E8%A1%A8.docx) to ascend@huawei.com. Once approved, you can get the resources. +- Framework + - [MindSpore](https://www.mindspore.cn/install/en) +- For more information, please check the resources below£º + - [MindSpore Tutorials](https://www.mindspore.cn/tutorial/training/en/master/index.html) + - [MindSpore Python API](https://www.mindspore.cn/doc/api_python/en/master/index.html) + +## [Script description](#contents) + +### [Script and sample code](#contents) + +```python +TNT +├── eval.py # inference entry +├── fig +│ └── tnt.png # the illustration of TNT network +├── readme.md # Readme +└── src + ├── config.py # config of model and data + ├── pet_dataset.py # dataset loader + └── tnt.py # TNT network +``` + +## [Training process](#contents) + +To Be Done + +## [Eval process](#contents) + +### Usage + +After installing MindSpore via the official website, you can start evaluation as follows: + +### Launch + +```bash +# infer example + GPU: python eval.py --model tnt-b --dataset_path ~/Pets/test.mindrecord --platform GPU --checkpoint_path [CHECKPOINT_PATH] +``` + +> checkpoint can be downloaded at https://www.mindspore.cn/resources/hub. + +### Result + +```bash +result: {'acc': 0.95} ckpt= ./tnt-b-pets.ckpt +``` + +## [Model Description](#contents) + +### [Performance](#contents) + +#### Evaluation Performance + +##### TNT on ImageNet2012 + +| Parameters | | | +| -------------------------- | -------------------------------------- |---------------------------------- | +| Model Version | TNT-B |TNT-S| +| uploaded Date | 21/03/2021 (month/day/year) | 21/03/2021 (month/day/year) | +| MindSpore Version | 1.1 | 1.1 | +| Dataset | ImageNet2012 | ImageNet2012| +| Input size | 224x224 | 224x224| +| Parameters (M) | 86.4 | 23.8 | +| FLOPs (M) | 14.1 | 5.2 | +| Accuracy (Top1) | 82.8 | 81.3 | + +###### TNT on Oxford-IIIT Pet + +| Parameters | | | +| -------------------------- | -------------------------------------- |---------------------------------- | +| Model Version | TNT-B |TNT-S| +| uploaded Date | 21/03/2021 (month/day/year) | 21/03/2021 (month/day/year) | +| MindSpore Version | 1.1 | 1.1 | +| Dataset | Oxford-IIIT Pet | Oxford-IIIT Pet| +| Input size | 384x384 | 384x384| +| Parameters (M) | 86.4 | 23.8 | +| Accuracy (Top1) | 95.0 | 94.7 | + +## [Description of Random Situation](#contents) + +In dataset.py, we set the seed inside "create_dataset" function. We also use random seed in train.py. + +## [ModelZoo Homepage](#contents) + +Please check the official [homepage](https://gitee.com/mindspore/mindspore/tree/master/model_zoo). \ No newline at end of file diff --git a/model_zoo/research/cv/TNT/src/config.py b/model_zoo/research/cv/TNT/src/config.py new file mode 100644 index 0000000000..f51d63efeb --- /dev/null +++ b/model_zoo/research/cv/TNT/src/config.py @@ -0,0 +1,54 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +""" +network config setting, will be used in train.py and eval.py +""" +from easydict import EasyDict as ed + +config_ascend = ed({ + "num_classes": 37, + "image_height": 384, + "image_width": 384, + "batch_size": 50, + "epoch_size": 300, + "warmup_epochs": 5, + "lr": 1e-3, + "momentum": 0.9, + "weight_decay": 0.05, + "label_smooth": 0.1, + "loss_scale": 1024, + "save_checkpoint": True, + "save_checkpoint_epochs": 1, + "keep_checkpoint_max": 200, + "save_checkpoint_path": "./checkpoint", +}) + +config_gpu = ed({ + "num_classes": 37, + "image_height": 384, + "image_width": 384, + "batch_size": 50, + "epoch_size": 300, + "warmup_epochs": 5, + "lr": 1e-3, + "momentum": 0.9, + "weight_decay": 0.05, + "label_smooth": 0.1, + "loss_scale": 1024, + "save_checkpoint": True, + "save_checkpoint_epochs": 1, + "keep_checkpoint_max": 500, + "save_checkpoint_path": "./checkpoint", +}) diff --git a/model_zoo/research/cv/TNT/src/pet_dataset.py b/model_zoo/research/cv/TNT/src/pet_dataset.py new file mode 100644 index 0000000000..1ce4025008 --- /dev/null +++ b/model_zoo/research/cv/TNT/src/pet_dataset.py @@ -0,0 +1,97 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +""" +create train or eval dataset. +""" +import os +import mindspore.common.dtype as mstype +import mindspore.dataset.engine as de +import mindspore.dataset.transforms.py_transforms as py_transforms +import mindspore.dataset.transforms.c_transforms as c_transforms +import mindspore.dataset.vision.py_transforms as py_vision +from mindspore.dataset.vision import Inter + +def create_dataset(dataset_path, do_train, config, platform, repeat_num=1, batch_size=1): + """ + create a train or eval dataset + + Args: + dataset_path(string): the path of dataset. + do_train(bool): whether dataset is used for train or eval. + repeat_num(int): the repeat times of dataset. Default: 1 + batch_size(int): the batch size of dataset. Default: 32 + + Returns: + dataset + """ + if platform == "Ascend": + rank_size = int(os.getenv("RANK_SIZE")) + rank_id = int(os.getenv("RANK_ID")) + if rank_size == 1: + ds = de.MindDataset(dataset_path, num_parallel_workers=8, shuffle=True) + else: + ds = de.MindDataset(dataset_path, num_parallel_workers=8, shuffle=True, + num_shards=rank_size, shard_id=rank_id) + elif platform == "GPU": + if do_train: + from mindspore.communication.management import get_rank, get_group_size + ds = de.MindDataset(dataset_path, num_parallel_workers=8, shuffle=True, + num_shards=get_group_size(), shard_id=get_rank()) + else: + ds = de.MindDataset(dataset_path, num_parallel_workers=8, shuffle=False) + else: + raise ValueError("Unsupported platform.") + + resize_height = config.image_height + resize_width = config.image_width + buffer_size = 1000 + + # define map operations + random_resize_crop_bicubic = py_vision.RandomResizedCrop(size=(resize_height, resize_width), + scale=(0.08, 1.0), ratio=(3./4., 4./3.), + interpolation=Inter.BICUBIC) + random_horizontal_flip_op = py_vision.RandomHorizontalFlip(0.5) + color_jitter = 0.4 + adjust_range = (max(0, 1 - color_jitter), 1 + color_jitter) + random_color_jitter_op = py_vision.RandomColorAdjust(brightness=adjust_range, + contrast=adjust_range, + saturation=adjust_range) + + decode_p = py_vision.Decode() + resize_p = py_vision.Resize(int(resize_height), interpolation=Inter.BICUBIC) + center_crop_p = py_vision.CenterCrop(resize_height) + totensor = py_vision.ToTensor() + normalize_p = py_vision.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + + if do_train: + trans = py_transforms.Compose([decode_p, random_resize_crop_bicubic, random_horizontal_flip_op, + random_color_jitter_op, totensor, normalize_p]) + else: + trans = py_transforms.Compose([decode_p, resize_p, center_crop_p, totensor, normalize_p]) + + type_cast_op = c_transforms.TypeCast(mstype.int32) + + ds = ds.map(input_columns="image", operations=trans, num_parallel_workers=8) + ds = ds.map(input_columns="label_list", operations=type_cast_op, num_parallel_workers=8) + + # apply shuffle operations + ds = ds.shuffle(buffer_size=buffer_size) + + # apply batch operations + ds = ds.batch(batch_size, drop_remainder=True) + + # apply dataset repeat operation + ds = ds.repeat(repeat_num) + return ds diff --git a/model_zoo/research/cv/TNT/src/tnt.py b/model_zoo/research/cv/TNT/src/tnt.py new file mode 100644 index 0000000000..6ff7253977 --- /dev/null +++ b/model_zoo/research/cv/TNT/src/tnt.py @@ -0,0 +1,390 @@ +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""TNT""" +import math +import copy +import numpy as np +import mindspore.common.dtype as mstype +from mindspore import nn +from mindspore.ops import operations as P +from mindspore.common.tensor import Tensor +from mindspore.common.parameter import Parameter + + +class MLP(nn.Cell): + """MLP""" + + def __init__(self, in_features, hidden_features=None, out_features=None, dropout=0.): + super(MLP, self).__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Dense(in_features, hidden_features) + self.dropout = nn.Dropout(1. - dropout) + self.fc2 = nn.Dense(hidden_features, out_features) + self.act = nn.GELU() + + def construct(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class Attention(nn.Cell): + """Multi-head Attention""" + + def __init__(self, dim, hidden_dim=None, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.): + super(Attention, self).__init__() + hidden_dim = hidden_dim or dim + self.hidden_dim = hidden_dim + self.num_heads = num_heads + head_dim = hidden_dim // num_heads + self.head_dim = head_dim + self.scale = head_dim ** -0.5 + + self.qk = nn.Dense(dim, hidden_dim * 2, has_bias=qkv_bias) + self.v = nn.Dense(dim, hidden_dim, has_bias=qkv_bias) + self.softmax = nn.Softmax(axis=-1) + self.batmatmul_trans_b = P.BatchMatMul(transpose_b=True) + self.attn_drop = nn.Dropout(1. - attn_drop) + self.batmatmul = P.BatchMatMul() + self.proj = nn.Dense(hidden_dim, dim) + self.proj_drop = nn.Dropout(1. - proj_drop) + + self.transpose = P.Transpose() + self.reshape = P.Reshape() + + def construct(self, x): + """Multi-head Attention""" + B, N, _ = x.shape + qk = self.transpose(self.reshape(self.qk(x), (B, N, 2, self.num_heads, self.head_dim)), (2, 0, 3, 1, 4)) + q, k = qk[0], qk[1] + v = self.transpose(self.reshape(self.v(x), (B, N, self.num_heads, self.head_dim)), (0, 2, 1, 3)) + + attn = self.softmax(self.batmatmul_trans_b(q, k) * self.scale) + attn = self.attn_drop(attn) + x = self.reshape(self.transpose(self.batmatmul(attn, v), (0, 2, 1, 3)), (B, N, -1)) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class DropConnect(nn.Cell): + """drop connect implementation""" + + def __init__(self, drop_connect_rate=0., seed0=0, seed1=0): + super(DropConnect, self).__init__() + self.shape = P.Shape() + self.dtype = P.DType() + self.keep_prob = 1 - drop_connect_rate + self.dropout = P.Dropout(keep_prob=self.keep_prob) + self.keep_prob_tensor = Tensor(self.keep_prob, dtype=mstype.float32) + + def construct(self, x): + shape = self.shape(x) + dtype = self.dtype(x) + ones_tensor = P.Fill()(dtype, (shape[0], 1, 1, 1), 1) + _, mask = self.dropout(ones_tensor) + x = x * mask + x = x / self.keep_prob_tensor + return x + + +class Pixel2Patch(nn.Cell): + """Projecting Pixel Embedding to Patch Embedding""" + + def __init__(self, outer_dim): + super(Pixel2Patch, self).__init__() + self.norm_proj = nn.LayerNorm([outer_dim]) + self.proj = nn.Dense(outer_dim, outer_dim) + self.fake = Parameter(Tensor(np.zeros((1, 1, outer_dim)), + mstype.float32), name='fake', requires_grad=False) + self.reshape = P.Reshape() + self.tile = P.Tile() + self.concat = P.Concat(axis=1) + + def construct(self, pixel_embed, patch_embed): + B, N, _ = patch_embed.shape + proj = self.reshape(pixel_embed, (B, N - 1, -1)) + proj = self.proj(self.norm_proj(proj)) + proj = self.concat((self.tile(self.fake, (B, 1, 1)), proj)) + patch_embed = patch_embed + proj + return patch_embed + + +class TNTBlock(nn.Cell): + """TNT Block""" + + def __init__(self, inner_config, outer_config, dropout=0., attn_dropout=0., drop_connect=0.): + super().__init__() + # inner transformer + inner_dim = inner_config['dim'] + num_heads = inner_config['num_heads'] + mlp_ratio = inner_config['mlp_ratio'] + self.inner_norm1 = nn.LayerNorm([inner_dim]) + self.inner_attn = Attention(inner_dim, num_heads=num_heads, qkv_bias=True, attn_drop=attn_dropout, + proj_drop=dropout) + self.inner_norm2 = nn.LayerNorm([inner_dim]) + self.inner_mlp = MLP(inner_dim, int(inner_dim * mlp_ratio), dropout=dropout) + # outer transformer + outer_dim = outer_config['dim'] + num_heads = outer_config['num_heads'] + mlp_ratio = outer_config['mlp_ratio'] + self.outer_norm1 = nn.LayerNorm([outer_dim]) + self.outer_attn = Attention(outer_dim, num_heads=num_heads, qkv_bias=True, attn_drop=attn_dropout, + proj_drop=dropout) + self.outer_norm2 = nn.LayerNorm([outer_dim]) + self.outer_mlp = MLP(outer_dim, int(outer_dim * mlp_ratio), dropout=dropout) + # pixel2patch + self.pixel2patch = Pixel2Patch(outer_dim) + # assistant + self.drop_connect = DropConnect(drop_connect) + self.reshape = P.Reshape() + self.tile = P.Tile() + self.concat = P.Concat(axis=1) + + def construct(self, pixel_embed, patch_embed): + """TNT Block""" + pixel_embed = pixel_embed + self.inner_attn(self.inner_norm1(pixel_embed)) + pixel_embed = pixel_embed + self.inner_mlp(self.inner_norm2(pixel_embed)) + + patch_embed = self.pixel2patch(pixel_embed, patch_embed) + + patch_embed = patch_embed + self.outer_attn(self.outer_norm1(patch_embed)) + patch_embed = patch_embed + self.outer_mlp(self.outer_norm2(patch_embed)) + return pixel_embed, patch_embed + + +def _get_clones(module, N): + """get_clones""" + return nn.CellList([copy.deepcopy(module) for i in range(N)]) + + +class TNTEncoder(nn.Cell): + """TNT""" + + def __init__(self, encoder_layer, num_layers): + super().__init__() + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + + def construct(self, pixel_embed, patch_embed): + """TNT""" + for layer in self.layers: + pixel_embed, patch_embed = layer(pixel_embed, patch_embed) + return pixel_embed, patch_embed + + +class _stride_unfold_(nn.Cell): + """Unfold with stride""" + + def __init__( + self, kernel_size, stride=-1): + super(_stride_unfold_, self).__init__() + if stride == -1: + self.stride = kernel_size + else: + self.stride = stride + self.kernel_size = kernel_size + self.reshape = P.Reshape() + self.transpose = P.Transpose() + self.unfold = _unfold_(kernel_size) + + def construct(self, x): + """TNT""" + N, C, H, W = x.shape + leftup_idx_x = [] + leftup_idx_y = [] + nh = int((H - self.kernel_size) / self.stride + 1) + nw = int((W - self.kernel_size) / self.stride + 1) + for i in range(nh): + leftup_idx_x.append(i * self.stride) + for i in range(nw): + leftup_idx_y.append(i * self.stride) + NumBlock_x = len(leftup_idx_x) + NumBlock_y = len(leftup_idx_y) + zeroslike = P.ZerosLike() + cc_2 = P.Concat(axis=2) + cc_3 = P.Concat(axis=3) + unf_x = P.Zeros()((N, C, NumBlock_x * self.kernel_size, + NumBlock_y * self.kernel_size), mstype.float32) + N, C, H, W = unf_x.shape + for i in range(NumBlock_x): + for j in range(NumBlock_y): + unf_i = i * self.kernel_size + unf_j = j * self.kernel_size + org_i = leftup_idx_x[i] + org_j = leftup_idx_y[j] + fill = x[:, :, org_i:org_i + self.kernel_size, + org_j:org_j + self.kernel_size] + unf_x += cc_3((cc_3((zeroslike(unf_x[:, :, :, :unf_j]), + cc_2((cc_2((zeroslike(unf_x[:, :, :unf_i, unf_j:unf_j + self.kernel_size]), fill)), + zeroslike(unf_x[:, :, unf_i + self.kernel_size:, + unf_j:unf_j + self.kernel_size]))))), + zeroslike(unf_x[:, :, :, unf_j + self.kernel_size:]))) + y = self.unfold(unf_x) + return y + + +class _unfold_(nn.Cell): + """Unfold""" + + def __init__( + self, kernel_size, stride=-1): + super(_unfold_, self).__init__() + if stride == -1: + self.stride = kernel_size + self.kernel_size = kernel_size + + self.reshape = P.Reshape() + self.transpose = P.Transpose() + + def construct(self, x): + """TNT""" + N, C, H, W = x.shape + numH = int(H / self.kernel_size) + numW = int(W / self.kernel_size) + if numH * self.kernel_size != H or numW * self.kernel_size != W: + x = x[:, :, :numH * self.kernel_size, :, numW * self.kernel_size] + output_img = self.reshape(x, (N, C, numH, self.kernel_size, W)) + + output_img = self.transpose(output_img, (0, 1, 2, 4, 3)) + + output_img = self.reshape(output_img, (N, C, int( + numH * numW), self.kernel_size, self.kernel_size)) + + output_img = self.transpose(output_img, (0, 2, 1, 4, 3)) + + output_img = self.reshape(output_img, (N, int(numH * numW), -1)) + return output_img + + +class PixelEmbed(nn.Cell): + """Image to Pixel Embedding""" + + def __init__(self, img_size, patch_size=16, in_channels=3, embedding_dim=768, stride=4): + super(PixelEmbed, self).__init__() + self.num_patches = (img_size // patch_size) * (img_size // patch_size) + new_patch_size = math.ceil(patch_size / stride) + self.new_patch_size = new_patch_size + self.inner_dim = embedding_dim // new_patch_size // new_patch_size + self.proj = nn.Conv2d(in_channels, self.inner_dim, kernel_size=7, pad_mode='pad', + padding=3, stride=stride, has_bias=True) + self.unfold = _unfold_(kernel_size=new_patch_size) + self.reshape = P.Reshape() + self.transpose = P.Transpose() + + def construct(self, x): + B = x.shape[0] + x = self.proj(x) # B, C, H, W + x = self.unfold(x) # B, N, Ck2 + x = self.reshape(x, (B * self.num_patches, self.inner_dim, -1)) # B*N, C, M + x = self.transpose(x, (0, 2, 1)) # B*N, M, C + return x + + +class TNT(nn.Cell): + """TNT""" + + def __init__( + self, + img_size, + patch_size, + num_channels, + embedding_dim, + num_heads, + num_layers, + hidden_dim, + num_class, + stride=4, + dropout=0, + attn_dropout=0, + drop_connect=0.1 + ): + super(TNT, self).__init__() + + assert embedding_dim % num_heads == 0 + assert img_size % patch_size == 0 + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.patch_size = patch_size + self.num_channels = num_channels + self.img_size = img_size + self.num_patches = int((img_size // patch_size) ** 2) + new_patch_size = math.ceil(patch_size / stride) + inner_dim = embedding_dim // new_patch_size // new_patch_size + + self.patch_pos = Parameter(Tensor(np.random.rand(1, self.num_patches + 1, embedding_dim), + mstype.float32), name='patch_pos', requires_grad=True) + self.pixel_pos = Parameter(Tensor(np.random.rand(1, inner_dim, new_patch_size * new_patch_size), + mstype.float32), name='pixel_pos', requires_grad=True) + self.cls_token = Parameter(Tensor(np.random.rand(1, 1, embedding_dim), + mstype.float32), requires_grad=True) + self.patch_embed = Parameter(Tensor(np.zeros((1, self.num_patches, embedding_dim)), + mstype.float32), name='patch_embed', requires_grad=False) + self.fake = Parameter(Tensor(np.zeros((1, 1, embedding_dim)), + mstype.float32), name='fake', requires_grad=False) + self.pos_drop = nn.Dropout(1. - dropout) + + self.pixel_embed = PixelEmbed(img_size, patch_size, num_channels, embedding_dim, stride) + self.pixel2patch = Pixel2Patch(embedding_dim) + + inner_config = {'dim': inner_dim, 'num_heads': 4, 'mlp_ratio': 4} + outer_config = {'dim': embedding_dim, 'num_heads': num_heads, 'mlp_ratio': hidden_dim / embedding_dim} + encoder_layer = TNTBlock(inner_config, outer_config, dropout=dropout, attn_dropout=attn_dropout, + drop_connect=drop_connect) + self.encoder = TNTEncoder(encoder_layer, num_layers) + + self.head = nn.SequentialCell( + nn.LayerNorm([embedding_dim]), + nn.Dense(embedding_dim, num_class) + ) + + self.add = P.TensorAdd() + self.reshape = P.Reshape() + self.concat = P.Concat(axis=1) + self.tile = P.Tile() + self.transpose = P.Transpose() + + def construct(self, x): + """TNT""" + B, _, _, _ = x.shape + pixel_embed = self.pixel_embed(x) + pixel_embed = pixel_embed + self.transpose(self.pixel_pos, (0, 2, 1)) # B*N, M, C + + patch_embed = self.concat((self.cls_token, self.patch_embed)) + patch_embed = self.tile(patch_embed, (B, 1, 1)) + patch_embed = self.pos_drop(patch_embed + self.patch_pos) + + patch_embed = self.pixel2patch(pixel_embed, patch_embed) + + pixel_embed, patch_embed = self.encoder(pixel_embed, patch_embed) + + y = self.head(patch_embed[:, 0]) + return y + + +def tnt_b(num_class): + return TNT(img_size=384, + patch_size=16, + num_channels=3, + embedding_dim=640, + num_heads=10, + num_layers=12, + hidden_dim=640*4, + stride=4, + num_class=num_class)