!1315 for second order codes

Merge pull request !1315 from zongha/master

example/resnet50_imagenet2012_THOR/config_imagenet.py

@@ -0,0 +1,45 @@
+# Copyright 2020 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 = ed({
+    "class_num": 1000,
+    "batch_size": 32,
+    "loss_scale": 128,
+    "momentum": 0.9,
+    "weight_decay": 5e-4,
+    "epoch_size": 50,
+    "buffer_size": 1000,
+    "image_height": 224,
+    "image_width": 224,
+    "save_checkpoint": True,
+    "save_checkpoint_steps": 5004,
+    "keep_checkpoint_max": 20,
+    "save_checkpoint_path": "./",
+    "lr_init": 0.01,
+    "lr_end": 0.00001,
+    "lr_max": 0.1,
+    "warmup_epochs": 0,
+    "lr_decay_mode": "cosine",
+    "label_smooth": 1,
+    "label_smooth_factor": 0.1,
+    "lr": 0.1,
+    "T_max": 90,
+    "eta_min": 0,
+    "frequency": 278
+})

example/resnet50_imagenet2012_THOR/crossentropy.py

@@ -0,0 +1,41 @@
+# Copyright 2020 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.
+# ============================================================================
+"""CrossEntropy"""
+import mindspore.nn as nn
+from mindspore import Tensor
+from mindspore.common import dtype as mstype
+from mindspore.nn.loss.loss import _Loss
+from mindspore.ops import functional as F
+from mindspore.ops import operations as P
+
+
+class CrossEntropy(_Loss):
+    """CrossEntropy"""
+    def __init__(self, smooth_factor=0., num_classes=1000):
+        super(CrossEntropy, self).__init__()
+        self.onehot = P.OneHot()
+        self.on_value = Tensor(1.0 - smooth_factor, mstype.float32)
+        self.off_value = Tensor(1.0 * smooth_factor / (num_classes - 1), mstype.float32)
+        # self.cast = P.Cast()
+        self.ce = nn.SoftmaxCrossEntropyWithLogits()
+        self.mean = P.ReduceMean(False)
+
+    def construct(self, logit, label):
+        # one_hot_label = self.onehot(self.cast(label, mstype.int32),
+        #                F.shape(logit)[1], self.on_value, self.off_value)、
+        one_hot_label = self.onehot(label, F.shape(logit)[1], self.on_value, self.off_value)
+        loss = self.ce(logit, one_hot_label)
+        loss = self.mean(loss, 0)
+        return loss

example/resnet50_imagenet2012_THOR/dataset_imagenet.py

@@ -0,0 +1,80 @@
+# Copyright 2020 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.c_transforms as C2
+import mindspore.dataset.transforms.vision.c_transforms as V_C
+
+
+def create_dataset(dataset_path, do_train, repeat_num=1, batch_size=32):
+    """
+    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
+    """
+
+    device_num = int(os.getenv("RANK_SIZE"))
+    rank_id = int(os.getenv("RANK_ID"))
+
+    if device_num == 1:
+        ds = de.ImageFolderDatasetV2(dataset_path, num_parallel_workers=8, shuffle=False)
+    else:
+        ds = de.ImageFolderDatasetV2(dataset_path, num_parallel_workers=8, shuffle=True,
+                                     num_shards=device_num, shard_id=rank_id)
+
+    image_size = 224
+    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
+    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
+    if do_train:
+        transform_img = [
+            V_C.RandomCropDecodeResize(image_size, scale=(0.08, 1.0), ratio=(0.75, 1.333)),
+            V_C.RandomHorizontalFlip(prob=0.5),
+            V_C.Normalize(mean=mean, std=std),
+            V_C.HWC2CHW()
+        ]
+    else:
+        transform_img = [
+            V_C.Decode(),
+            V_C.Resize((256, 256)),
+            V_C.CenterCrop(image_size),
+            V_C.Normalize(mean=mean, std=std),
+            V_C.HWC2CHW()
+        ]
+    # type_cast_op = C2.TypeCast(mstype.float16)
+    type_cast_op = C2.TypeCast(mstype.int32)
+
+    ds = ds.map(input_columns="image", operations=transform_img, num_parallel_workers=8)
+    ds = ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=8)
+
+    # apply shuffle operations
+    # ds = ds.shuffle(buffer_size=config.buffer_size)
+
+    # apply batch operations
+    ds = ds.batch(batch_size, drop_remainder=True)
+
+    # apply dataset repeat operation
+    ds = ds.repeat(repeat_num)
+
+    return ds

example/resnet50_imagenet2012_THOR/lr_generator.py

@@ -0,0 +1,126 @@
+# Copyright 2020 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.
+# ============================================================================
+"""learning rate generator"""
+import math
+
+import numpy as np
+
+
+def linear_warmup_lr(current_step, warmup_steps, base_lr, init_lr):
+    """linear_warmup_lr"""
+    lr_inc = (float(base_lr) - float(init_lr)) / float(warmup_steps)
+    lr = float(init_lr) + lr_inc * current_step
+    return lr
+
+
+def cosine_annealing_lr(lr, steps_per_epoch, warmup_epochs, max_epoch, T_max, eta_min=0, num_periods=0.5):
+    """linear_warmup_lr"""
+    base_lr = lr
+    warmup_init_lr = 0
+    total_steps = int(max_epoch * steps_per_epoch)
+    warmup_steps = int(warmup_epochs * steps_per_epoch)
+    decay_steps = total_steps - warmup_steps
+    lr_each_step = []
+    for i in range(total_steps):
+        if i < warmup_steps:
+            lr = linear_warmup_lr(i + 1, warmup_steps, base_lr, warmup_init_lr)
+        else:
+            # linear_decay = (total_steps - i) / decay_steps
+            cosine_decay = 0.5 * (1 + math.cos(math.pi * i / decay_steps))
+            decayed = cosine_decay
+            lr = base_lr * decayed
+        lr_each_step.append(lr)
+    return np.array(lr_each_step).astype(np.float32)
+
+
+def warmup_cosine_annealing_lr(lr, steps_per_epoch, warmup_epochs, max_epoch, T_max, eta_min=0, num_periods=0.5):
+    """warmup_cosine_annealing_lr"""
+    base_lr = lr
+    warmup_init_lr = 0
+    total_steps = int(max_epoch * steps_per_epoch * 0.99)
+    warmup_steps = int(warmup_epochs * steps_per_epoch)
+    decay_steps = total_steps - warmup_steps
+    lr_each_step = []
+    for i in range(total_steps):
+        if i < warmup_steps:
+            lr = linear_warmup_lr(i + 1, warmup_steps, base_lr, warmup_init_lr)
+        else:
+            linear_decay = (total_steps - i) / decay_steps
+            cosine_decay = 0.5 * (1 + math.cos(math.pi * 2 * num_periods * i / decay_steps))
+            decayed = linear_decay * cosine_decay
+            lr = base_lr * decayed + 0.000005
+        lr_each_step.append(lr)
+    return np.array(lr_each_step).astype(np.float32)
+
+
+def get_lr(global_step, lr_init, lr_end, lr_max, warmup_epochs, total_epochs, steps_per_epoch, lr_decay_mode):
+    """
+    generate learning rate array
+
+    Args:
+       global_step(int): total steps of the training
+       lr_init(float): init learning rate
+       lr_end(float): end learning rate
+       lr_max(float): max learning rate
+       warmup_epochs(int): number of warmup epochs
+       total_epochs(int): total epoch of training
+       steps_per_epoch(int): steps of one epoch
+       lr_decay_mode(string): learning rate decay mode, including steps, poly or default
+
+    Returns:
+       np.array, learning rate array
+    """
+    lr_each_step = []
+    total_steps = steps_per_epoch * total_epochs
+    warmup_steps = steps_per_epoch * warmup_epochs
+    if lr_decay_mode == 'steps':
+        decay_epoch_index = [0.3 * total_steps, 0.6 * total_steps, 0.8 * total_steps]
+        for i in range(total_steps):
+            if i < decay_epoch_index[0]:
+                lr = lr_max
+            elif i < decay_epoch_index[1]:
+                lr = lr_max * 0.1
+            elif i < decay_epoch_index[2]:
+                lr = lr_max * 0.01
+            else:
+                lr = lr_max * 0.001
+            lr_each_step.append(lr)
+    elif lr_decay_mode == 'poly':
+        if warmup_steps != 0:
+            inc_each_step = (float(lr_max) - float(lr_init)) / float(warmup_steps)
+        else:
+            inc_each_step = 0
+        for i in range(total_steps):
+            if i < warmup_steps:
+                lr = float(lr_init) + inc_each_step * float(i)
+            else:
+                base = (1.0 - (float(i) - float(warmup_steps)) / (float(total_steps) - float(warmup_steps)))
+                lr = float(lr_max) * base * base
+                if lr < 0.0:
+                    lr = 0.0
+            lr_each_step.append(lr)
+    else:
+        for i in range(total_steps):
+            if i < warmup_steps:
+                lr = lr_init + (lr_max - lr_init) * i / warmup_steps
+            else:
+                lr = lr_max - (lr_max - lr_end) * (i - warmup_steps) / (total_steps - warmup_steps)
+            lr_each_step.append(lr)
+
+    current_step = global_step
+    lr_each_step = np.array(lr_each_step).astype(np.float32)
+    learning_rate = lr_each_step[current_step:]
+
+    return learning_rate

example/resnet50_imagenet2012_THOR/model/dataset_helper.py

