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: 30ffcd8a1f
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '30ffcd8a1f'
${ noResults }
mindspore2022/tests/ut/python/parallel/test_one_hot_net.py

310 lines
12 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 numpy as np

  16. 

  17. import mindspore as ms

  18. import mindspore.nn as nn

  19. from mindspore import Tensor, Parameter

  20. from mindspore import context

  21. from mindspore.common import dtype as mstype

  22. from mindspore.common.api import _executor

  23. from mindspore.nn.cell import Cell

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

  25. from mindspore.ops import composite as C

  26. from mindspore.ops import functional as F

  27. from mindspore.ops import operations as P

  28. from mindspore.train import Model, ParallelMode

  29. from tests.dataset_mock import MindData

  30. from tests.ut.python.ops.test_math_ops import VirtualLoss

  31. 

  32. device_num = 16

  33. device_id = 2

  34. 

  35. 

  36. class StrategyModel():

  37.     onehot_strategy = ((1, device_num), (), ())

  38.     twod_strategy = ((1, device_num),)

  39.     twod_strategy_m = ((device_num, 1),)

  40.     scalar_twod_strategy = ((), (1, device_num))

  41.     twod_scalar_strategy = ((1, device_num), ())

  42.     scalar_strategy = ((),)

  43.     oned_strategy = ((1,),)

  44.     scalar_scalar_strategy = ((), ())

  45.     twod_twod_strategy = ((1, device_num), (1, device_num))

  46.     twod_twodbc_strategy = ((1, device_num), (1, 1))

  47.     twodbc_twod_strategy = ((1, 1), (device_num, 1))

  48. 

  49. 

  50. class StrategyBatch():

  51.     onehot_strategy = ((device_num, 1), (), ())

  52.     twod_strategy = ((1, device_num),)

  53.     twod_strategy_m = ((device_num, 1),)

  54.     scalar_twod_strategy = ((), (1, device_num))

  55.     twod_scalar_strategy = ((1, device_num), ())

  56.     scalar_strategy = ((),)

  57.     oned_strategy = ((1,),)

  58.     scalar_scalar_strategy = ((), ())

  59.     twod_twod_strategy = ((1, device_num), (1, device_num))

  60.     twod_twodbc_strategy = ((1, device_num), (1, 1))

  61.     twodbc_twod_strategy = ((1, 1), (device_num, 1))

  62. 

  63. 

  64. class Args():

  65.     a = 1

  66.     b = 2

  67.     c = 3

  68.     d = 4

  69.     e = 5

  70.     num_classes = 512

  71.     emb_size = 512

  72. 

  73. 

  74. class SemiAutoOneHotNet(Cell):

  75.     def __init__(self, args, strategy):

  76.         super(SemiAutoOneHotNet, self).__init__()

  77.         self.a = args.a

  78.         self.b = args.b

  79.         self.c = args.c

  80.         self.d = args.d

  81.         self.e = args.e

  82.         self.cast = P.Cast()

  83.         self.cast.set_strategy(strategy=strategy.twod_strategy)

  84.         self.cast1 = P.Cast()

  85.         self.cast1.set_strategy(strategy=strategy.twod_strategy)

  86.         self.cast2 = P.Cast()

  87.         self.cast2.set_strategy(strategy=strategy.twod_strategy)

  88.         self.cast3 = P.Cast()

  89.         self.cast3.set_strategy(strategy=strategy.scalar_strategy)

  90.         self.cast4 = P.Cast()

  91.         self.cast4.set_strategy(strategy=strategy.scalar_strategy)

  92.         self.a_const = Tensor(self.a, dtype=mstype.float32)

  93.         self.b_const = Tensor(self.b, dtype=mstype.float32)

  94.         self.c_const = Tensor(self.c, dtype=mstype.float32)

  95.         self.d_const = Tensor(self.d, dtype=mstype.float32)

  96.         self.e_const = Tensor(self.e, dtype=mstype.float32)

  97.         self.m_const_zero = Tensor(0, dtype=mstype.float32)

  98.         self.a_const_one = Tensor(1, dtype=mstype.float32)

  99.         self.onehot = P.OneHot()

  100.         self.onehot.set_strategy(strategy=strategy.onehot_strategy)

  101.         self.exp = P.Exp()

  102.         self.exp.set_strategy(strategy=strategy.twod_strategy)

  103.         self.exp2 = P.Exp()

  104.         self.exp2.set_strategy(strategy=strategy.twod_strategy)

  105.         self.exp3 = P.Exp()

  106.         self.exp3.set_strategy(strategy=strategy.twod_strategy)

  107.         self.mul_const = P.Mul()

  108.         self.mul_const.set_strategy(strategy=strategy.scalar_twod_strategy)

  109.         self.mul_const2 = P.TensorAdd()

  110.         self.mul_const2.set_strategy(strategy=strategy.scalar_twod_strategy)

  111.         self.mul_const3 = P.Sub()

  112.         self.mul_const3.set_strategy(strategy=strategy.twod_scalar_strategy)

  113.         self.mul_const4 = P.Sub()

  114.         self.mul_const4.set_strategy(strategy=strategy.scalar_twod_strategy)

  115.         self.mul_const5 = P.Mul()

  116.         self.mul_const5.set_strategy(strategy=strategy.twod_scalar_strategy)

  117.         self.mul = P.Mul()

  118.         self.mul.set_strategy(strategy=strategy.twod_twod_strategy)

  119.         self.mul2 = P.Mul()

  120.         self.mul2.set_strategy(strategy=strategy.twod_twod_strategy)

  121.         self.mul3 = P.TensorAdd()

  122.         self.mul3.set_strategy(strategy=strategy.twod_twod_strategy)

  123.         self.mul4 = P.Sub()

  124.         self.mul4.set_strategy(strategy=strategy.twod_twodbc_strategy)

  125.         self.mul5 = P.RealDiv()

  126.         self.mul5.set_strategy(strategy=strategy.twod_twodbc_strategy)

  127.         self.mul6 = P.Mul()

  128.         self.mul6.set_strategy(strategy=strategy.twod_twod_strategy)

  129.         self.mul7 = P.Mul()

  130.         self.mul7.set_strategy(strategy=strategy.twod_scalar_strategy)

  131.         self.mul8 = P.RealDiv()

  132.         self.mul8.set_strategy(strategy=strategy.scalar_scalar_strategy)

  133.         self.mul9 = P.TensorAdd()

  134.         self.mul9.set_strategy(strategy=strategy.twod_scalar_strategy)

  135. 

  136.         self.reduce_max = P.ReduceMax(keep_dims=True)

  137.         self.reduce_max.set_strategy(strategy=strategy.twod_strategy)

  138. 

  139.         self.reduce_sum = P.ReduceSum(keep_dims=False)

  140.         self.reduce_sum.set_strategy(strategy=strategy.twod_strategy)

  141.         self.reduce_sum_2 = P.ReduceSum(keep_dims=False)

  142.         self.reduce_sum_2.set_strategy(strategy=strategy.twod_strategy)

  143.         self.reduce_sum_3 = P.ReduceSum(keep_dims=False)

  144.         self.reduce_sum_3.set_strategy(strategy=strategy.oned_strategy)

  145. 

  146.         self.reshape = P.Reshape()

  147.         self.log = P.Log()

  148.         self.log.set_strategy(strategy=strategy.twod_strategy)

  149. 

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

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

  152.         self.normalize = P.L2Normalize(axis=1)

  153.         self.normalize.set_strategy(strategy=strategy.twod_strategy_m)

  154.         self.normalize2 = P.L2Normalize(axis=1)

  155.         self.normalize2.set_strategy(strategy=strategy.twod_strategy_m)

  156.         self.fc = P.MatMul(transpose_b=True)

  157.         self.fc.set_strategy(strategy=strategy.twodbc_twod_strategy)

  158.         weight_shape = [args.num_classes, args.emb_size]

  159.         weight_np = np.zeros(weight_shape, np.float32)

  160.         self.weight = Parameter(Tensor(weight_np), name='model_parallel_weight')

  161. 

  162.     def construct(self, input_, label):

  163.         input_n = self.normalize(input_)

  164.         w = self.normalize2(self.weight)

  165.         fc_o = self.fc(input_n, w)

  166.         fc_o_shape = F.shape(fc_o)

  167.         one_hot_float = self.onehot(label, fc_o_shape[1], self.on_value, self.off_value)

  168.         local_label = self.cast(one_hot_float, mstype.int32)

  169. 

  170.         exp_o = self.exp(fc_o)

  171.         mul_const_o = self.mul_const(self.a_const, exp_o)

  172.         mul_const2_o = self.mul_const2(self.b_const, mul_const_o)

  173.         exp2_o = self.exp2(mul_const2_o)

  174.         mul_const3_o = self.mul_const3(exp2_o, self.c_const)

  175.         mul_const4_o = self.mul_const4(F.scalar_to_array(1), local_label)

  176.         mul6_o = self.mul6(self.mul(mul_const3_o, one_hot_float),

  177.                            self.mul2(fc_o, self.cast2(mul_const4_o, mstype.float32)))

  178.         mul_const5_o = self.mul_const5(mul6_o, self.d_const)

  179. 

  180.         max_o = self.reduce_max(mul_const5_o, -1)

  181.         mul4_o = self.mul4(mul_const5_o, max_o)

  182.         exp3_o = self.exp3(mul4_o)

  183.         sum_o = self.reduce_sum(exp3_o, -1)

  184.         reshape_o = self.reshape(sum_o, (F.shape(sum_o)[0], 1))

  185.         mul5_o = self.mul5(exp3_o, reshape_o)

  186.         log_o = self.log(self.mul9(mul5_o, self.e_const))

  187.         mul3_o = self.mul3(log_o, one_hot_float)

  188.         mul7_o = self.mul7(mul3_o, self.cast3(F.scalar_to_array(-1), mstype.float32))

  189.         sum2_o = self.reduce_sum_2(mul7_o, -1)

  190.         loss = self.mul8(self.reduce_sum_3(sum2_o, -1),

  191.                          self.cast4(F.scalar_to_array(F.shape(mul_const5_o)[0]), mstype.float32))

  192.         return loss

  193. 

  194. 

  195. class Dataset(MindData):

  196.     def __init__(self, predict, label, length=3, input_num=2):

  197.         super(Dataset, self).__init__(size=length)

  198.         self.predict = predict

  199.         self.label = label

  200.         self.index = 0

  201.         self.length = length

  202.         self.input_num = input_num

  203. 

  204.     def __iter__(self):

  205.         return self

  206. 

  207.     def __next__(self):

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

  209.             raise StopIteration

  210.         self.index += 1

  211.         if self.input_num == 2:

  212.             return (self.predict, self.label)

  213.         return (self.predict,)

  214. 

  215.     def reset(self):

  216.         self.index = 0

  217. 

  218. 

  219. class NetWithLoss(nn.Cell):

  220.     def __init__(self, network):

  221.         super(NetWithLoss, self).__init__()

  222.         self.loss = VirtualLoss()

  223.         self.network = network

  224. 

  225.     def construct(self, x, b):

  226.         predict = self.network(x, b)

  227.         return self.loss(predict)

  228. 

  229. 

  230. class GradWrap(nn.Cell):

  231.     def __init__(self, network):

  232.         super(GradWrap, self).__init__()

  233.         self.network = network

  234. 

  235.     def construct(self, x, b):

  236.         return C.grad_all(self.network)(x, b)

  237. 

  238. 

  239. def bn_with_initialize(out_channels):

  240.     bn = nn.BatchNorm2d(out_channels, momentum=0.3, eps=1e-5).add_flags_recursive(fp32=True)

  241.     return bn

  242. 

  243. 

  244. def fc_with_initialize(input_channels, out_channels):

  245.     return nn.Dense(input_channels, out_channels)

  246. 

  247. 

  248. class BNReshapeDenseBNNet(nn.Cell):

  249.     def __init__(self):

  250.         super(BNReshapeDenseBNNet, self).__init__()

  251.         self.batch_norm = bn_with_initialize(2)

  252.         self.reshape = P.Reshape()

  253.         self.batch_norm2 = nn.BatchNorm1d(512, affine=False)

  254.         self.fc = fc_with_initialize(2 * 32 * 32, 512)

  255.         self.loss = SemiAutoOneHotNet(args=Args(), strategy=StrategyBatch())

  256. 

  257.     def construct(self, x, label):

  258.         x = self.batch_norm(x)

  259.         x = self.reshape(x, (16, 2 * 32 * 32))

  260.         x = self.fc(x)

  261.         x = self.batch_norm2(x)

  262.         loss = self.loss(x, label)

  263.         return loss

  264. 

  265. 

  266. def test_bn_reshape_dense_bn_train_loss():

  267.     batch_size = 16

  268.     context.set_auto_parallel_context(device_num=device_num, global_rank=0)

  269.     input_ = Tensor(np.ones([batch_size, 2, 32, 32]).astype(np.float32) * 0.01)

  270.     label = Tensor(np.ones([batch_size]), dtype=ms.int32)

  271. 

  272.     net = GradWrap(NetWithLoss(BNReshapeDenseBNNet()))

  273.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  274.     net.set_auto_parallel()

  275. 

  276.     _executor.compile(net, input_, label)

  277. 

  278. 

  279. def test_semi_one_hot_net_batch():

  280.     batch_size = 16

  281.     context.set_auto_parallel_context(device_num=device_num, global_rank=0)

  282.     input_ = Tensor(np.ones([batch_size * 1, 512]).astype(np.float32) * 0.01)

  283.     label = Tensor(np.ones([batch_size]), dtype=ms.int32)

  284. 

  285.     net = SemiAutoOneHotNet(args=Args(), strategy=StrategyBatch())

  286.     net = GradWrap(NetWithLoss(net))

  287.     context.set_auto_parallel_context(parallel_mode="semi_auto_parallel")

  288.     net.set_auto_parallel()

  289. 

  290.     _executor.compile(net, input_, label)

  291. 

  292. 

  293. def test_semi_one_hot_net_model():

  294.     batch_size = 16

  295.     learning_rate = 0.1

  296.     momentum = 0.9

  297.     epoch_size = 2

  298. 

  299.     predict = Tensor(np.ones([batch_size, 512]), dtype=ms.float32)

  300.     label = Tensor(np.ones([batch_size]), dtype=ms.int32)

  301.     dataset = Dataset(predict, label, 2, input_num=2)

  302. 

  303.     net = SemiAutoOneHotNet(args=Args(), strategy=StrategyModel())

  304.     opt = Momentum(net.trainable_params(), learning_rate, momentum)

  305.     context.reset_auto_parallel_context()

  306.     context.set_auto_parallel_context(parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL, device_num=16)

  307.     context.set_context(mode=context.GRAPH_MODE)

  308.     model = Model(net, optimizer=opt)

  309.     model.train(epoch_size, dataset, dataset_sink_mode=False)