zzy34407230
/
mindspore2022
Not watched
1
0
Fork 0
Code
Releases 22 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 9cee0d2143
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '9cee0d2143'
${ noResults }
mindspore2022/tests/ut/python/parallel/test_auto_parallel_resnet.py

735 lines
31 kB
Raw Blame History 

  1. # Copyright 2019 Huawei Technologies Co., Ltd

  2. #

  3. # Licensed under the Apache License, Version 2.0 (the "License");

  4. # you may not use this file except in compliance with the License.

  5. # You may obtain a copy of the License at

  6. #

  7. # http://www.apache.org/licenses/LICENSE-2.0

  8. #

  9. # Unless required by applicable law or agreed to in writing, software

  10. # distributed under the License is distributed on an "AS IS" BASIS,

  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  12. # See the License for the specific language governing permissions and

  13. # limitations under the License.

  14. 

  15. import re

  16. import numpy as np

  17. 

  18. import mindspore.common.dtype as mstype

  19. import mindspore.nn as nn

  20. import mindspore.ops.functional as F

  21. from mindspore import Tensor

  22. from mindspore import context

  23. from mindspore.common.api import _executor

  24. from mindspore.common.initializer import TruncatedNormal

  25. from mindspore.communication.management import init

  26. from mindspore.nn.loss.loss import _Loss

  27. from mindspore.nn.optim.momentum import Momentum

  28. from mindspore.ops import operations as P

  29. from mindspore.parallel import _cost_model_context as cost_model_context

  30. from mindspore.parallel import set_algo_parameters

  31. from mindspore.parallel._utils import _reset_op_id as resset_op_id

  32. from mindspore.train.model import Model

  33. from mindspore.context import ParallelMode

  34. 

  35. context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")

  36. context.set_context(device_id=0)

  37. init()

  38. 

  39. 

  40. def weight_variable():

  41.     return TruncatedNormal(0.02)

  42. 

  43. 

  44. def _conv3x3(in_channels, out_channels, stride=1, padding=0, pad_mode='same'):

  45.     """Get a conv2d layer with 3x3 kernel size."""

  46.     init_value = weight_variable()

  47.     return nn.Conv2d(in_channels, out_channels,

  48.                      kernel_size=3, stride=stride, padding=padding, pad_mode=pad_mode, weight_init=init_value)

  49. 

  50. 

  51. def _conv1x1(in_channels, out_channels, stride=1, padding=0, pad_mode='same'):

  52.     """Get a conv2d layer with 1x1 kernel size."""

  53.     init_value = weight_variable()

  54.     return nn.Conv2d(in_channels, out_channels,

  55.                      kernel_size=1, stride=stride, padding=padding, pad_mode=pad_mode, weight_init=init_value)

  56. 

  57. 

  58. def _conv7x7(in_channels, out_channels, stride=1, padding=0, pad_mode='same'):

  59.     """Get a conv2d layer with 7x7 kernel size."""

  60.     init_value = weight_variable()

  61.     return nn.Conv2d(in_channels, out_channels,

  62.                      kernel_size=7, stride=stride, padding=padding, pad_mode=pad_mode, weight_init=init_value)

  63. 

  64. 

  65. def _fused_bn(channels, momentum=0.9):

  66.     """Get a fused batchnorm"""

  67.     return nn.BatchNorm2d(channels, momentum=momentum)

  68. 

  69. 

  70. class ResidualBlock(nn.Cell):

  71.     expansion = 4

  72. 

  73.     def __init__(self,

  74.                  in_channels,

  75.                  out_channels,

  76.                  stride=1,

  77.                  momentum=0.9):

  78.         super(ResidualBlock, self).__init__()

  79. 

  80.         out_chls = out_channels // self.expansion

  81.         self.conv1 = _conv1x1(in_channels, out_chls, stride=1)

  82.         self.bn1 = _fused_bn(out_chls, momentum=momentum)

  83. 

  84.         self.conv2 = _conv3x3(out_chls, out_chls, stride=stride)

  85.         self.bn2 = _fused_bn(out_chls, momentum=momentum)

  86. 

  87.         self.conv3 = _conv1x1(out_chls, out_channels, stride=1)

  88.         self.bn3 = _fused_bn(out_channels, momentum=momentum)

  89. 

  90.         self.relu = P.ReLU()

  91.         self.downsample = (in_channels != out_channels)

  92.         self.stride = stride

  93.         if self.downsample:

  94.             self.conv_down_sample = _conv1x1(in_channels, out_channels,

  95.                                              stride=stride)

  96.             self.bn_down_sample = _fused_bn(out_channels, momentum=momentum)

  97.         elif self.stride != 1:

  98.             self.maxpool_down = nn.MaxPool2d(kernel_size=1, stride=2, pad_mode='same')

  99. 

  100.         self.add = P.TensorAdd()

  101. 

  102.     def construct(self, x):

  103.         identity = x

  104. 

  105.         out = self.conv1(x)

  106.         out = self.bn1(out)

  107.         out = self.relu(out)

  108. 

  109.         out = self.conv2(out)

  110.         out = self.bn2(out)

  111.         out = self.relu(out)

  112. 

  113.         out = self.conv3(out)

  114.         out = self.bn3(out)

  115. 

  116.         if self.downsample:

  117.             identity = self.conv_down_sample(identity)

  118.             identity = self.bn_down_sample(identity)

  119.         elif self.stride != 1:

  120.             identity = self.maxpool_down(identity)

  121. 

  122.         out = self.add(out, identity)

  123.         out = self.relu(out)

  124. 

  125.         return out

  126. 

  127. 

  128. class ResNet(nn.Cell):

  129.     def __init__(self,

  130.                  block,

  131.                  layer_nums,

  132.                  in_channels,

  133.                  out_channels,

  134.                  strides=None,

  135.                  num_classes=100):

  136.         super(ResNet, self).__init__()

  137.         if strides is None:

  138.             strides = [1, 2, 2, 2]

  139.         if not len(layer_nums) == len(in_channels) == len(out_channels) == 4:

  140.             raise ValueError("the length of "

  141.                              "layer_num, inchannel, outchannel list must be 4!")

  142. 

  143.         self.conv1 = _conv7x7(3, 64, stride=2)

  144.         self.bn1 = _fused_bn(64)

  145.         self.relu = P.ReLU()

  146.         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')

  147. 

  148.         self.layer1 = self._make_layer(block,

  149.                                        layer_nums[0],

  150.                                        in_channel=in_channels[0],

  151.                                        out_channel=out_channels[0],

  152.                                        stride=strides[0])

  153.         self.layer2 = self._make_layer(block,

  154.                                        layer_nums[1],

  155.                                        in_channel=in_channels[1],

  156.                                        out_channel=out_channels[1],

  157.                                        stride=strides[1])

  158.         self.layer3 = self._make_layer(block,

  159.                                        layer_nums[2],

  160.                                        in_channel=in_channels[2],

  161.                                        out_channel=out_channels[2],

  162.                                        stride=strides[2])

  163.         self.layer4 = self._make_layer(block,

  164.                                        layer_nums[3],

  165.                                        in_channel=in_channels[3],

  166.                                        out_channel=out_channels[3],

  167.                                        stride=strides[3])

  168. 

  169.         self.mean = P.ReduceMean(keep_dims=True)

  170.         self.end_point = nn.Dense(2048, num_classes, has_bias=True,

  171.                                   weight_init=weight_variable(),

  172.                                   bias_init=weight_variable()).add_flags_recursive(fp16=True)

  173.         self.squeeze = P.Squeeze()

  174.         self.cast = P.Cast()

  175. 

  176.     def _make_layer(self, block, layer_num, in_channel, out_channel, stride):

  177.         layers = []

  178. 

  179.         resblk = block(in_channel, out_channel, stride=1)

  180.         layers.append(resblk)

  181. 

  182.         for _ in range(1, layer_num - 1):

  183.             resblk = block(out_channel, out_channel, stride=1)

  184.             layers.append(resblk)

  185. 

  186.         resblk = block(out_channel, out_channel, stride=stride)

  187.         layers.append(resblk)

  188. 

  189.         return nn.SequentialCell(layers)

  190. 

  191.     def construct(self, x):

  192.         x = self.conv1(x)

  193.         x = self.bn1(x)

  194.         x = self.relu(x)

  195.         c1 = self.maxpool(x)

  196.         c2 = self.layer1(c1)

  197.         c3 = self.layer2(c2)

  198.         c4 = self.layer3(c3)

  199.         c5 = self.layer4(c4)

  200.         out = self.mean(c5, (2, 3))

  201.         out = self.squeeze(out)

  202.         out = self.end_point(out)

  203. 

  204.         return out

  205. 

  206. 

  207. def resnet50(class_num=10):

  208.     return ResNet(ResidualBlock,

  209.                   [3, 4, 6, 3],

  210.                   [64, 256, 512, 1024],

  211.                   [256, 512, 1024, 2048],

  212.                   [2, 2, 2, 1],

  213.                   class_num)

  214. 

  215. 

  216. class SoftmaxCrossEntropyExpand(_Loss):

  217.     def __init__(self, sparse=False):

  218.         super(SoftmaxCrossEntropyExpand, self).__init__()

  219.         self.exp = P.Exp()

  220.         self.sum = P.ReduceSum(keep_dims=True)

  221.         self.onehot = P.OneHot()

  222.         self.on_value = Tensor(1.0, mstype.float32)

  223.         self.off_value = Tensor(0.0, mstype.float32)

  224.         self.div = P.Div()

  225.         self.log = P.Log()

  226.         self.sum_cross_entropy = P.ReduceSum(keep_dims=False)

  227.         self.mul = P.Mul()

  228.         self.mul2 = P.Mul()

  229.         self.cast = P.Cast()

  230.         self.mean = P.ReduceMean(keep_dims=False).add_prim_attr("cross_batch", True)

  231.         self.sparse = sparse

  232.         self.max = P.ReduceMax(keep_dims=True)

  233.         self.sub = P.Sub()

  234.         self.cast1 = P.Cast()

  235. 

  236.     def construct(self, logit, label):

  237.         logit = self.cast1(logit, mstype.float32)

  238.         logit_max = self.max(logit)

  239.         exp = self.exp(self.sub(logit, logit_max))

  240.         exp_sum = self.sum(exp, -1)

  241.         softmax_result = self.div(exp, exp_sum)

  242.         if self.sparse:

  243.             label = self.onehot(label, F.shape(logit)[1], self.on_value, self.off_value)

  244. 

  245.         softmax_result_log = self.log(softmax_result)

  246.         loss = self.sum_cross_entropy((self.mul(softmax_result_log, label)), -1)

  247.         loss = self.mul2(F.scalar_to_array(-1.0), loss)

  248.         loss = self.mean(loss, -1)

  249. 

  250.         return loss

  251. 

  252. 

  253. class DatasetLenet():

  254.     def __init__(self, predict, label, length=3):

  255.         self.predict = predict

  256.         self.label = label

  257.         self.index = 0

  258.         self.length = length

  259. 

  260.     def __iter__(self):

  261.         return self

  262. 

  263.     def __next__(self):

  264.         if self.index >= self.length:

  265.             raise StopIteration

  266.         self.index += 1

  267.         return self.predict, self.label

  268. 

  269.     def reset(self):

  270.         self.index = 0

  271. 

  272.     def get_dataset_size(self):

  273.         return 32

  274. 

  275.     def get_repeat_count(self):

  276.         return 1

  277. 

  278.     def create_tuple_iterator(self, num_epochs=-1):

  279.         return self

  280. 

  281. 

  282. def test_train_32k_8p(batch_size=32, num_classes=32768):

  283.     dev_num = 8

  284.     context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL, device_num=dev_num)

  285.     set_algo_parameters(elementwise_op_strategy_follow=True)

  286.     resset_op_id()

  287.     np.random.seed(6)

  288.     input_np = np.ones([batch_size, 3, 224, 224]).astype(np.float32)

  289.     label_np = np.zeros([batch_size]).astype(np.int32)

  290.     for i in range(0, batch_size):

  291.         label_np[i] = i % num_classes

  292.     dataset = DatasetLenet(Tensor(input_np), Tensor(label_np), 1)

  293.     net = resnet50(num_classes)

  294.     loss = SoftmaxCrossEntropyExpand(sparse=True)

  295.     opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)

  296.     model = Model(net, loss_fn=loss, optimizer=opt)

  297.     model.train(5, dataset, dataset_sink_mode=False)

  298.     strategies = _executor._get_shard_strategy(model._train_network)

  299.     for (k, v) in strategies.items():

  300.         if re.search('Conv2D-op', k) is not None:

  301.             assert v[0][0] == dev_num

  302.         elif re.search('MatMul-op', k) is not None:

  303.             assert v == [[dev_num, 1], [1, 1]]

  304.         elif re.search('ReduceSum-op', k) is not None:

  305.             assert v == [[dev_num, 1]]

  306. 

  307.     allreduce_fusion_dict = _executor._get_allreduce_fusion(model._train_network)

  308.     print(allreduce_fusion_dict)

  309.     return allreduce_fusion_dict

  310. 

  311. 

  312. def train_32k_8p_fusion1(batch_size=32, num_classes=32768):  # 1048576 #131072 #32768 #8192

  313.     cost_model_context.set_cost_model_context(costmodel_gamma=0.001, costmodel_beta=400.0)

  314.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_algorithm=1)

  315.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_times=2)

  316.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_tail_percent=0.5)

  317.     allreduce_fusion_dict = test_train_32k_8p(batch_size, num_classes)

  318.     expect_dict = {'end_point.bias': 2,

  319.                    'end_point.weight': 2,

  320.                    'layer4.2.bn3.beta': 2,

  321.                    'layer4.2.bn3.gamma': 2,

  322.                    'layer4.2.conv3.weight': 2,

  323.                    'layer4.2.bn2.beta': 2,

  324.                    'layer4.2.bn2.gamma': 2,

  325.                    'layer4.2.conv2.weight': 2,

  326.                    'layer4.2.bn1.beta': 2,

  327.                    'layer4.2.bn1.gamma': 2,

  328.                    'layer4.2.conv1.weight': 2,

  329.                    'layer4.1.bn3.beta': 2,

  330.                    'layer4.1.bn3.gamma': 2,

  331.                    'layer4.1.conv3.weight': 2,

  332.                    'layer4.1.bn2.beta': 2,

  333.                    'layer4.1.bn2.gamma': 2,

  334.                    'layer4.1.conv2.weight': 2,

  335.                    'layer4.1.bn1.beta': 2,

  336.                    'layer4.1.bn1.gamma': 2,

  337.                    'layer4.1.conv1.weight': 2,

  338.                    'layer4.0.bn_down_sample.beta': 2,

  339.                    'layer4.0.bn_down_sample.gamma': 2,

  340.                    'layer4.0.conv_down_sample.weight': 2,

  341.                    'layer4.0.bn3.beta': 2,

  342.                    'layer4.0.bn3.gamma': 2,

  343.                    'layer4.0.conv3.weight': 2,

  344.                    'layer4.0.bn2.beta': 2,

  345.                    'layer4.0.bn2.gamma': 2,

  346.                    'layer4.0.conv2.weight': 2,

  347.                    'layer4.0.bn1.beta': 2,

  348.                    'layer4.0.bn1.gamma': 2,

  349.                    'layer4.0.conv1.weight': 2,

  350.                    'layer3.5.bn3.beta': 2,

  351.                    'layer3.5.bn3.gamma': 2,

  352.                    'layer3.5.conv3.weight': 2,

  353.                    'layer3.5.bn2.beta': 2,

  354.                    'layer3.5.bn2.gamma': 2,

  355.                    'layer3.5.conv2.weight': 2,

  356.                    'layer3.5.bn1.beta': 2,

  357.                    'layer3.5.bn1.gamma': 2,

  358.                    'layer3.5.conv1.weight': 2,

  359.                    'layer3.4.bn3.beta': 2,

  360.                    'layer3.4.bn3.gamma': 2,

  361.                    'layer3.4.conv3.weight': 2,

  362.                    'layer3.4.bn2.beta': 2,

  363.                    'layer3.4.bn2.gamma': 2,

  364.                    'layer3.4.conv2.weight': 2,

  365.                    'layer3.4.bn1.beta': 2,

  366.                    'layer3.4.bn1.gamma': 2,

  367.                    'layer3.4.conv1.weight': 2,

  368.                    'layer3.3.bn3.beta': 2,

  369.                    'layer3.3.bn3.gamma': 2,

  370.                    'layer3.3.conv3.weight': 2,

  371.                    'layer3.3.bn2.beta': 2,

  372.                    'layer3.3.bn2.gamma': 2,

  373.                    'layer3.3.conv2.weight': 2,

  374.                    'layer3.3.bn1.beta': 2,

  375.                    'layer3.3.bn1.gamma': 2,

  376.                    'layer3.3.conv1.weight': 2,

  377.                    'layer3.2.bn3.beta': 2,

  378.                    'layer3.2.bn3.gamma': 2,

  379.                    'layer3.2.conv3.weight': 2,

  380.                    'layer3.2.bn2.beta': 2,

  381.                    'layer3.2.bn2.gamma': 2,

  382.                    'layer3.2.conv2.weight': 2,

  383.                    'layer3.2.bn1.beta': 2,

  384.                    'layer3.2.bn1.gamma': 2,

  385.                    'layer3.2.conv1.weight': 2,

  386.                    'layer3.1.bn3.beta': 2,

  387.                    'layer3.1.bn3.gamma': 2,

  388.                    'layer3.1.conv3.weight': 2,

  389.                    'layer3.1.bn2.beta': 2,

  390.                    'layer3.1.bn2.gamma': 2,

  391.                    'layer3.1.conv2.weight': 2,

  392.                    'layer3.1.bn1.beta': 2,

  393.                    'layer3.1.bn1.gamma': 2,

  394.                    'layer3.1.conv1.weight': 2,

  395.                    'layer3.0.bn_down_sample.beta': 2,

  396.                    'layer3.0.bn_down_sample.gamma': 2,

  397.                    'layer3.0.conv_down_sample.weight': 2,

  398.                    'layer3.0.bn3.beta': 2,

  399.                    'layer3.0.bn3.gamma': 2,

  400.                    'layer3.0.conv3.weight': 2,

  401.                    'layer3.0.bn2.beta': 2,

  402.                    'layer3.0.bn2.gamma': 2,

  403.                    'layer3.0.conv2.weight': 2,

  404.                    'layer3.0.bn1.beta': 2,

  405.                    'layer3.0.bn1.gamma': 2,

  406.                    'layer3.0.conv1.weight': 2,

  407.                    'layer2.3.bn3.beta': 2,

  408.                    'layer2.3.bn3.gamma': 2,

  409.                    'layer2.3.conv3.weight': 2,

  410.                    'layer2.3.bn2.beta': 2,

  411.                    'layer2.3.bn2.gamma': 2,

  412.                    'layer2.3.conv2.weight': 2,

  413.                    'layer2.3.bn1.beta': 2,

  414.                    'layer2.3.bn1.gamma': 2,

  415.                    'layer2.3.conv1.weight': 2,

  416.                    'layer2.2.bn3.beta': 2,

  417.                    'layer2.2.bn3.gamma': 2,

  418.                    'layer2.2.conv3.weight': 2,

  419.                    'layer2.2.bn2.beta': 2,

  420.                    'layer2.2.bn2.gamma': 2,

  421.                    'layer2.2.conv2.weight': 2,

  422.                    'layer2.2.bn1.beta': 2,

  423.                    'layer2.2.bn1.gamma': 2,

  424.                    'layer2.2.conv1.weight': 2,

  425.                    'layer2.1.bn3.beta': 2,

  426.                    'layer2.1.bn3.gamma': 2,

  427.                    'layer2.1.conv3.weight': 2,

  428.                    'layer2.1.bn2.beta': 2,

  429.                    'layer2.1.bn2.gamma': 2,

  430.                    'layer2.1.conv2.weight': 2,

  431.                    'layer2.1.bn1.beta': 2,

  432.                    'layer2.1.bn1.gamma': 2,

  433.                    'layer2.1.conv1.weight': 2,

  434.                    'layer2.0.bn_down_sample.beta': 2,

  435.                    'layer2.0.bn_down_sample.gamma': 2,

  436.                    'layer2.0.conv_down_sample.weight': 2,

  437.                    'layer2.0.bn3.beta': 2,

  438.                    'layer2.0.bn3.gamma': 2,

  439.                    'layer2.0.conv3.weight': 2,

  440.                    'layer2.0.bn2.beta': 2,

  441.                    'layer2.0.bn2.gamma': 2,

  442.                    'layer2.0.conv2.weight': 2,

  443.                    'layer2.0.bn1.beta': 2,

  444.                    'layer2.0.bn1.gamma': 2,

  445.                    'layer2.0.conv1.weight': 2,

  446.                    'layer1.2.bn3.beta': 2,

  447.                    'layer1.2.bn3.gamma': 2,

  448.                    'layer1.2.conv3.weight': 2,

  449.                    'layer1.2.bn2.beta': 2,

  450.                    'layer1.2.bn2.gamma': 2,

  451.                    'layer1.2.conv2.weight': 2,

  452.                    'layer1.2.bn1.beta': 2,

  453.                    'layer1.2.bn1.gamma': 2,

  454.                    'layer1.2.conv1.weight': 2,

  455.                    'layer1.1.bn3.beta': 2,

  456.                    'layer1.1.bn3.gamma': 2,

  457.                    'layer1.1.conv3.weight': 2,

  458.                    'layer1.1.bn2.beta': 2,

  459.                    'layer1.1.bn2.gamma': 2,

  460.                    'layer1.1.conv2.weight': 2,

  461.                    'layer1.1.bn1.beta': 2,

  462.                    'layer1.1.bn1.gamma': 2,

  463.                    'layer1.1.conv1.weight': 2,

  464.                    'layer1.0.bn_down_sample.beta': 2,

  465.                    'layer1.0.bn_down_sample.gamma': 2,

  466.                    'layer1.0.conv_down_sample.weight': 2,

  467.                    'layer1.0.bn3.beta': 2,

  468.                    'layer1.0.bn3.gamma': 2,

  469.                    'layer1.0.conv3.weight': 2,

  470.                    'layer1.0.bn2.beta': 2,

  471.                    'layer1.0.bn2.gamma': 2,

  472.                    'layer1.0.conv2.weight': 2,

  473.                    'layer1.0.bn1.beta': 2,

  474.                    'layer1.0.bn1.gamma': 2,

  475.                    'layer1.0.conv1.weight': 2,

  476.                    'bn1.beta': 1,

  477.                    'bn1.gamma': 1,

  478.                    'conv1.weight': 1}

  479. 

  480.     assert allreduce_fusion_dict == expect_dict

  481.     cost_model_context.reset_cost_model_context()

  482. 

  483. 

  484. def train_32k_8p_fusion2(batch_size=32, num_classes=32768):  # 1048576 #131072 #32768 #8192

  485.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_algorithm=2)

  486.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_tail_time=0.1)

  487.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_allreduce_inherent_time=0.05)

  488.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_allreduce_bandwidth=0.000001)

  489.     cost_model_context.set_cost_model_context(costmodel_allreduce_fusion_computation_time_parameter=0.0000015)

  490.     allreduce_fusion_dict = test_train_32k_8p(batch_size, num_classes)

  491.     expect_dict = {'end_point.bias': 2,

  492.                    'end_point.weight': 2,

  493.                    'layer4.2.bn3.beta': 2,

  494.                    'layer4.2.bn3.gamma': 2,

  495.                    'layer4.2.conv3.weight': 2,

  496.                    'layer4.2.bn2.beta': 2,

  497.                    'layer4.2.bn2.gamma': 2,

  498.                    'layer4.2.conv2.weight': 2,

  499.                    'layer4.2.bn1.beta': 2,

  500.                    'layer4.2.bn1.gamma': 2,

  501.                    'layer4.2.conv1.weight': 2,

  502.                    'layer4.1.bn3.beta': 2,

  503.                    'layer4.1.bn3.gamma': 2,

  504.                    'layer4.1.conv3.weight': 2,

  505.                    'layer4.1.bn2.beta': 2,

  506.                    'layer4.1.bn2.gamma': 2,

  507.                    'layer4.1.conv2.weight': 2,

  508.                    'layer4.1.bn1.beta': 2,

  509.                    'layer4.1.bn1.gamma': 2,

  510.                    'layer4.1.conv1.weight': 2,

  511.                    'layer4.0.bn_down_sample.beta': 2,

  512.                    'layer4.0.bn_down_sample.gamma': 2,

  513.                    'layer4.0.conv_down_sample.weight': 2,

  514.                    'layer4.0.bn3.beta': 2,

  515.                    'layer4.0.bn3.gamma': 2,

  516.                    'layer4.0.conv3.weight': 2,

  517.                    'layer4.0.bn2.beta': 2,

  518.                    'layer4.0.bn2.gamma': 2,

  519.                    'layer4.0.conv2.weight': 2,

  520.                    'layer4.0.bn1.beta': 2,

  521.                    'layer4.0.bn1.gamma': 2,

  522.                    'layer4.0.conv1.weight': 2,

  523.                    'layer3.5.bn3.beta': 2,

  524.                    'layer3.5.bn3.gamma': 2,

  525.                    'layer3.5.conv3.weight': 2,

  526.                    'layer3.5.bn2.beta': 2,

  527.                    'layer3.5.bn2.gamma': 2,

  528.                    'layer3.5.conv2.weight': 2,

  529.                    'layer3.5.bn1.beta': 2,

  530.                    'layer3.5.bn1.gamma': 2,

  531.                    'layer3.5.conv1.weight': 2,

  532.                    'layer3.4.bn3.beta': 2,

  533.                    'layer3.4.bn3.gamma': 2,

  534.                    'layer3.4.conv3.weight': 2,

  535.                    'layer3.4.bn2.beta': 2,

  536.                    'layer3.4.bn2.gamma': 2,

  537.                    'layer3.4.conv2.weight': 2,

  538.                    'layer3.4.bn1.beta': 2,

  539.                    'layer3.4.bn1.gamma': 2,

  540.                    'layer3.4.conv1.weight': 2,

  541.                    'layer3.3.bn3.beta': 2,

  542.                    'layer3.3.bn3.gamma': 2,

  543.                    'layer3.3.conv3.weight': 2,

  544.                    'layer3.3.bn2.beta': 2,

  545.                    'layer3.3.bn2.gamma': 2,

  546.                    'layer3.3.conv2.weight': 2,

  547.                    'layer3.3.bn1.beta': 2,

  548.                    'layer3.3.bn1.gamma': 2,

  549.                    'layer3.3.conv1.weight': 2,

  550.                    'layer3.2.bn3.beta': 2,

  551.                    'layer3.2.bn3.gamma': 2,

  552.                    'layer3.2.conv3.weight': 2,

  553.                    'layer3.2.bn2.beta': 2,

  554.                    'layer3.2.bn2.gamma': 2,

  555.                    'layer3.2.conv2.weight': 2,

  556.                    'layer3.2.bn1.beta': 2,

  557.                    'layer3.2.bn1.gamma': 2,

  558.                    'layer3.2.conv1.weight': 2,

  559.                    'layer3.1.bn3.beta': 2,

  560.                    'layer3.1.bn3.gamma': 2,

  561.                    'layer3.1.conv3.weight': 2,

  562.                    'layer3.1.bn2.beta': 2,

  563.                    'layer3.1.bn2.gamma': 2,

  564.                    'layer3.1.conv2.weight': 2,

  565.                    'layer3.1.bn1.beta': 2,

  566.                    'layer3.1.bn1.gamma': 2,

  567.                    'layer3.1.conv1.weight': 2,

  568.                    'layer3.0.bn_down_sample.beta': 2,

  569.                    'layer3.0.bn_down_sample.gamma': 2,

  570.                    'layer3.0.conv_down_sample.weight': 2,

  571.                    'layer3.0.bn3.beta': 2,

  572.                    'layer3.0.bn3.gamma': 2,

  573.                    'layer3.0.conv3.weight': 2,

  574.                    'layer3.0.bn2.beta': 2,

  575.                    'layer3.0.bn2.gamma': 2,

  576.                    'layer3.0.conv2.weight': 2,

  577.                    'layer3.0.bn1.beta': 2,

  578.                    'layer3.0.bn1.gamma': 2,

  579.                    'layer3.0.conv1.weight': 2,

  580.                    'layer2.3.bn3.beta': 2,

  581.                    'layer2.3.bn3.gamma': 2,

  582.                    'layer2.3.conv3.weight': 2,

  583.                    'layer2.3.bn2.beta': 2,

  584.                    'layer2.3.bn2.gamma': 2,

  585.                    'layer2.3.conv2.weight': 2,

  586.                    'layer2.3.bn1.beta': 2,

  587.                    'layer2.3.bn1.gamma': 2,

  588.                    'layer2.3.conv1.weight': 2,

  589.                    'layer2.2.bn3.beta': 2,

  590.                    'layer2.2.bn3.gamma': 2,

  591.                    'layer2.2.conv3.weight': 2,

  592.                    'layer2.2.bn2.beta': 2,

  593.                    'layer2.2.bn2.gamma': 2,

  594.                    'layer2.2.conv2.weight': 2,

  595.                    'layer2.2.bn1.beta': 2,

  596.                    'layer2.2.bn1.gamma': 2,

  597.                    'layer2.2.conv1.weight': 2,

  598.                    'layer2.1.bn3.beta': 2,

  599.                    'layer2.1.bn3.gamma': 2,

  600.                    'layer2.1.conv3.weight': 2,

  601.                    'layer2.1.bn2.beta': 2,

  602.                    'layer2.1.bn2.gamma': 2,

  603.                    'layer2.1.conv2.weight': 2,

  604.                    'layer2.1.bn1.beta': 2,

  605.                    'layer2.1.bn1.gamma': 2,

  606.                    'layer2.1.conv1.weight': 2,

  607.                    'layer2.0.bn_down_sample.beta': 2,

  608.                    'layer2.0.bn_down_sample.gamma': 2,

  609.                    'layer2.0.conv_down_sample.weight': 2,

  610.                    'layer2.0.bn3.beta': 2,

  611.                    'layer2.0.bn3.gamma': 2,

  612.                    'layer2.0.conv3.weight': 2,

  613.                    'layer2.0.bn2.beta': 2,

  614.                    'layer2.0.bn2.gamma': 2,

  615.                    'layer2.0.conv2.weight': 2,

  616.                    'layer2.0.bn1.beta': 2,

  617.                    'layer2.0.bn1.gamma': 2,

  618.                    'layer2.0.conv1.weight': 2,

  619.                    'layer1.2.bn3.beta': 2,

  620.                    'layer1.2.bn3.gamma': 2,

  621.                    'layer1.2.conv3.weight': 2,

  622.                    'layer1.2.bn2.beta': 2,

  623.                    'layer1.2.bn2.gamma': 2,

  624.                    'layer1.2.conv2.weight': 2,

  625.                    'layer1.2.bn1.beta': 2,

  626.                    'layer1.2.bn1.gamma': 2,

  627.                    'layer1.2.conv1.weight': 2,

  628.                    'layer1.1.bn3.beta': 2,

  629.                    'layer1.1.bn3.gamma': 2,

  630.                    'layer1.1.conv3.weight': 2,

  631.                    'layer1.1.bn2.beta': 2,

  632.                    'layer1.1.bn2.gamma': 2,

  633.                    'layer1.1.conv2.weight': 2,

  634.                    'layer1.1.bn1.beta': 2,

  635.                    'layer1.1.bn1.gamma': 2,

  636.                    'layer1.1.conv1.weight': 2,

  637.                    'layer1.0.bn_down_sample.beta': 2,

  638.                    'layer1.0.bn_down_sample.gamma': 2,

  639.                    'layer1.0.conv_down_sample.weight': 2,

  640.                    'layer1.0.bn3.beta': 2,

  641.                    'layer1.0.bn3.gamma': 2,

  642.                    'layer1.0.conv3.weight': 2,

  643.                    'layer1.0.bn2.beta': 2,

  644.                    'layer1.0.bn2.gamma': 2,

  645.                    'layer1.0.conv2.weight': 1,

  646.                    'layer1.0.bn1.beta': 1,

  647.                    'layer1.0.bn1.gamma': 1,

  648.                    'layer1.0.conv1.weight': 1,

  649.                    'bn1.beta': 1,

  650.                    'bn1.gamma': 1,

  651.                    'conv1.weight': 1}

  652. 

  653.     assert allreduce_fusion_dict == expect_dict

  654.     cost_model_context.reset_cost_model_context()

  655. 

  656. 

  657. def test_train_64k_8p(batch_size=32, num_classes=65536):  # 1048576 #131072 #32768 #8192

  658.     dev_num = 8

  659.     context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL, device_num=dev_num)

  660.     cost_model_context.set_cost_model_context(costmodel_gamma=0.001, costmodel_beta=400.0)

  661.     set_algo_parameters(elementwise_op_strategy_follow=True)

  662.     resset_op_id()

  663.     np.random.seed(6)

  664.     input_np = np.ones([batch_size, 3, 224, 224]).astype(np.float32)

  665.     label_np = np.zeros([batch_size]).astype(np.int32)

  666.     for i in range(0, batch_size):

  667.         label_np[i] = i % num_classes

  668.     dataset = DatasetLenet(Tensor(input_np), Tensor(label_np), 1)

  669.     net = resnet50(num_classes)

  670.     loss = SoftmaxCrossEntropyExpand(sparse=True)

  671.     opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)

  672.     model = Model(net, loss_fn=loss, optimizer=opt)

  673.     model.train(5, dataset, dataset_sink_mode=False)

  674.     strategies = _executor._get_shard_strategy(model._train_network)

  675.     for (k, v) in strategies.items():

  676.         if re.search('Conv2D-op', k) is not None:

  677.             assert v[0][0] == dev_num

  678.         elif re.search('MatMul-op', k) is not None:

  679.             assert v == [[1, 1], [dev_num, 1]]

  680.         elif re.search('ReduceSum-op', k) is not None:

  681.             assert v == [[1, dev_num]]

  682. 

  683. 

  684. def test_train_8k_8p_gpu(batch_size=32, num_classes=8192):

  685.     dev_num = 8

  686.     context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

  687.     context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL, device_num=dev_num)

  688.     set_algo_parameters(elementwise_op_strategy_follow=True)

  689.     resset_op_id()

  690.     np.random.seed(6)

  691.     input_np = np.ones([batch_size, 3, 224, 224]).astype(np.float32)

  692.     label_np = np.zeros([batch_size]).astype(np.int32)

  693.     for i in range(0, batch_size):

  694.         label_np[i] = i % num_classes

  695.     dataset = DatasetLenet(Tensor(input_np), Tensor(label_np), 1)

  696.     net = resnet50(num_classes)

  697.     loss = SoftmaxCrossEntropyExpand(sparse=True)

  698.     opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)

  699.     model = Model(net, loss_fn=loss, optimizer=opt)

  700.     model.train(5, dataset, dataset_sink_mode=False)

  701.     strategies = _executor._get_shard_strategy(model._train_network)

  702.     for (k, v) in strategies.items():

  703.         if re.search('Conv2D-op', k) is not None:

  704.             assert v[0][0] == dev_num

  705.         elif re.search('MatMul-op', k) is not None:

  706.             assert v == [[1, 1], [dev_num, 1]]

  707.         elif re.search('ReduceSum-op', k) is not None:

  708.             assert v == [[1, dev_num]]

  709. 

  710. def test_train_4k_8p_gpu(batch_size=32, num_classes=4096):

  711.     dev_num = 8

  712.     context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

  713.     context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL, device_num=dev_num)

  714.     set_algo_parameters(elementwise_op_strategy_follow=True)

  715.     resset_op_id()

  716.     np.random.seed(6)

  717.     input_np = np.ones([batch_size, 3, 224, 224]).astype(np.float32)

  718.     label_np = np.zeros([batch_size]).astype(np.int32)

  719.     for i in range(0, batch_size):

  720.         label_np[i] = i % num_classes

  721.     dataset = DatasetLenet(Tensor(input_np), Tensor(label_np), 1)

  722.     net = resnet50(num_classes)

  723.     loss = SoftmaxCrossEntropyExpand(sparse=True)

  724.     opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)

  725.     model = Model(net, loss_fn=loss, optimizer=opt)

  726.     model.train(5, dataset, dataset_sink_mode=False)

  727.     strategies = _executor._get_shard_strategy(model._train_network)

  728.     for (k, v) in strategies.items():

  729.         if re.search('Conv2D-op', k) is not None:

  730.             assert v[0][0] == dev_num

  731.         elif re.search('MatMul-op', k) is not None:

  732.             assert v == [[dev_num, 1], [1, 1]]

  733.         elif re.search('ReduceSum-op', k) is not None:

  734.             assert v == [[dev_num, 1]]