@@ -0,0 +1,191 @@
+# Copyright 2020 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.
+# ============================================================================
+"""Dataset help for minddata dataset"""
+from mindspore import context
+from mindspore._checkparam import check_bool
+from mindspore.nn.wrap import GetNextSingleOp
+from mindspore.parallel._utils import _get_device_num, _get_global_rank, _get_parallel_mode
+from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes, _to_tensor, \
+    _construct_tensor_list, _to_full_shapes, _to_full_tensor
+from mindspore.train.parallel_utils import ParallelMode
+
+
+class DatasetHelper:
+    """
+    Help function to use the Minddata dataset.
+
+    According to different context, change the iter of dataset, to use the same for loop in different context.
+
+    Note:
+        The iter of DatasetHelper will give one epoch data.
+
+    Args:
+        dataset (DataSet): The dataset.
+        dataset_sink_mode (bool): If true use GetNext to fetch the data, or else feed the data from host.
+            Default: True.
+
+    Examples:
+        >>> dataset_helper = DatasetHelper(dataset)
+        >>> for inputs in dataset_helper:
+        >>>     outputs = network(*inputs)
+    """
+
+    def __init__(self, dataset, first_order_iter=0, dataset_sink_mode=True):
+        check_bool(dataset_sink_mode)
+
+        iterclass = _DatasetIterGE
+        if not dataset_sink_mode:
+            iterclass = _DatasetIterFeed
+        elif not context.get_context("enable_ge"):
+            if context.get_context("enable_loop_sink"):
+                iterclass = _DatasetIterMSLoopSink
+            else:
+                iterclass = _DatasetIterMS
+
+        self.iter = iterclass(dataset, first_order_iter)
+
+    def __iter__(self):
+        return self.iter.__iter__()
+
+    # A temp solution for loop sink. Delete later
+    def types_shapes(self):
+        """Get the types and shapes from dataset on current config."""
+        return self.iter.types_shapes()
+
+    def loop_size(self):
+        """Get loop_size for every iteration."""
+        return self.iter.loop_size
+
+
+class _DatasetIter:
+    """Base iter for dataset help"""
+
+    def __init__(self, dataset):
+        self.loop_size = 1
+        if not hasattr(dataset, '__ME_INITED__'):
+            if not hasattr(dataset, '__loop_size__'):
+                self.loop_size = dataset.get_dataset_size()
+            else:
+                self.loop_size = dataset.__loop_size__
+            dataset.__ME_INITED__ = _exec_datagraph(dataset, self.loop_size).queue_name
+
+        self.ind = 0
+        self.dataset = dataset
+        dataset_types, dataset_shapes = _get_types_and_shapes(dataset)
+        self.dataset_types, self.dataset_shapes = dataset_types, dataset_shapes
+        # for self._parallel_mode equal to semi_auto_parallel or auto_parallel, use a complete tensor to
+        # compile, and slice tensor to run. The batch dimension of tensors for compile is device_number
+        # times the batch dimension of tensors for run
+        if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL):
+            device_num = _get_device_num()
+            self.dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
+
+    def __iter__(self):
+        self.ind = 0
+        return self
+
+    def __next__(self):
+        if self.ind >= self.loop_count:
+            raise StopIteration()
+        self.ind += 1
+        return self.op()
+
+    def types_shapes(self):
+        return self.dataset_types, self.dataset_shapes
+
+    def get_loop_count(self, dataset):
+        loop_count = 1
+        if hasattr(dataset, '__loop_size__'):
+            loop_size = dataset.__loop_size__
+            loop_count = int(dataset.get_dataset_size() / loop_size)
+        return loop_count
+
+
+class _DatasetIterMSLoopSink(_DatasetIter):
+    """Iter for context (enable_loop_sink=True)"""
+
+    def __init__(self, dataset, first_order_iter):
+        super(_DatasetIterMSLoopSink, self).__init__(dataset)
+        # self.loop_count = self.get_loop_count(dataset)
+        loop_size = dataset.__loop_size__ + first_order_iter
+        self.loop_count = int(dataset.get_dataset_size() / loop_size) * 2
+
+        def op():
+            return tuple()
+
+        self.op = op
+
+
+class _DatasetIterMS(_DatasetIter):
+    """Iter for context (enable_loop_sink=False)"""
+
+    def __init__(self, dataset, first_order_order):
+        super(_DatasetIterMS, self).__init__(dataset)
+        self.loop_count = dataset.get_dataset_size()
+        self.loop_size = 1
+        queue_name = dataset.__ME_INITED__
+        self.op = GetNextSingleOp(self.dataset_types, self.dataset_shapes, queue_name)
+
+
+class _DatasetIterGE(_DatasetIter):
+    """Iter for ge"""
+
+    def __init__(self, dataset):
+        super(_DatasetIterGE, self).__init__(dataset)
+        self.loop_count = self.get_loop_count(dataset)
+        parallel_mode = _get_parallel_mode()
+        self.need_to_full = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
+        batch_expand_num = 1
+        if self.need_to_full:
+            batch_expand_num = _get_device_num()
+        tensor_list_run = _construct_tensor_list(self.dataset_types, self.dataset_shapes, batch_expand_num)
+
+        def op():
+            return tensor_list_run
+
+        self.op = op
+
+
+class _DatasetIterFeed:
+    """Iter for feed data"""
+
+    def __init__(self, dataset, first_order_order):
+        self.dataset = dataset
+        self.device_num = _get_device_num()
+        self.global_rank = _get_global_rank()
+        self.repeat_count = dataset.get_repeat_count()
+        self.repeat_ind = 0
+        self.loop_count = dataset.get_dataset_size()
+        self.ind = 0
+
+        parallel_mode = context.get_auto_parallel_context("parallel_mode")
+        self.need_to_full = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
+
+    def __iter__(self):
+        if self.repeat_ind % self.repeat_count == 0:
+            self.iter = self.dataset.__iter__()
+
+        self.repeat_ind += 1
+        self.ind = 0
+        return self
+
+    def __next__(self):
+        if self.ind >= self.loop_count:
+            raise StopIteration()
+        self.ind += 1
+        data = self.iter.__next__()
+        if self.need_to_full:
+            return _to_full_tensor(data, self.device_num, self.global_rank)
+        return _to_tensor(data)

example/resnet50_imagenet2012_THOR/model/grad_reducer_thor.py

@@ -0,0 +1,183 @@
+# Copyright 2020 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.
+# ============================================================================
+"""grad_reducer_thor"""
+import mindspore.common.dtype as mstype
+from mindspore.communication.management import GlobalComm, get_group_size
+from mindspore.nn.cell import Cell
+from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore.ops.operations.comm_ops import AllReduce, ReduceOp
+
+reduce_opt = C.MultitypeFuncGraph("reduce_opt")
+
+_all_reduce_A = AllReduce()
+
+
+def _init_optimizer_allreduce(group):
+    global _all_reduce_A
+    _all_reduce_A = AllReduce(ReduceOp.SUM, GlobalComm.WORLD_COMM_GROUP)
+    _all_reduce_A.add_prim_attr('fusion', group)
+
+
+@reduce_opt.register("Function", "Number", "Tensor")
+def _tensors_allreduce_mean(mul, degree, grad):
+    degree = F.scalar_cast(degree, F.dtype(grad))
+    grad = _all_reduce_A(grad)
+    cast_op = P.Cast()
+    return mul(grad, cast_op(F.scalar_to_array(1.0 / degree), F.dtype(grad)))
+
+
+@reduce_opt.register("Bool", "Tensor")
+def _tensors_allreduce(allreduce_filter, grad):
+    if allreduce_filter:
+        return _all_reduce_A(grad)
+    return grad
+
+
+_get_datatype = C.MultitypeFuncGraph("_get_datatype")
+
+
+@_get_datatype.register("Tensor")
+def _tensors_get_datatype(grad):
+    """
+    Acquire gradient datatype.
+
+    Args:
+        grad (Tensor): The gradient tensor before operation.
+
+    Returns:
+        mstype, the datatype of gradient.
+    """
+    return F.dtype(grad)
+
+
+_cast_datatype = C.MultitypeFuncGraph("_cast_datatype")
+
+
+@_cast_datatype.register("TypeType", "Tensor")
+def _tensors_cast_datatype(datatype, grad):
+    """
+    Cast gradient to datatype.
+
+    Args:
+        datatype (mstype): the destination datatype of gradient.
+        grad (Tensor): The gradient tensor before operation.
+
+    Returns:
+        Tensor, the gradient tensor after operation.
+    """
+    return F.cast(grad, datatype)
+
+
+class DistributedGradReducerThor(Cell):
+    """
+    A distributed optimizer.
+
+    Constructs a gradient reducer Cell, which applies communication and average operations on
+    single-process gradient values.
+
+    Args:
+        parameters (list): the parameters to be updated.
+        mean (bool): When mean is true, the mean coefficient (degree) would apply on gradients. Default: False.
+        degree (int): The mean coefficient. Usually it equals to device number. Default: None.
+
+    Raises:
+        ValueError: If degree is not a int or less than 0.
+
+    Examples:
+        >>> from mindspore.communication import init, get_group_size
+        >>> from mindspore.ops import composite as C
+        >>> from mindspore.ops import operations as P
+        >>> from mindspore.ops import functional as F
+        >>> from mindspore import context
+        >>> from mindspore import nn
+        >>> from mindspore import ParallelMode, ParameterTuple
+        >>>
+        >>> device_id = int(os.environ["DEVICE_ID"])
+        >>> context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=True,
+        >>>                     device_id=int(device_id), enable_hccl=True)
+        >>> init()
+        >>> context.reset_auto_parallel_context()
+        >>> context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL)
+        >>>
+        >>>
+        >>> class TrainingWrapper(nn.Cell):
+        >>>     def __init__(self, network, optimizer, sens=1.0):
+        >>>         super(TrainingWrapper, self).__init__(auto_prefix=False)
+        >>>         self.network = network
+        >>>         self.network.add_flags(defer_inline=True)
+        >>>         self.weights = ParameterTuple(network.trainable_params())
+        >>>         self.optimizer = optimizer
+        >>>         self.grad = C.GradOperation('grad', get_by_list=True, sens_param=True)
+        >>>         self.sens = sens
+        >>>         self.reducer_flag = False
+        >>>         self.grad_reducer = None
+        >>>         self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
+        >>>         if self.parallel_mode in [ParallelMode.DATA_PARALLEL,
+        >>>                                            ParallelMode.HYBRID_PARALLEL]:
+        >>>             self.reducer_flag = True
+        >>>         if self.reducer_flag:
+        >>>             mean = context.get_auto_parallel_context("mirror_mean")
+        >>>             if mean.get_device_num_is_set():
+        >>>                 degree = context.get_auto_parallel_context("device_num")
+        >>>             else:
+        >>>                 degree = get_group_size()
+        >>>             self.grad_reducer = nn.DistributedGradReducer(optimizer.parameters, mean, degree)
+        >>>
+        >>>     def construct(self, *args):
+        >>>         weights = self.weights
+        >>>         loss = self.network(*args)
+        >>>         sens = P.Fill()(P.DType()(loss), P.Shape()(loss), self.sens)
+        >>>         grads = self.grad(self.network, weights)(*args, sens)
+        >>>         if self.reducer_flag:
+        >>>             # apply grad reducer on grads
+        >>>             grads = self.grad_reducer(grads)
+        >>>         return F.depend(loss, self.optimizer(grads))
+        >>>
+        >>> network = Net()
+        >>> optimizer = nn.Momentum(network.trainable_params(), learning_rate=0.1, momentum=0.9)
+        >>> train_cell = TrainingWrapper(network, optimizer)
+        >>> inputs = Tensor(np.ones([16, 16]).astype(np.float32))
+        >>> label = Tensor(np.zeros([16, 16]).astype(np.float32))
+        >>> grads = train_cell(inputs, label)
+    """
+
+    def __init__(self, parameters, group, mean=True, degree=None):
+        super(DistributedGradReducerThor, self).__init__(auto_prefix=False)
+        self.hyper_map = C.HyperMap()
+        self.mul = P.Mul()
+        if degree is None:
+            self.degree = get_group_size()
+        else:
+            if not isinstance(degree, int) or degree <= 0:
+                raise ValueError("Parameter 'degree' in DistributedGradReducer should large than 0 and be int")
+            self.degree = degree
+        self.mean = mean
+        self.allreduce_filter = tuple(x.layerwise_parallel is False for x in parameters)
+        _init_optimizer_allreduce(group)
+
+    def construct(self, grads):
+        # In some circumstances, the data precision of grads could be mixed with float16 and float32. Thus, the
+        # result of AllReduce is unreliable. To solve the problem, grads should be cast to float32 before AllReduce,
+        # and cast back after the operation.
+        datatypes = self.hyper_map(F.partial(_get_datatype), grads)
+        grads = self.hyper_map(F.partial(_cast_datatype, mstype.float32), grads)
+
+        if self.mean:
+            new_grad = self.hyper_map(F.partial(reduce_opt, self.mul, self.degree), grads)
+        else:
+            new_grad = self.hyper_map(F.partial(reduce_opt), self.allreduce_filter, grads)
+
+        new_grad = self.hyper_map(F.partial(_cast_datatype), datatypes, new_grad)
+        return new_grad

example/resnet50_imagenet2012_THOR/model/model_thor.py

@@ -0,0 +1,630 @@
+# Copyright 2020 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.
+# ============================================================================
+"""Model."""
+from mindspore import context
+from mindspore import log as logger
+from mindspore._c_expression import init_exec_dataset
+from mindspore._checkparam import check_input_data, check_output_data, check_int_positive, check_bool
+from mindspore.common import dtype as mstype
+from mindspore.common.dtype import pytype_to_dtype
+from mindspore.common.tensor import Tensor
+from mindspore.nn.metrics import Loss
+from mindspore.nn.metrics import get_metrics
+from mindspore.nn.wrap.cell_wrapper import _VirtualDatasetCell
+from mindspore.parallel._utils import _get_parallel_mode, _get_device_num, _get_global_rank, \
+    _get_parameter_broadcast, _device_number_check, _parameter_broadcast_check
+from mindspore.train import amp
+from mindspore.train.callback import _InternalCallbackParam, RunContext, _build_callbacks
+from mindspore.train.parallel_utils import ParallelMode
+import mindspore.nn as nn
+from second_order.dataset_helper import DatasetHelper
+import numpy as np
+
+def _convert_type(types):
+    """
+    Convert from numpy type to tensor type.
+
+    Args:
+        types (list): Numpy type list of element in dataset.
+
+    Returns:
+        list, list of element in dataset.
+    """
+    ms_types = []
+    for np_type in types:
+        ms_type = pytype_to_dtype(np_type)
+        ms_types.append(ms_type)
+    return ms_types
+
+
+def _get_types_and_shapes(dataset):
+    """Get dataset types and shapes."""
+    dataset_types = _convert_type(dataset.output_types())
+    dataset_shapes = dataset.output_shapes()
+    return dataset_types, dataset_shapes
+
+
+def _exec_datagraph(exec_dataset, dataset_size, phase='dataset'):
+    """Initialize and execute the dataset graph."""
+    batch_size = exec_dataset.get_batch_size()
+    input_indexs = exec_dataset.input_indexs
+
+    # transform data format
+    dataset_types, dataset_shapes = _get_types_and_shapes(exec_dataset)
+    init_exec_dataset(exec_dataset.__ME_INITED__,
+                      dataset_size,
+                      batch_size,
+                      dataset_types,
+                      dataset_shapes,
+                      input_indexs,
+                      phase=phase)
+
+
+class Model:
+    """
+    High-Level API for Training or Testing.
+
+    `Model` groups layers into an object with training and inference features.
+
+    Args:
+        network (Cell): The training or testing network.
+        loss_fn (Cell): Objective function, if loss_fn is None, the
+                             network should contain the logic of loss and grads calculation, and the logic
+                             of parallel if needed. Default: None.
+        optimizer (Cell): Optimizer for updating the weights. Default: None.
+        metrics (Union[dict, set]): Dict or set of metrics to be evaluated by the model during
+                        training and testing. eg: {'accuracy', 'recall'}. Default: None.
+        eval_network (Cell): Network for evaluation. If not defined, `network` and `loss_fn` would be wrapped as
+                             `eval_network`. Default: None.
+        eval_indexes (list): In case of defining the `eval_network`, if `eval_indexes` is None, all outputs of
+                             `eval_network` would be passed to metrics, otherwise `eval_indexes` must contain three
+                             elements, representing the positions of loss value, predict value and label, the loss
+                             value would be passed to `Loss` metric, predict value and label would be passed to other
+                             metric. Default: None.
+        amp_level (str): Option for argument `level` in `mindspore.amp.build_train_network`, level for mixed
+            precision training. Supports [O0, O2]. Default: "O0".
+
+            - O0: Do not change.
+            - O2: Cast network to float16, keep batchnorm run in float32, using dynamic loss scale.
+
+        loss_scale_manager (Union[None, LossScaleManager]): If None, not scale the loss, or else
+            scale the loss by LossScaleManager. If it is set, overwrite the level setting. It's a eyword argument.
+            e.g. Use `loss_scale_manager=None` to set the value.
+        keep_batchnorm_fp32 (bool): Keep Batchnorm run in `float32`. If set, overwrite the level setting. Default: True.
+
+    Examples:
+        >>> class Net(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.conv = nn.Conv2d(3, 64, 3, has_bias=False, weight_init='normal')
+        >>>         self.bn = nn.BatchNorm2d(64)
+        >>>         self.relu = nn.ReLU()
+        >>>         self.flatten = nn.Flatten()
+        >>>         self.fc = nn.Dense(64*224*224, 12) # padding=0
+        >>>
+        >>>     def construct(self, x):
+        >>>         x = self.conv(x)
+        >>>         x = self.bn(x)
+        >>>         x = self.relu(x)
+        >>>         x = self.flatten(x)
+        >>>         out = self.fc(x)
+        >>>         return out
+        >>>
+        >>> net = Net()
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+        >>> optim = Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
+        >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics=None)
+        >>> dataset = get_dataset()
+        >>> model.train(2, dataset)
+    """
+
+    def __init__(self, network, loss_fn=None, optimizer=None, metrics=None, eval_network=None,
+                 eval_indexes=None, amp_level="O0", frequency=278, **kwargs):
+        self._network = network
+        self._loss_fn = loss_fn
+        self._optimizer = optimizer
+        self._loss_scale_manager = None
+        self._loss_scale_manager_set = False
+        self._keep_bn_fp32 = True
+        self._frequency = frequency
+        self._check_kwargs(kwargs)
+        if 'keep_batchnorm_fp32' in kwargs:
+            self._keep_bn_fp32 = kwargs['keep_batchnorm_fp32']
+        if 'loss_scale_manager' in kwargs:
+            self._loss_scale_manager = kwargs['loss_scale_manager']
+            self._loss_scale_manager_set = True
+        self._amp_level = amp_level
+        self._parallel_mode = _get_parallel_mode()
+        self._device_number = _get_device_num()
+        self._global_rank = _get_global_rank()
+        self._parameter_broadcast = _get_parameter_broadcast()
+
+        self._train_network = self._build_train_network()
+        self._build_eval_network(metrics, eval_network, eval_indexes)
+        self._build_predict_network()
+
+    def _check_kwargs(self, kwargs):
+        for arg in kwargs:
+            if arg not in ['loss_scale_manager', 'keep_batchnorm_fp32']:
+                raise ValueError(f"Unsupport arg '{arg}'")
+
+    def _build_train_network(self):
+        """Build train network"""
+        network = self._network
+        if self._optimizer:
+            if self._loss_scale_manager_set:
+                network = amp.build_train_network(network,
+                                                  self._optimizer,
+                                                  self._loss_fn,
+                                                  level=self._amp_level,
+                                                  loss_scale_manager=self._loss_scale_manager,
+                                                  keep_batchnorm_fp32=self._keep_bn_fp32)
+            else:
+                network = amp.build_train_network(network,
+                                                  self._optimizer,
+                                                  self._loss_fn,
+                                                  level=self._amp_level,
+                                                  keep_batchnorm_fp32=self._keep_bn_fp32)
+        elif self._loss_fn:
+            network = nn.WithLossCell(network, self._loss_fn)
+        # If need to check if loss_fn is not None, but optimizer is None
+        return network
+
+    def _build_eval_network(self, metrics, eval_network, eval_indexes):
+        """Build the network for evaluation."""
+        self._metric_fns = get_metrics(metrics)
+        if not self._metric_fns:
+            return
+
+        if eval_network is not None:
+            if eval_indexes is not None and not (isinstance(eval_indexes, list) and len(eval_indexes) == 3):
+                raise ValueError("Eval_indexes must be a list or None. If eval_indexes is a list, length of it \
+                                 must be three. But got {}".format(eval_indexes))
+
+            self._eval_network = eval_network
+            self._eval_indexes = eval_indexes
+        else:
+            if self._loss_fn is None:
+                raise ValueError("loss_fn can not be None.")
+            self._eval_network = nn.WithEvalCell(self._network, self._loss_fn)
+            self._eval_indexes = [0, 1, 2]
+
+    def _build_predict_network(self):
+        """Build the network for prediction."""
+        self._predict_network = self._network
+        if self._parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL):
+            self._predict_network = _VirtualDatasetCell(self._network)
+
+    def _clear_metrics(self):
+        """Clear metrics local values."""
+        for metric in self._metric_fns.values():
+            metric.clear()
+
+    def _update_metrics(self, outputs):
+        """Update metrics local values."""
+        if not isinstance(outputs, tuple):
+            raise ValueError("The `outputs` is not tuple.")
+
+        if self._eval_indexes is not None and len(outputs) < 3:
+            raise ValueError("The length of `outputs` must be greater than or equal to 3, \
+                             but got {}".format(len(outputs)))
+
+        for metric in self._metric_fns.values():
+            if self._eval_indexes is None:
+                metric.update(*outputs)
+            else:
+                if isinstance(metric, Loss):
+                    metric.update(outputs[self._eval_indexes[0]])
+                else:
+                    metric.update(outputs[self._eval_indexes[1]], outputs[self._eval_indexes[2]])
+
+    def _get_metrics(self):
+        """Get metrics local values."""
+        metrics = dict()
+        for key, value in self._metric_fns.items():
+            metrics[key] = value.eval()
+        return metrics
+
+    def _get_scaling_sens(self):
+        """get the scaling sens"""
+        scaling_sens = 1
+        if self._loss_scale_manager is not None:
+            scaling_sens = self._loss_scale_manager.get_loss_scale()
+        if self._parallel_mode == ParallelMode.DATA_PARALLEL:
+            scaling_sens /= self._device_number
+        return scaling_sens
+
+    def _train(self, epoch, train_dataset, callbacks=None, dataset_sink_mode=True):
+        """
+        Training.
+
+        Args:
+            epoch (int): Total number of iterations on the data.
+            train_dataset (Dataset): A training dataset iterator. If there is no
+                                     loss_fn, a tuple with multiply data (data1, data2, data3, ...) will be
+                                     returned and passed to the network. Otherwise, a tuple (data, label) will
+                                     be returned, and the data and label are passed to the network and loss
+                                     function respectively.
+            callbacks (list): List of callback object. Callbacks which should be executed while training. Default: None.
+            dataset_sink_mode (bool): Determines whether to pass the data through dataset channel. Default: True.
+                                      Configure pynative mode, the training process will be performed with
+                                      dataset not sink.
+        """
+        epoch = check_int_positive(epoch)
+        self._train_network.set_train()
+
+        if self._parameter_broadcast:
+            self._train_network.set_broadcast_flag()
+
+        # build callback list
+        list_callback = _build_callbacks(callbacks)
+        cb_params = _InternalCallbackParam()
+        cb_params.train_network = self._train_network
+        cb_params.epoch_num = epoch
+        cb_params.batch_num = train_dataset.get_dataset_size()
+        cb_params.mode = "train"
+        cb_params.loss_fn = self._loss_fn
+        cb_params.optimizer = self._optimizer
+        cb_params.parallel_mode = self._parallel_mode
+        cb_params.device_number = self._device_number
+        cb_params.train_dataset = train_dataset
+        cb_params.list_callback = list_callback
+
+        if dataset_sink_mode:
+            if context.get_context("mode") == context.PYNATIVE_MODE:
+                logger.warning("The pynative mode cannot support dataset sink mode currently."
+                               "So the training process will be performed with dataset not sink.")
+                self._train_process(epoch, train_dataset, list_callback, cb_params)
+            else:
+                self._train_dataset_sink_process(epoch, train_dataset, list_callback, cb_params)
+        else:
+            self._train_process(epoch, train_dataset, list_callback, cb_params)
+
+    def _train_dataset_sink_process(self, epoch, train_dataset, list_callback=None, cb_params=None):
+        """
+        Training process. The data would be passed to network through dataset channel.
+
+        Args:
+            epoch (int): Total number of iterations on the data.
+            train_dataset (Dataset): A training dataset iterator. If there is no
+                                     loss_fn, a tuple with multiply data (data1, data2, data3, ...) should be
+                                     returned and passed to the network. Otherwise, a tuple (data, label) should
+                                     be returned, and the data and label are passed to the network and loss
+                                     function respectively.
+            list_callback (_ListCallback): Executor of callback list. Default: None.
+            cb_params (_InternalCallbackParam): Callback parameters. Default: None.
+        """
+        # remove later to deal with loop sink
+        iter_first_order = 277
+        iter_second_order = 1
+        train_dataset.__loop_size__ = iter_second_order
+        need_wrap = False
+        if not hasattr(train_dataset, '__ME_INITED__') and context.get_context("enable_loop_sink") \
+                and not context.get_context("enable_ge"):
+            need_wrap = True
+
+        dataset_helper = DatasetHelper(train_dataset, iter_first_order)
+        # remove later to deal with loop sink
+        if need_wrap:
+            self._train_network = nn.DataWrapper(self._train_network, *(dataset_helper.types_shapes()),
+                                                 train_dataset.__ME_INITED__)
+            cb_params.train_network = self._train_network
+            self._train_network.set_train()
+
+        cb_params.cur_step_num = 0
+        loop_size = dataset_helper.loop_size()
+        run_context = RunContext(cb_params)
+        list_callback.begin(run_context)
+
+        # used to stop training for early stop, such as stopAtTIme or stopATStep
+        should_stop = False
+        has_do_train1_dataset = False
+        checkpoint_branch_one = True
+        for i in range(epoch):
+            cb_params.cur_epoch_num = i + 1
+            list_callback.epoch_begin(run_context)
+
+            # for data sink dataset_helper only iter once, other wise iter epoch_size times.
+            for inputs in dataset_helper:
+                list_callback.step_begin(run_context)
+                if checkpoint_branch_one:
+                    cb_params.cur_step_num += loop_size
+                    self._train_network.set_second_order(True)
+                    self._train_network.phase = 'train0'
+                else:
+                    cb_params.cur_step_num += iter_first_order
+                    self._train_network.set_second_order(False)
+                    self._train_network.phase = 'train1'
+                    if not has_do_train1_dataset:
+                        _exec_datagraph(train_dataset, iter_first_order, phase='train1_dataset')
+                        has_do_train1_dataset = True
+                checkpoint_branch_one = not checkpoint_branch_one
+                outputs = self._train_network(*inputs)
+                cb_params.net_outputs = outputs
+                list_callback.step_end(run_context)
+
+            list_callback.epoch_end(run_context)
+            should_stop = should_stop or run_context.get_stop_requested()
+            if should_stop:
+                break
+
+        list_callback.end(run_context)
+
+    def _train_process(self, epoch, train_dataset, list_callback=None, cb_params=None):
+        """
+        Training process. The data would be passed to network directly.
+
+        Args:
+            epoch (int): Total number of iterations on the data.
+            train_dataset (Dataset): A training dataset iterator. If there is no
+                                     loss_fn, a tuple with multiply data (data1, data2, data3, ...) should be
+                                     returned and passed to the network. Otherwise, a tuple (data, label) should
+                                     be returned, and the data and label are passed to the network and loss
+                                     function respectively.
+            list_callback (_ListCallback): Executor of callback list. Default: None.
+            cb_params (_InternalCallbackParam): Callback parameters. Default: None.
+        """
+        dataset_helper = DatasetHelper(train_dataset, dataset_sink_mode=False)
+        cb_params.cur_step_num = 0
+        run_context = RunContext(cb_params)
+        _callback_wrapper(list_callback, run_context, "begin")
+        # used to stop training for early stop, such as stopAtTIme or stopATStep
+        should_stop = False
+
+        for i in range(epoch):
+            cb_params.cur_epoch_num = i + 1
+
+            _callback_wrapper(list_callback, run_context, "epoch_begin")
+
+            for next_element in dataset_helper:
+                len_element = len(next_element)
+                if self._loss_fn and len_element != 2:
+                    raise ValueError("when loss_fn is not None, train_dataset should"
+                                     "return two elements, but got {}".format(len_element))
+                cb_params.cur_step_num += 1
+                _callback_wrapper(list_callback, run_context, "step_begin")
+
+                overflow = False
+                if self._loss_scale_manager and self._loss_scale_manager.get_drop_overflow_update():
+                    scaling_sens = self._get_scaling_sens()
+                    next_element = tuple(next_element) + (Tensor(scaling_sens, mstype.float32),)
+
+                outputs = self._train_network(*next_element)
+                cb_params.net_outputs = outputs
+                if self._loss_scale_manager and self._loss_scale_manager.get_drop_overflow_update():
+                    _, overflow, _ = outputs
+                    overflow = np.all(overflow.asnumpy())
+                    self._loss_scale_manager.update_loss_scale(overflow)
+
+                _callback_wrapper(list_callback, run_context, "step_end")
+                should_stop = should_stop or run_context.get_stop_requested()
+                if should_stop:
+                    break
+
+            train_dataset.reset()
+
+            _callback_wrapper(list_callback, run_context, "epoch_end")
+            should_stop = should_stop or run_context.get_stop_requested()
+            if should_stop:
+                break
+
+        _callback_wrapper(list_callback, run_context, "end")
+
+    def train(self, epoch, train_dataset, callbacks=None, dataset_sink_mode=True):
+        """
+        Training API where the iteration is controlled by python front-end.
+
+        When setting pynative mode, the training process will be performed with dataset not sink.
+
+        Note:
+            CPU is not supported when dataset_sink_mode is true.
+            If dataset_sink_mode is True, epoch of training should be equal to the count of repeat
+            operation in dataset processing. Otherwise, errors could occur since the amount of data
+            is not the amount training requires.
+            If dataset_sink_mode is True, data will be sent to device. If device is Ascend, features
+            of data will be transferred one by one. The limitation of data transmission per time is 256M.
+
+        Args:
+            epoch (int): Total number of iterations on the data.
+            train_dataset (Dataset): A training dataset iterator. If there is no
+                                     loss_fn, a tuple with multiply data (data1, data2, data3, ...) should be
+                                     returned and passed to the network. Otherwise, a tuple (data, label) should
+                                     be returned, and the data and label are passed to the network and loss
+                                     function respectively.
+            callbacks (list): List of callback object. Callbacks which should be excuted while training. Default: None.
+            dataset_sink_mode (bool): Determines whether to pass the data through dataset channel. Default: True.
+                                      Configure pynative mode, the training process will be performed with
+                                      dataset not sink.
+
+
+        Examples:
+            >>> dataset = get_dataset()
+            >>> net = Net()
+            >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+            >>> loss_scale_manager = FixedLossScaleManager()
+            >>> optim = Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
+            >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics=None, loss_scale_manager=loss_scale_manager)
+            >>> model.train(2, dataset)
+        """
+        repeat_count = train_dataset.get_repeat_count()
+        if epoch != repeat_count and dataset_sink_mode is True:
+            logger.warning(f"The epoch_size {epoch} is not the same with dataset repeat_count {repeat_count}")
+        check_bool(dataset_sink_mode)
+        _device_number_check(self._parallel_mode, self._device_number)
+        _parameter_broadcast_check(self._parallel_mode, self._parameter_broadcast)
+
+        if context.get_context("device_target") in ["CPU", "GPU"] and context.get_context("enable_loop_sink"):
+            raise ValueError("CPU and GPU can't support loop sink, please set enable_loop_sink=False.")
+
+        self._train(epoch,
+                    train_dataset,
+                    callbacks=callbacks,
+                    dataset_sink_mode=dataset_sink_mode)
+
+    def _eval_dataset_sink_process(self, valid_dataset, list_callback=None, cb_params=None):
+        """
+        Evaluation. The data would be passed to network through dataset channel.
+
+        Args:
+            valid_dataset (Dataset): Dataset to evaluate the model.
+            list_callback (ListCallback): Executor of callback list. Default: None.
+            cb_params (_InternalCallbackParam): Callback parameters. Default: None.
+
+        Returns:
+            Dict, returns the loss value & metrics values for the model in test mode.
+        """
+        _device_number_check(self._parallel_mode, self._device_number)
+
+        run_context = RunContext(cb_params)
+
+        # remove later to deal with loop sink
+        need_wrap = False
+        if not hasattr(valid_dataset, '__ME_INITED__') and context.get_context("enable_loop_sink") \
+                and not context.get_context("enable_ge"):
+            need_wrap = True
+
+        valid_dataset.__loop_size__ = 1
+        dataset_helper = DatasetHelper(valid_dataset)
+
+        # remove later to deal with loop sink
+        if need_wrap:
+            self._eval_network = nn.DataWrapper(self._eval_network, *(dataset_helper.types_shapes()),
+                                                valid_dataset.__ME_INITED__)
+            self._eval_network.set_train(mode=False)
+            self._eval_network.phase = 'eval'
+        list_callback.begin(run_context)
+
+        for inputs in dataset_helper:
+            cb_params.cur_step_num += 1
+            list_callback.step_begin(run_context)
+
+            outputs = self._eval_network(*inputs)
+
+            cb_params.net_outputs = outputs
+            list_callback.step_end(run_context)
+            self._update_metrics(outputs)
+
+        metrics = self._get_metrics()
+        cb_params.metrics = metrics
+        list_callback.end(run_context)
+
+        return metrics
+
+    def _eval_process(self, valid_dataset, list_callback=None, cb_params=None):
+        """
+        Evaluation. The data would be passed to network directly.
+
+        Args:
+            valid_dataset (Dataset): Dataset to evaluate the model.
+            list_callback (ListCallback): Executor of callback list. Default: None.
+            cb_params (_InternalCallbackParam): Callback parameters. Default: None.
+
+        Returns:
+            Dict, returns the loss value & metrics values for the model in test mode.
+        """
+        run_context = RunContext(cb_params)
+        list_callback.begin(run_context)
+
+        dataset_helper = DatasetHelper(valid_dataset, dataset_sink_mode=False)
+        for next_element in dataset_helper:
+            cb_params.cur_step_num += 1
+            list_callback.step_begin(run_context)
+            outputs = self._eval_network(*next_element)
+            cb_params.net_outputs = outputs
+            list_callback.step_end(run_context)
+            self._update_metrics(outputs)
+
+        metrics = self._get_metrics()
+        cb_params.metrics = metrics
+        list_callback.end(run_context)
+        return metrics
+
+    def eval(self, valid_dataset, callbacks=None, dataset_sink_mode=True):
+        """
+        Evaluation API where the iteration is controlled by python front-end.
+
+        Configure to pynative mode, the evaluation will be performed with dataset non-sink mode.
+
+        Note:
+            CPU is not supported when dataset_sink_mode is true.
+            If dataset_sink_mode is True, data will be sent to device. If device is Ascend, features
+            of data will be transferred one by one. The limitation of data transmission per time is 256M.
+
+        Args:
+            valid_dataset (Dataset): Dataset to evaluate the model.
+            callbacks (list): List of callback object. Callbacks which should be excuted
+                              while training. Default: None.
+            dataset_sink_mode (bool): Determines whether to pass the data through dataset channel. Default: True.
+
+        Returns:
+            Dict, returns the loss value & metrics values for the model in test mode.
+
+        Examples:
+            >>> dataset = get_dataset()
+            >>> net = Net()
+            >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+            >>> model = Model(net, loss_fn=loss, optimizer=None, metrics={'acc'})
+            >>> model.eval(dataset)
+        """
+        check_bool(dataset_sink_mode)
+        if not self._metric_fns:
+            raise ValueError("metric fn can not be None or empty.")
+
+        list_callback = _build_callbacks(callbacks)
+        cb_params = _InternalCallbackParam()
+        cb_params.eval_network = self._eval_network
+        cb_params.valid_dataset = valid_dataset
+        cb_params.batch_num = valid_dataset.get_dataset_size()
+        cb_params.mode = "eval"
+        cb_params.cur_step_num = 0
+
+        self._eval_network.set_train(mode=False)
+        self._eval_network.phase = 'eval'
+
+        self._clear_metrics()
+
+        if dataset_sink_mode:
+            return self._eval_dataset_sink_process(valid_dataset, list_callback, cb_params)
+        return self._eval_process(valid_dataset, list_callback, cb_params)
+
+    def predict(self, *predict_data):
+        """
+        Generates output predictions for the input samples.
+
+        Data could be single tensor, or list of tensor, tuple of tensor.
+
+        Note:
+            Batch data should be put together in one tensor.
+
+        Args:
+           predict_data (Tensor): Tensor of predict data. can be array, list or tuple.
+
+        Returns:
+            Tensor, array(s) of predictions.
+
+        Examples:
+            >>> input_data = Tensor(np.random.randint(0, 255, [1, 3, 224, 224]), mindspore.float32)
+            >>> model = Model(Net())
+            >>> model.predict(input_data)
+        """
+        self._predict_network.set_train(False)
+        check_input_data(*predict_data, data_class=Tensor)
+        result = self._predict_network(*predict_data)
+
+        check_output_data(result)
+        return result
+
+
+__all__ = ["Model"]

example/resnet50_imagenet2012_THOR/model/resnet.py

@@ -0,0 +1,359 @@
+# Copyright 2020 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.
+# ============================================================================
+"""ResNet."""
+import math
+
+import mindspore.nn as nn
+import numpy as np
+from mindspore.common.tensor import Tensor
+from mindspore.ops import operations as P
+from second_order.thor_layer import Conv2d_Thor, Dense_Thor
+
+
+def calculate_gain(nonlinearity, param=None):
+    linear_fns = ['linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d']
+    if nonlinearity in linear_fns or nonlinearity == 'sigmoid':
+        return 1
+    elif nonlinearity == 'tanh':
+        return 5.0 / 3
+    elif nonlinearity == 'relu':
+        return math.sqrt(2.0)
+    elif nonlinearity == 'leaky_relu':
+        if param is None:
+            negative_slope = 0.01
+        elif not isinstance(param, bool) and isinstance(param, int) or isinstance(param, float):
+            # True/False are instances of int, hence check above
+            negative_slope = param
+        else:
+            raise ValueError("negative_slope {} not a valid number".format(param))
+        return math.sqrt(2.0 / (1 + negative_slope ** 2))
+    else:
+        raise ValueError("Unsupported nonlinearity {}".format(nonlinearity))
+
+
+def _calculate_fan_in_and_fan_out(tensor):
+    dimensions = len(tensor)
+    if dimensions < 2:
+        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")
+
+    if dimensions == 2:  # Linear
+        fan_in = tensor[1]
+        fan_out = tensor[0]
+    else:
+        num_input_fmaps = tensor[1]
+        num_output_fmaps = tensor[0]
+        receptive_field_size = 1
+        if dimensions > 2:
+            receptive_field_size = tensor[2] * tensor[3]
+        fan_in = num_input_fmaps * receptive_field_size
+        fan_out = num_output_fmaps * receptive_field_size
+    return fan_in, fan_out
+
+
+def _calculate_correct_fan(tensor, mode):
+    mode = mode.lower()
+    valid_modes = ['fan_in', 'fan_out']
+    if mode not in valid_modes:
+        raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes))
+
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    return fan_in if mode == 'fan_in' else fan_out
+
+
+def kaiming_normal(inputs_shape, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    fan = _calculate_correct_fan(inputs_shape, mode)
+    gain = calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    return np.random.normal(0, std, size=inputs_shape).astype(np.float32)
+
+
+def kaiming_uniform(inputs_shape, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    fan = _calculate_correct_fan(inputs_shape, mode)
+    gain = calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    bound = math.sqrt(3.0) * std  # Calculate uniform bounds from standard deviation
+    return np.random.uniform(-bound, bound, size=inputs_shape).astype(np.float32)
+
+
+def _conv3x3(in_channel, out_channel, stride=1, damping=0.03, loss_scale=1, frequency=278):
+    weight_shape = (out_channel, in_channel, 3, 3)
+    weight = Tensor(kaiming_normal(weight_shape, mode="fan_out", nonlinearity='relu'))
+    return Conv2d_Thor(in_channel, out_channel,
+                       kernel_size=3, stride=stride, padding=0, pad_mode='same', weight_init=weight,
+                       damping=damping, loss_scale=loss_scale, frequency=frequency)
+    # return nn.Conv2d(in_channel, out_channel,
+    #                  kernel_size=3, stride=stride, padding=0, pad_mode='same', weight_init=weight)
+
+
+def _conv1x1(in_channel, out_channel, stride=1, damping=0.03, loss_scale=1, frequency=278):
+    weight_shape = (out_channel, in_channel, 1, 1)
+    weight = Tensor(kaiming_normal(weight_shape, mode="fan_out", nonlinearity='relu'))
+    return Conv2d_Thor(in_channel, out_channel,
+                       kernel_size=1, stride=stride, padding=0, pad_mode='same', weight_init=weight,
+                       damping=damping, loss_scale=loss_scale, frequency=frequency)
+
+
+def _conv7x7(in_channel, out_channel, stride=1, damping=0.03, loss_scale=1, frequency=278):
+    weight_shape = (out_channel, in_channel, 7, 7)
+    weight = Tensor(kaiming_normal(weight_shape, mode="fan_out", nonlinearity='relu'))
+    return Conv2d_Thor(in_channel, out_channel,
+                       kernel_size=7, stride=stride, padding=0, pad_mode='same', weight_init=weight,
+                       damping=damping, loss_scale=loss_scale, frequency=frequency)
+
+
+def _bn(channel):
+    return nn.BatchNorm2d(channel, eps=1e-4, momentum=0.9,
+                          gamma_init=1, beta_init=0, moving_mean_init=0, moving_var_init=1)
+
+
+def _bn_last(channel):
+    return nn.BatchNorm2d(channel, eps=1e-4, momentum=0.9,
+                          gamma_init=1, beta_init=0, moving_mean_init=0, moving_var_init=1)
+
+
+def _fc(in_channel, out_channel, damping, loss_scale, frequency):
+    weight_shape = (out_channel, in_channel)
+    weight = Tensor(kaiming_uniform(weight_shape, a=math.sqrt(5))
+    return Dense_Thor(in_channel, out_channel, has_bias=False, weight_init=weight,
+                      bias_init=0, damping=damping, loss_scale=loss_scale, frequency=frequency)
+
+
+class ResidualBlock(nn.Cell):
+    """
+    ResNet V1 residual block definition.
+ 
+    Args:
+        in_channel (int): Input channel.
+        out_channel (int): Output channel.
+        stride (int): Stride size for the first convolutional layer. Default: 1.
+ 
+    Returns:
+        Tensor, output tensor.
+ 
+    Examples:
+        >>> ResidualBlock(3, 256, stride=2)
+    """
+    expansion = 4
+
+    def __init__(self,
+                 in_channel,
+                 out_channel,
+                 stride=1,
+                 damping=0.03,
+                 loss_scale=1,
+                 frequency=278):
+        super(ResidualBlock, self).__init__()
+
+        channel = out_channel // self.expansion
+        self.conv1 = _conv1x1(in_channel, channel, stride=1, damping=damping, loss_scale=loss_scale,
+                              frequency=frequency)
+        self.bn1 = _bn(channel)
+
+        self.conv2 = _conv3x3(channel, channel, stride=stride, damping=damping, loss_scale=loss_scale,
+                              frequency=frequency)
+        self.bn2 = _bn(channel)
+
+        self.conv3 = _conv1x1(channel, out_channel, stride=1, damping=damping, loss_scale=loss_scale,
+                              frequency=frequency)
+        self.bn3 = _bn_last(out_channel)
+
+        self.relu = nn.ReLU()
+
+        self.down_sample = False
+
+        if stride != 1 or in_channel != out_channel:
+            self.down_sample = True
+        self.down_sample_layer = None
+
+        if self.down_sample:
+            self.down_sample_layer = nn.SequentialCell([_conv1x1(in_channel, out_channel, stride,
+                                                                 damping=damping, loss_scale=loss_scale,
+                                                                 frequency=frequency),
+                                                        _bn(out_channel)])
+        self.add = P.TensorAdd()
+
+    def construct(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.down_sample:
+            identity = self.down_sample_layer(identity)
+
+        out = self.add(out, identity)
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Cell):
+    """
+    ResNet architecture.
+ 
+    Args:
+        block (Cell): Block for network.
+        layer_nums (list): Numbers of block in different layers.
+        in_channels (list): Input channel in each layer.
+        out_channels (list): Output channel in each layer.
+        strides (list):  Stride size in each layer.
+        num_classes (int): The number of classes that the training images are belonging to.
+    Returns:
+        Tensor, output tensor.
+ 
+    Examples:
+        >>> ResNet(ResidualBlock,
+        >>>        [3, 4, 6, 3],
+        >>>        [64, 256, 512, 1024],
+        >>>        [256, 512, 1024, 2048],
+        >>>        [1, 2, 2, 2],
+        >>>        10)
+    """
+
+    def __init__(self,
+                 block,
+                 layer_nums,
+                 in_channels,
+                 out_channels,
+                 strides,
+                 num_classes,
+                 damping,
+                 loss_scale,
+                 frequency):
+        super(ResNet, self).__init__()
+
+        if not len(layer_nums) == len(in_channels) == len(out_channels) == 4:
+            raise ValueError("the length of layer_num, in_channels, out_channels list must be 4!")
+
+        self.conv1 = _conv7x7(3, 64, stride=2, damping=damping, loss_scale=loss_scale, frequency=frequency)
+        self.bn1 = _bn(64)
+        self.relu = P.ReLU()
+        self.maxpool = P.MaxPoolWithArgmax(padding="same", ksize=3, strides=2)
+
+        self.layer1 = self._make_layer(block,
+                                       layer_nums[0],
+                                       in_channel=in_channels[0],
+                                       out_channel=out_channels[0],
+                                       stride=strides[0],
+                                       damping=damping,
+                                       loss_scale=loss_scale,
+                                       frequency=frequency)
+        self.layer2 = self._make_layer(block,
+                                       layer_nums[1],
+                                       in_channel=in_channels[1],
+                                       out_channel=out_channels[1],
+                                       stride=strides[1],
+                                       damping=damping,
+                                       loss_scale=loss_scale,
+                                       frequency=frequency)
+        self.layer3 = self._make_layer(block,
+                                       layer_nums[2],
+                                       in_channel=in_channels[2],
+                                       out_channel=out_channels[2],
+                                       stride=strides[2], damping=damping,
+                                       loss_scale=loss_scale,
+                                       frequency=frequency)
+        self.layer4 = self._make_layer(block,
+                                       layer_nums[3],
+                                       in_channel=in_channels[3],
+                                       out_channel=out_channels[3],
+                                       stride=strides[3],
+                                       damping=damping,
+                                       loss_scale=loss_scale,
+                                       frequency=frequency)
+
+        self.mean = P.ReduceMean(keep_dims=True)
+        self.flatten = nn.Flatten()
+        self.end_point = _fc(out_channels[3], num_classes, damping=damping, loss_scale=loss_scale, frequency=frequency)
+
+    def _make_layer(self, block, layer_num, in_channel, out_channel, stride,
+                    damping, loss_scale, frequency):
+        """
+        Make stage network of ResNet.
+ 
+        Args:
+            block (Cell): Resnet block.
+            layer_num (int): Layer number.
+            in_channel (int): Input channel.
+            out_channel (int): Output channel.
+            stride (int): Stride size for the first convolutional layer.
+ 
+        Returns:
+            SequentialCell, the output layer.
+ 
+        Examples:
+            >>> _make_layer(ResidualBlock, 3, 128, 256, 2)
+        """
+        layers = []
+
+        resnet_block = block(in_channel, out_channel, stride=stride,
+                             damping=damping, loss_scale=loss_scale, frequency=frequency)
+        layers.append(resnet_block)
+
+        for _ in range(1, layer_num):
+            resnet_block = block(out_channel, out_channel, stride=1,
+                                 damping=damping, loss_scale=loss_scale, frequency=frequency)
+            layers.append(resnet_block)
+
+        return nn.SequentialCell(layers)
+
+    def construct(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        c1, argmax = self.maxpool(x)
+
+        c2 = self.layer1(c1)
+        c3 = self.layer2(c2)
+        c4 = self.layer3(c3)
+        c5 = self.layer4(c4)
+
+        out = self.mean(c5, (2, 3))
+        out = self.flatten(out)
+        out = self.end_point(out)
+
+        return out
+
+
+def resnet50(class_num=10, damping=0.03, loss_scale=1, frequency=278):
+    """
+    Get ResNet50 neural network.
+ 
+    Args:
+        class_num (int): Class number.
+ 
+    Returns:
+        Cell, cell instance of ResNet50 neural network.
+ 
+    Examples:
+        >>> net = resnet50(10)
+    """
+    return ResNet(ResidualBlock,
+                  [3, 4, 6, 3],
+                  [64, 256, 512, 1024],
+                  [256, 512, 1024, 2048],
+                  [1, 2, 2, 2],
+                  class_num,
+                  damping,
+                  loss_scale,
+                  frequency)

example/resnet50_imagenet2012_THOR/model/thor.py

@@ -0,0 +1,202 @@
+# Copyright 2020 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.
+# ============================================================================
+"""momentum"""
+import mindspore.common.dtype as mstype
+from mindspore.common.initializer import initializer
+from mindspore.common.parameter import Parameter
+from mindspore.common.parameter import ParameterTuple
+from mindspore.common.tensor import Tensor
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore.parallel._utils import _get_device_num, _get_mirror_mean
+
+from cus_ops.cus_matmul_cube_dense_right import CusMatMulCubeDenseRight
+from cus_ops.cus_matmul_cube_fracz_left_cast import CusMatMulCubeFraczLeftCast
+from cus_ops.cus_matmul_cube_dense_left import CusMatMulCubeDenseLeft
+from cus_ops.cus_matmul_cube_fracz_right_mul import CusMatMulCubeFraczRightMul
+from model.grad_reducer_thor import DistributedGradReducerThor
+
+momentum_opt = C.MultitypeFuncGraph("momentum_opt")
+
+
+@momentum_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _tensor_run_opt_ext(opt, learning_rate, momentum, gradient, weight, moment):
+    """Apply momentum optimizer to the weight parameter using Tensor."""
+    success = True
+    success = F.depend(success, opt(weight, moment, learning_rate, gradient, momentum))
+    return success
+
+
+op_add = P.AddN()
+apply_decay = C.MultitypeFuncGraph("apply_decay")
+
+
+@apply_decay.register("Number", "Bool", "Tensor", "Tensor")
+def _tensor_apply_decay(weight_decay, if_apply, weight, gradient):
+    """Get grad with weight_decay."""
+    if if_apply:
+        return op_add((weight * weight_decay, gradient))
+    return gradient
+
+
+class THOR(Optimizer):
+    """THOR"""
+    def __init__(self, params, learning_rate, momentum, matrix_A, matrix_G, A_inv_max, G_inv_max, weight_decay=0.0,
+                 loss_scale=1.0,
+                 decay_filter=lambda x: x.name not in []):
+        super(THOR, self).__init__(learning_rate, params, weight_decay, loss_scale)
+        if isinstance(momentum, float) and momentum < 0.0:
+            raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum))
+        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
+        self.params = self.parameters
+        self.moments = self.params.clone(prefix="moments", init='zeros')
+        self.hyper_map = C.HyperMap()
+        self.opt = P.ApplyMomentum()
+        self.matrix_A = ParameterTuple(matrix_A)
+        self.matrix_G = ParameterTuple(matrix_G)
+        self.A_inv_max = ParameterTuple(A_inv_max)
+        self.G_inv_max = ParameterTuple(G_inv_max)
+        self.cube_matmul_left = CusMatMulCubeFraczLeftCast()
+        self.cube_matmul_left_fc = CusMatMulCubeDenseLeft()
+        self.cube_matmul_right_fc = CusMatMulCubeDenseRight()
+        self.cube_matmul_right_mul = CusMatMulCubeFraczRightMul()
+        self.transpose = P.Transpose()
+        self.shape = P.Shape()
+        self.reshape = P.Reshape()
+        self.mul = P.Mul()
+        self.weight_idx = []
+        for i in range(len(self.params)):
+            if "conv" in self.params[i].name or "end_point" in self.params[i].name:
+                self.weight_idx.append(i)
+        self.weight_idx.append(len(self.params))
+        self.feature_map = [1.0 / 12544, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136,
+                            1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136, 1.0 / 3136,
+                            1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784,
+                            1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784, 1.0 / 784,
+                            1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196,
+                            1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196, 1.0 / 196,
+                            1.0 / 196, 1.0 / 196, 1.0 / 196,
+                            1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49, 1.0 / 49,
+                            1.0]
+        mean = _get_mirror_mean()
+        degree = _get_device_num()
+        self.grad_reducer_Amax = DistributedGradReducerThor(self.parameters, 2, mean, degree)
+        self.grad_reducer_Gmax = DistributedGradReducerThor(self.parameters, 5, mean, degree)
+        self.grad_reducer_A = DistributedGradReducerThor(self.parameters, 3, mean, degree)
+        self.grad_reducer_G = DistributedGradReducerThor(self.parameters, 4, mean, degree)
+        self.matrix_A_inv = ()
+        self.matrix_G_inv = ()
+        self.matrix_max_inv = ()
+
+        for i in range(54):
+            self.matrix_max_inv = self.matrix_max_inv + (
+                Parameter(initializer(1, [1], mstype.float32), name="matrix_max" + str(i), requires_grad=False),)
+        self.log = P.Log()
+        self.exp = P.Exp()
+        self.sqrt = P.Sqrt()
+        self.matrix_max_inv = ParameterTuple(self.matrix_max_inv)
+        self.assign = P.Assign()
+        self.cast = P.Cast()
+        self.thor = True
+        self.weight_decay = weight_decay * loss_scale
+        self.decay_flags = tuple(decay_filter(x) for x in self.parameters)
+
+    def construct(self, gradients):
+        params = self.params
+        moments = self.moments
+        if self.thor:
+            matrix_A_allreduce = ()
+            matrix_G_allreduce = ()
+            matrix_A_max_allreduce = ()
+            matrix_G_max_allreduce = ()
+            for i in range(54):
+                g = gradients[i * 3]
+                matrix_A = self.matrix_A[i]
+                matrix_G = self.matrix_G[i]
+                A_max = self.A_inv_max[i]
+                G_max = self.G_inv_max[i]
+                matrix_A = F.depend(matrix_A, g)
+                matrix_G = F.depend(matrix_G, g)
+                A_max = F.depend(A_max, g)
+                G_max = F.depend(G_max, g)
+                matrix_A_allreduce = matrix_A_allreduce + (matrix_A,)
+                matrix_G_allreduce = matrix_G_allreduce + (matrix_G,)
+                matrix_A_max_allreduce = matrix_A_max_allreduce + (A_max,)
+                matrix_G_max_allreduce = matrix_G_max_allreduce + (G_max,)
+            matrix_A_allreduce = self.grad_reducer_A(matrix_A_allreduce)
+            matrix_G_allreduce = self.grad_reducer_G(matrix_G_allreduce)
+            matrix_A_max_allreduce = self.grad_reducer_Amax(matrix_A_max_allreduce)
+            matrix_G_max_allreduce = self.grad_reducer_Gmax(matrix_G_max_allreduce)
+            new_grads = ()
+            for i in range(54):
+                g = gradients[i * 3]
+                temp_a = matrix_A_allreduce[i]
+                temp_g = matrix_G_allreduce[i]
+                temp_a = self.cast(temp_a, mstype.float32)
+                temp_g = self.cast(temp_g, mstype.float32)
+                matrix_A_inv_max = self.log(matrix_A_max_allreduce[i])
+                matrix_A_inv_max = self.mul(matrix_A_inv_max, -1)
+                matrix_A_inv_max = self.exp(matrix_A_inv_max)
+                temp_a = self.mul(temp_a, matrix_A_inv_max)
+                matrix_G_inv_max = self.log(matrix_G_max_allreduce[i])
+                matrix_G_inv_max = self.mul(matrix_G_inv_max, -1)
+                matrix_G_inv_max = self.exp(matrix_G_inv_max)
+                temp_g = self.mul(temp_g, matrix_G_inv_max)
+                temp_max = self.mul(matrix_A_max_allreduce[i], matrix_G_max_allreduce[i])
+                temp_max = self.mul(temp_max, self.feature_map[i])
+                if i == 53:
+                    g = self.cube_matmul_left_fc(temp_g, g)
+                    g = self.cube_matmul_right_fc(g, temp_a, temp_max)
+                else:
+                    g = self.cube_matmul_left(temp_g, g)
+                    g = self.cube_matmul_right_mul(g, temp_a, temp_max)
+                fake_A = self.assign(self.matrix_A[i], temp_a)
+                fake_G = self.assign(self.matrix_G[i], temp_g)
+                fake_max = self.assign(self.matrix_max_inv[i], temp_max)
+                g = F.depend(g, fake_A)
+                g = F.depend(g, fake_G)
+                g = F.depend(g, fake_max)
+                if i == 53:
+                    new_grads = new_grads + (g,)
+                else:
+                    new_grads = new_grads + (g, gradients[i * 3 + 1], gradients[i * 3 + 2])
+            gradients = new_grads
+        else:
+            new_grads = ()
+            for i in range(54):
+                g = gradients[i * 3]
+                matrix_A = self.matrix_A[i]
+                matrix_G = self.matrix_G[i]
+                matrix_max = self.matrix_max_inv[i]
+                matrix_A = F.depend(matrix_A, g)
+                matrix_G = F.depend(matrix_G, g)
+                matrix_max = F.depend(matrix_max, g)
+                if i == 53:
+                    g = self.cube_matmul_left_fc(matrix_G, g)
+                    g = self.cube_matmul_right_fc(g, matrix_A, matrix_max)
+                    new_grads = new_grads + (g,)
+                else:
+                    g = self.cube_matmul_left(matrix_G, g)
+                    g = self.cube_matmul_right_mul(g, matrix_A, matrix_max)
+                    new_grads = new_grads + (g, gradients[i * 3 + 1], gradients[i * 3 + 2])
+            gradients = new_grads
+
+        if self.weight_decay > 0:
+            gradients = self.hyper_map(F.partial(apply_decay, self.weight_decay), self.decay_flags,
+                                       params, gradients)
+        gradients = self.scale_grad(gradients)
+        lr = self.get_lr()
+        success = self.hyper_map(F.partial(momentum_opt, self.opt, lr, self.momentum), gradients, params, moments)
+        return success

example/resnet50_imagenet2012_THOR/model/thor_layer.py

@@ -0,0 +1,488 @@
+# Copyright 2020 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.
+# ============================================================================
+"""thor_layer"""
+import mindspore as ms
+import mindspore.common.dtype as mstype
+from mindspore._checkparam import check_bool, twice, check_int_positive
+from mindspore._extends import cell_attr_register
+from mindspore.common.initializer import initializer
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+from mindspore.nn.cell import Cell
+from mindspore.nn.layer.activation import get_activation
+from mindspore.ops import operations as P
+
+from cus_ops.cus_batch_matmul import CusBatchMatMul
+from cus_ops.cus_cholesky_trsm import CusCholeskyTrsm
+from cus_ops.cus_fused_abs_max1 import CusFusedAbsMax1
+from cus_ops.cus_img2col import CusImg2Col
+from cus_ops.cus_matmul_cube import CusMatMulCube
+from cus_ops.cus_matrix_combine import CusMatrixCombine
+from cus_ops.cus_transpose02314 import CusTranspose02314
+
+import numpy as np
+C0 = 16
+
+
+def caculate_device_shape(matrix_dim, channel, is_A):
+    ll = (0)
+    if is_A:
+        if channel // C0 == 0:
+            matrix_dim = (matrix_dim / channel) * C0
+        ll = (int(matrix_dim // C0), int(matrix_dim // C0), C0, C0), int(matrix_dim)
+    else:
+        ll = (int(matrix_dim // C0), int(matrix_dim // C0), C0, C0), int(matrix_dim)
+    return ll
+
+class _Conv(Cell):
+    r"""Applies a N-D convolution over an input signal composed of several input
+       planes.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 pad_mode,
+                 padding,
+                 dilation,
+                 group,
+                 data_format,
+                 has_bias,
+                 weight_init,
+                 bias_init,
+                 ):
+        super(_Conv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.pad_mode = pad_mode
+        self.padding = padding
+        self.dilation = dilation
+        self.group = group
+        self.data_format = data_format
+        self.has_bias = has_bias
+        if not (isinstance(in_channels, int) and in_channels > 0):
+            raise ValueError('Attr \'in_channels\' of \'Conv2D\' Op passed '
+                             + str(in_channels) + ', should be a int and greater than 0.')
+        if (not isinstance(kernel_size, tuple)) or len(kernel_size) != 2 or \
+                (not isinstance(kernel_size[0], int)) or (not isinstance(kernel_size[1], int)) or \
+                kernel_size[0] < 1 or kernel_size[1] < 1:
+            raise ValueError('Attr \'kernel_size\' of \'Conv2D\' Op passed '
+                             + str(self.kernel_size) + ', should be a int or tuple and equal to or greater than 1.')
+        if in_channels % group != 0:
+            raise ValueError('Attr \'in_channels\' of \'Conv2D\' Op must be divisible by '
+                             'attr \'group\' of \'Conv2D\' Op.')
+        if out_channels % group != 0:
+            raise ValueError('Attr \'out_channels\' of \'Conv2D\' Op must be divisible by '
+                             'attr \'group\' of \'Conv2D\' Op.')
+
+        self.weight = Parameter(initializer(
+            weight_init, [out_channels, in_channels // group, *kernel_size]), name='weight')
+
+        if check_bool(has_bias):
+            self.bias = Parameter(_initializer(
+                bias_init, [out_channels]), name='bias')
+        else:
+            if bias_init != 'zeros':
+                logger.warning("Value of 'has_bias' is False, value of 'bias_init' will be ignored.")
+            self.bias = None
+
+    def construct(self, *inputs):
+        raise NotImplementedError
+
+
+class Conv2d_Thor(_Conv):
+    """Conv2d_Thor"""
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 pad_mode='same',
+                 padding=0,
+                 dilation=1,
+                 group=1,
+                 data_format='NCHW',
+                 has_bias=False,
+                 weight_init='normal',
+                 damping=0.03,
+                 loss_scale=1,
+                 frequency=278,
+                 bias_init='zeros'):
+        self.thor = True
+        ksizes = (1, kernel_size, kernel_size, 1)
+        self.hw = kernel_size * kernel_size
+        strides = (1, stride, stride, 1)
+        kernel_size = twice(kernel_size)
+        super(Conv2d_Thor, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            pad_mode,
+            padding,
+            dilation,
+            group,
+            data_format,
+            has_bias,
+            weight_init,
+            bias_init,
+        )
+        self.conv2d = P.Conv2D(out_channel=self.out_channels,
+                               kernel_size=self.kernel_size,
+                               mode=1,
+                               pad_mode=self.pad_mode,
+                               pad=self.padding,
+                               stride=self.stride,
+                               dilation=self.dilation,
+                               group=self.group
+                               )
+
+        self.img2col = CusImg2Col(ksizes=ksizes, strides=strides)
+        self.cube_matmul = CusMatMulCube(transpose_a=True)
+        self.matrix_combine = CusMatrixCombine()
+        self.cholesky = CusCholeskyTrsm()
+        self.transpose02314 = CusTranspose02314()
+        self.matrix_A_dim = self.in_channels * self.kernel_size[0] * self.kernel_size[1]
+        self.matrix_G_dim = self.out_channels
+        self.matrix_A_device_shape, self.matrix_A_device_dim = caculate_device_shape(self.matrix_A_dim,
+                                                                                     self.in_channels, True)
+        self.matrix_G_device_shape, self.matrix_G_device_dim = caculate_device_shape(self.matrix_G_dim,
+                                                                                     self.in_channels, False)
+        self.matrix_A_device_temp_shape = (
+            self.matrix_A_device_shape[0], self.matrix_A_device_shape[2], self.matrix_A_device_shape[1],
+            self.matrix_A_device_shape[3])
+        self.matrix_G_device_temp_shape = (
+            self.matrix_G_device_shape[0], self.matrix_G_device_shape[2], self.matrix_G_device_shape[1],
+            self.matrix_G_device_shape[3])
+        self.matrix_A_inv = Parameter(
+            Tensor(np.reshape(np.identity(self.matrix_A_device_dim).astype(np.float16), self.matrix_A_device_shape)),
+            name='matrix_A_inv', requires_grad=False)
+        self.A_inv_max = Parameter(initializer(0, [1], mstype.float32), name="A_inv_max", requires_grad=False)
+        self.matrix_G_inv = Parameter(
+            Tensor(np.reshape(np.identity(self.matrix_G_device_dim).astype(np.float16), self.matrix_G_device_shape)),
+            name="matrix_G_inv", requires_grad=False)
+
+        self.G_inv_max = Parameter(initializer(0, [1], mstype.float32), name="G_inv_max", requires_grad=False)
+        self.fake_G = Tensor(
+            np.reshape(np.identity(self.matrix_G_device_dim).astype(np.float16), self.matrix_G_device_shape))
+        self.fake_G_inv_max = Tensor(np.zeros([1,]).astype(np.float32))
+
+        self.shape = P.Shape()
+        self.reshape = P.Reshape()
+        self.transpose = P.Transpose()
+        self.cov_step = Parameter(initializer(0, [1], mstype.int32), name="cov_step", requires_grad=False)
+        self.mul = P.Mul()
+        self.cast = P.Cast()
+        self.damping = Tensor(damping)
+        self.vector_matmul = CusBatchMatMul()
+        self.diag_block_dim = 128
+        self.channels_slice_flag = False
+        if self.in_channels % C0 != 0:
+            self.channels_slice_flag = True
+
+        self.padA_flag = False
+        if (self.matrix_A_dim // self.diag_block_dim) * self.diag_block_dim != self.matrix_A_dim \
+            and self.matrix_A_dim > self.diag_block_dim:
+            self.padA_flag = True
+            pad_dim = self.diag_block_dim - self.matrix_A_dim % self.diag_block_dim
+            self.padA = P.Pad(((0, pad_dim), (0, pad_dim)))
+        self.device_shape_pad_flag = False
+        if self.matrix_A_dim != self.matrix_A_device_dim:
+            self.device_shape_pad_flag = True
+            self.device_shape_pad = P.Pad(((0, 0), (0, C0 - self.in_channels), (0, 0), (0, C0 - self.in_channels)))
+        self.slice = P.Slice()
+        self.gather = P.GatherV2()
+        self.freq = Tensor(frequency, mstype.int32)
+        self.loss_scale = Tensor(1 / loss_scale, mstype.float16)
+        self.axis = 0
+
+        dampingA_dim = self.matrix_A_dim
+        if (self.matrix_A_dim % self.diag_block_dim) != 0 and self.matrix_A_dim > self.diag_block_dim:
+            dampingA_dim = (self.matrix_A_dim // self.diag_block_dim + 1) * self.diag_block_dim
+        dampingG_dim = self.matrix_G_dim
+        if (self.matrix_G_dim % self.diag_block_dim) != 0 and self.matrix_G_dim > self.diag_block_dim:
+            dampingG_dim = (self.matrix_G_dim // self.diag_block_dim + 1) * self.diag_block_dim
+
+        self.dampingA = Tensor(np.identity(dampingA_dim), mstype.float32)
+        self.dampingG = Tensor(np.identity(dampingG_dim), mstype.float32)
+        self.fused_abs_max1 = CusFusedAbsMax1([self.matrix_A_dim, self.matrix_A_dim])
+        self.fused_abs_max2 = CusFusedAbsMax1()
+        self.log = P.Log()
+        self.exp = P.Exp()
+        self.sqrt = P.Sqrt()
+        self.getG = P.InsertGradientOf(self.save_gradient)
+
+    def save_gradient(self, dout):
+        """save_gradient"""
+        out = dout
+        dout = self.mul(dout, self.loss_scale)
+        dout = self.mul(dout, 32.0)
+        dout = self.transpose02314(dout)
+        dout_shape = self.shape(dout)
+        normalizer = dout_shape[0]
+
+        matrix_G = self.cube_matmul(dout, dout)
+        normalizer = self.cast(normalizer, ms.float32)
+        matrix_G = self.mul(matrix_G, 1.0 / normalizer)
+        damping_step = self.gather(self.damping, self.cov_step, 0)
+        self.cov_step = self.cov_step + self.freq
+        damping_step = self.cast(damping_step, mstype.float32)
+        damping = self.mul(damping_step, 32.0 / normalizer)
+        damping = self.sqrt(damping)
+        dampingG = self.cast(self.dampingG, mstype.float32)
+        matrix_G = matrix_G + damping * dampingG
+
+        matrix_G_inv = self.cholesky(matrix_G)
+        matrix_G_inv = self.vector_matmul(matrix_G_inv, matrix_G_inv)
+        matrix_G_inv_max = self.fused_abs_max2(matrix_G_inv)
+        matrix_G_inv_max = self.fused_abs_max2(matrix_G_inv_max)
+        self.G_inv_max = matrix_G_inv_max
+        matrix_G_inv = self.matrix_combine(matrix_G_inv)
+        matrix_G_inv = self.reshape(matrix_G_inv, self.matrix_G_device_temp_shape)
+        matrix_G_inv = self.transpose(matrix_G_inv, (2, 0, 1, 3))
+        matrix_G = self.cast(matrix_G_inv, mstype.float16)
+        self.matrix_G_inv = matrix_G
+        return out
+
+    def construct(self, x):
+        if self.thor:
+            matrix_A = self.img2col(x)
+            matrix_A_shape = self.shape(matrix_A)
+            normalizer = matrix_A_shape[0]
+            matrix_A = self.cube_matmul(matrix_A, matrix_A)
+
+            if self.channels_slice_flag:
+                matrix_A = self.reshape(matrix_A, (self.hw, C0, self.hw, C0))
+                matrix_A = self.slice(matrix_A, (0, 0, 0, 0), (self.hw, self.in_channels, self.hw, self.in_channels))
+                matrix_A = self.reshape(matrix_A, (self.matrix_A_dim, self.matrix_A_dim))
+            normalizer = self.cast(normalizer, ms.float32)
+            matrix_A = self.mul(matrix_A, 1.0 / normalizer)
+            if self.padA_flag:
+                matrix_A = self.padA(matrix_A)
+            damping_step = self.gather(self.damping, self.cov_step, self.axis)
+            damping_step = self.cast(damping_step, mstype.float32)
+            damping = self.mul(damping_step, 32.0 / normalizer)
+            damping = self.sqrt(damping)
+            damping_A = self.cast(self.dampingA, mstype.float32)
+            matrix_A = matrix_A + damping * damping_A
+            matrix_A_inv = self.cholesky(matrix_A)
+            matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv)
+            matrix_A_inv_max = self.fused_abs_max1(matrix_A_inv)
+            matrix_A_inv_max = self.fused_abs_max2(matrix_A_inv_max)
+            self.A_inv_max = matrix_A_inv_max
+            matrix_A_inv = self.matrix_combine(matrix_A_inv)
+            matrix_A_inv = self.cast(matrix_A_inv, mstype.float16)
+            if self.padA_flag:
+                matrix_A_inv = self.slice(matrix_A_inv, (0, 0), (self.matrix_A_dim, self.matrix_A_dim))
+
+            if self.device_shape_pad_flag:
+                matrix_A_inv = self.reshape(matrix_A_inv, (self.hw, self.in_channels, self.hw, self.in_channels))
+                matrix_A_inv = self.device_shape_pad(matrix_A_inv)
+            matrix_A_inv = self.reshape(matrix_A_inv, self.matrix_A_device_temp_shape)
+            matrix_A_inv = self.transpose(matrix_A_inv, (2, 0, 1, 3))
+            self.G_inv_max = self.fake_G_inv_max
+            self.matrix_A_inv = matrix_A_inv
+            self.matrix_G_inv = self.fake_G
+            out = self.conv2d(x, self.weight)
+            out = self.getG(out)
+        else:
+            out = self.conv2d(x, self.weight)
+
+        return out
+
+    def extra_repr(self):
+        """extra_repr"""
+        s = 'input_channels={}, output_channels={}, kernel_size={},' \
+            'stride={},  pad_mode={}, padding={}, dilation={}, ' \
+            'group={}, data_format={}, has_bias={},' \
+            'weight_init={}, bias_init={}'.format(
+                self.in_channels,
+                self.out_channels,
+                self.kernel_size,
+                self.stride,
+                self.pad_mode,
+                self.padding,
+                self.dilation,
+                self.group,
+                self.data_format,
+                self.has_bias,
+                self.weight,
+                self.bias)
+
+        if self.has_bias:
+            s += ', bias={}'.format(self.bias)
+        return s
+
+
+class Dense_Thor(Cell):
+    """Dense_Thor"""
+    @cell_attr_register(attrs=['has_bias', 'activation'])
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 weight_init='normal',
+                 bias_init='zeros',
+                 damping=0.03,
+                 loss_scale=1,
+                 frequency=278,
+                 has_bias=True,
+                 activation=None):
+        super(Dense_Thor, self).__init__()
+        self.in_channels = check_int_positive(in_channels)
+        self.out_channels = check_int_positive(out_channels)
+        self.has_bias = check_bool(has_bias)
+        self.thor = True
+        if isinstance(weight_init, Tensor):
+            if weight_init.dim() != 2 or weight_init.shape()[0] != out_channels or \
+                    weight_init.shape()[1] != in_channels:
+                raise ValueError("weight_init shape error")
+
+        self.weight = Parameter(initializer(weight_init, [out_channels, in_channels]), name="weight")
+
+        if self.has_bias:
+            if isinstance(bias_init, Tensor):
+                if bias_init.dim() != 1 or bias_init.shape()[0] != out_channels:
+                    raise ValueError("bias_init shape error")
+
+            self.bias = Parameter(initializer(bias_init, [out_channels]), name="bias")
+
+        self.matmul = P.MatMul(transpose_b=True)
+        self.bias_add = P.BiasAdd()
+
+        self.activation = get_activation(activation)
+        self.activation_flag = self.activation is not None
+
+        self.matrix_A_inv = Parameter(Tensor(np.zeros([128, 128, 16, 16]).astype(np.float16)), name='matrix_A_inv',
+                                      requires_grad=False)
+        self.matrix_G_inv = Parameter(Tensor(np.zeros([63, 63, 16, 16]).astype(np.float16)), name="matrix_G_inv",
+                                      requires_grad=False)
+        self.fake_G = Tensor(np.zeros([63, 63, 16, 16]).astype(np.float16))
+
+        self.matmul = P.MatMul(transpose_b=True)
+        self.cube_matmul = CusMatMulCube(transpose_a=True)
+        self.matrix_combine = CusMatrixCombine()
+        self.cholesky = CusCholeskyTrsm()
+        self.shape = P.Shape()
+        self.reshape = P.Reshape()
+        self.transpose = P.Transpose()
+        self.cov_step = Parameter(initializer(0, [1], mstype.int32), name="cov_step", requires_grad=False)
+        self.mul = P.Mul()
+        self.cast = P.Cast()
+        self.damping = Tensor(damping)
+        self.loss_scale = Tensor(1 / loss_scale, mstype.float16)
+        self.vector_matmul = CusBatchMatMul()
+        self.pad = P.Pad(((0, 24), (0, 24)))
+        self.pad1 = P.Pad(((0, 8), (0, 8)))
+        self.slice = P.Slice()
+        self.gather = P.GatherV2()
+        self.assignadd = P.AssignAdd()
+        self.freq = Tensor(frequency, mstype.int32)
+        self.axis = 0
+        self.A_inv_max = Parameter(initializer(0, [1], mstype.float32), name="A_inv_max", requires_grad=False)
+        self.G_inv_max = Parameter(initializer(0, [1], mstype.float32), name="G_inv_max", requires_grad=False)
+        self.fused_abs_max1 = CusFusedAbsMax1([1000, 1000])
+        self.fused_abs_max2 = CusFusedAbsMax1()
+        self.log = P.Log()
+        self.exp = P.Exp()
+        self.dampingA = Tensor(np.identity(2048), mstype.float32)
+        self.dampingG = Tensor(np.identity(1024), mstype.float32)
+        self.add = P.TensorAdd()
+        self.sqrt = P.Sqrt()
+        self.getG = P.InsertGradientOf(self.save_gradient)
+
+    def save_gradient(self, dout):
+        """save_gradient"""
+        out = dout
+        dout = self.mul(dout, self.loss_scale)
+        dout = self.mul(dout, 32.0)
+        normalizer = 32
+        matrix_G = self.cube_matmul(dout, dout)
+        normalizer = self.cast(normalizer, ms.float32)
+        matrix_G = self.mul(matrix_G, 1.0 / normalizer)
+        matrix_G = self.pad(matrix_G)
+        damping_step = self.gather(self.damping, self.cov_step, 0)
+        damping_step = self.cast(damping_step, mstype.float32)
+        self.cov_step = self.cov_step + self.freq
+        damping = self.sqrt(damping_step)
+        dampingG = self.cast(self.dampingG, mstype.float32)
+        matrix_G = matrix_G + damping * dampingG
+        matrix_G_inv = self.cholesky(matrix_G)
+        matrix_G_inv = self.vector_matmul(matrix_G_inv, matrix_G_inv)
+        matrix_G_inv_max = self.fused_abs_max1(matrix_G_inv)
+        matrix_G_inv_max = self.fused_abs_max2(matrix_G_inv_max)
+        self.G_inv_max = matrix_G_inv_max
+        matrix_G_inv = self.matrix_combine(matrix_G_inv)
+        matrix_G_inv = self.slice(matrix_G_inv, (0, 0), (1000, 1000))
+        matrix_G_inv = self.pad1(matrix_G_inv)
+        matrix_G_inv_shape = self.shape(matrix_G_inv)
+        matrix_G_inv = self.reshape(matrix_G_inv, (matrix_G_inv_shape[0] / 16, 16, matrix_G_inv_shape[0] / 16, 16))
+        matrix_G_inv = self.transpose(matrix_G_inv, (2, 0, 1, 3))
+        matrix_G_inv = self.cast(matrix_G_inv, mstype.float16)
+        self.matrix_G_inv = matrix_G_inv
+        return out
+
+    def construct(self, x):
+        """construct"""
+        if self.thor:
+            inputs = self.cube_matmul(x, x)
+            normalizer = 32
+            normalizer = self.cast(normalizer, ms.float32)
+            matrix_A = self.mul(inputs, 1.0 / normalizer)
+
+            damping_step = self.gather(self.damping, self.cov_step, self.axis)
+            damping_step = self.cast(damping_step, mstype.float32)
+            damping = self.sqrt(damping_step)
+            dampingA = self.cast(self.dampingA, mstype.float32)
+            matrix_A = matrix_A + damping * dampingA
+            matrix_A_inv = self.cholesky(matrix_A)
+            matrix_A_inv = self.vector_matmul(matrix_A_inv, matrix_A_inv)
+
+            matrix_A_inv_max = self.fused_abs_max2(matrix_A_inv)
+            matrix_A_inv_max = self.fused_abs_max2(matrix_A_inv_max)
+            self.A_inv_max = matrix_A_inv_max
+
+            matrix_A_inv = self.matrix_combine(matrix_A_inv)
+            matrix_A_inv_shape = self.shape(matrix_A_inv)
+            matrix_A_inv = self.reshape(matrix_A_inv, (matrix_A_inv_shape[0] / 16, 16, matrix_A_inv_shape[0] / 16, 16))
+            matrix_A_inv = self.transpose(matrix_A_inv, (2, 0, 1, 3))
+            matrix_A_inv = self.cast(matrix_A_inv, mstype.float16)
+            self.matrix_A_inv = matrix_A_inv
+            self.matrix_G_inv = self.fake_G
+            output = self.matmul(x, self.weight)
+            output = self.getG(output)
+        else:
+            output = self.matmul(x, self.weight)
+
+        if self.has_bias:
+            output = self.bias_add(output, self.bias)
+        if self.activation_flag:
+            return self.activation(output)
+        return output
+
+    def extend_repr(self):
+        """extend_repr"""
+        str_info = 'in_channels={}, out_channels={}, weight={}, has_bias={}' \
+            .format(self.in_channels, self.out_channels, self.weight, self.has_bias)
+        if self.has_bias:
+            str_info = str_info + ', bias={}'.format(self.bias)
+
+        if self.activation_flag:
+            str_info = str_info + ', activation={}'.format(self.activation)
+
+        return str_info

example/resnet50_imagenet2012_THOR/run_distribute_train_new.sh

@@ -0,0 +1,56 @@
+#!/bin/bash
+# Copyright 2020 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.
+# ============================================================================
+
+if [ $# != 3 ]
+then
+    echo "Usage: sh run_distribute_train.sh [MINDSPORE_HCCL_CONFIG_PATH] [DATASET_PATH] [DEVICE_NUM]"
+exit 1
+fi
+
+if [ ! -f $1 ]
+then
+    echo "error: DMINDSPORE_HCCL_CONFIG_PATH=$1 is not a file"
+exit 1
+fi
+
+if [ ! -d $2 ]
+then
+    echo "error: DATASET_PATH=$2 is not a directory"
+exit 1
+fi
+
+ulimit -u unlimited
+export DEVICE_NUM=$3
+export RANK_SIZE=$3
+export MINDSPORE_HCCL_CONFIG_PATH=$1
+
+for((i=0; i<${DEVICE_NUM}; i++))
+do
+    export DEVICE_ID=$i
+    export RANK_ID=$i
+    rm -rf ./train_parallel$i
+    mkdir ./train_parallel$i
+    cp *.py ./train_parallel$i
+    cp *.sh ./train_parallel$i
+    cp -r second_order ./train_parallel$i/second_order
+    cp -r test_ops ./train_parallel$i/test_ops
+    cd ./train_parallel$i || exit
+    echo "start training for rank $RANK_ID, device $DEVICE_ID"
+
+    env > env.log
+    python train_0517_1.py --do_train=True --run_distribute=True --device_num=$DEVICE_NUM --dataset_path=$2 > log 2>&1 &
+    cd ..
+done

example/resnet50_imagenet2012_THOR/train.py

@@ -0,0 +1,143 @@
+# Copyright 2020 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.
+# ============================================================================
+"""train_imagenet."""
+import argparse
+import os
+import random
+
+import mindspore.dataset.engine as de
+from mindspore import Tensor
+from mindspore import context
+from mindspore.communication.management import init
+from mindspore.parallel._auto_parallel_context import auto_parallel_context
+from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
+from mindspore.train.loss_scale_manager import FixedLossScaleManager
+from mindspore.train.model import ParallelMode
+from second_order.model_second_order import Model
+from second_order.resnet import resnet50
+from second_order.thor import THOR
+
+import numpy as np
+from config_imagenet import config
+from crossentropy import CrossEntropy
+from dataset_imagenet import create_dataset
+from lr_generator import  warmup_cosine_annealing_lr
+
+random.seed(1)
+np.random.seed(1)
+de.config.set_seed(1)
+
+parser = argparse.ArgumentParser(description='Image classification')
+parser.add_argument('--run_distribute', type=bool, default=False, help='Run distribute')
+parser.add_argument('--device_num', type=int, default=1, help='Device num.')
+parser.add_argument('--do_train', type=bool, default=True, help='Do train or not.')
+parser.add_argument('--do_eval', type=bool, default=False, help='Do eval or not.')
+parser.add_argument('--dataset_path', type=str, default=None, help='Dataset path')
+
+args_opt = parser.parse_args()
+device_id = int(os.getenv('DEVICE_ID'))
+
+context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=True, device_id=device_id)
+context.set_context(enable_task_sink=True)
+context.set_context(enable_loop_sink=True)
+context.set_context(enable_mem_reuse=True)
+
+
+def get_second_order_lr(global_step, lr_init, decay, total_epochs, steps_per_epoch):
+    """get_second_order_lr"""
+    lr_each_step = []
+    total_steps = steps_per_epoch * total_epochs
+    for i in range(total_steps):
+        epoch = (i + 1) / steps_per_epoch
+        base = (1.0 - float(epoch) / total_epochs) ** decay
+        lr_local = lr_init * base
+        lr_each_step.append(lr_local)
+    current_step = global_step
+    lr_each_step = np.array(lr_each_step).astype(np.float32)
+    print("learning_rate_is=====", lr_each_step)
+    learning_rate = lr_each_step[current_step:]
+    return learning_rate
+
+
+def get_second_order_damping(global_step, damping_init, decay_rate, total_epochs, steps_per_epoch):
+    """get_second_order_damping"""
+    damping_each_step = []
+    total_steps = steps_per_epoch * total_epochs
+    for step in range(total_steps):
+        epoch = (step + 1) / steps_per_epoch
+        damping_here = damping_init * (decay_rate ** (epoch / 10))
+        damping_each_step.append(damping_here)
+
+    current_step = global_step
+    damping_each_step = np.array(damping_each_step).astype(np.float32)
+    damping_now = damping_each_step[current_step:]
+    print("damping_is=========", damping_now)
+    return damping_now
+
+
+if __name__ == '__main__':
+    if args_opt.do_eval:
+        print("eval")
+    else:
+        if args_opt.run_distribute:
+            context.set_auto_parallel_context(device_num=args_opt.device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
+                                              mirror_mean=True, parameter_broadcast=True)
+            auto_parallel_context().set_all_reduce_fusion_split_indices([80], "hccl_world_groupsum1")
+            auto_parallel_context().set_all_reduce_fusion_split_indices([27], "hccl_world_groupsum3")
+            auto_parallel_context().set_all_reduce_fusion_split_indices([27], "hccl_world_groupsum4")
+            init()
+        else:
+            print(" ")
+
+    epoch_size = config.epoch_size
+    damping = get_second_order_damping(0, 0.03, 0.87, 50, 5004)
+    net = resnet50(class_num=config.class_num, damping=damping, loss_scale=config.loss_scale,
+                   frequency=config.frequency)
+
+    if not config.label_smooth:
+        config.label_smooth_factor = 0.0
+    loss = CrossEntropy(smooth_factor=config.label_smooth_factor, num_classes=config.class_num)
+    if args_opt.do_train:
+        dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=True,
+                                 repeat_num=epoch_size, batch_size=config.batch_size)
+        step_size = dataset.get_dataset_size()
+
+        loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False)
+        lr = Tensor(warmup_cosine_annealing_lr(0.035,
+                                               step_size,
+                                               config.warmup_epochs,
+                                               50,
+                                               config.T_max,
+                                               config.eta_min))
+        opt = THOR(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
+                   config.momentum, damping, config.frequency,
+                   filter(lambda x: 'matrix_A' in x.name, net.get_parameters()),
+                   filter(lambda x: 'matrix_G' in x.name, net.get_parameters()),
+                   filter(lambda x: 'spatial_norm' in x.name, net.get_parameters()),
+                   config.weight_decay, config.loss_scale)
+
+        model = Model(net, loss_fn=loss, optimizer=opt, amp_level='O2', loss_scale_manager=loss_scale,
+                      keep_batchnorm_fp32=False, metrics={'acc'}, frequency=config.frequency)
+
+        time_cb = TimeMonitor(data_size=step_size)
+        loss_cb = LossMonitor()
+        cb = [time_cb, loss_cb]
+        if config.save_checkpoint:
+            config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_steps,
+                                         keep_checkpoint_max=config.keep_checkpoint_max)
+            ckpt_cb = ModelCheckpoint(prefix="resnet", directory=config.save_checkpoint_path, config=config_ck)
+            cb += [ckpt_cb]
+
+        model.train(epoch_size, dataset, callbacks=cb)