1# Copyright 2020 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""" test dynamic shape """ 16import numpy as np 17 18from mindspore import Tensor, context, nn, Parameter 19from mindspore import dtype as mstype 20from mindspore.ops import operations as P 21 22context.set_context(mode=context.GRAPH_MODE) 23 24 25def test_sparse_apply_proximal_ada_grad(): 26 class Net(nn.Cell): 27 def __init__(self): 28 super(Net, self).__init__() 29 self.sparse_apply_proximal_adagrad = P.SparseApplyProximalAdagrad() 30 self.var = Parameter(Tensor(np.random.rand(7800, 80).astype(np.float32)), name="var") 31 self.accum = Parameter(Tensor(np.random.rand(7800, 80).astype(np.float32)), name="accum") 32 self.lr = 0.01 33 self.l1 = 0.0 34 self.l2 = 0.0 35 36 def construct(self, grad, indices): 37 out = self.sparse_apply_proximal_adagrad(self.var, self.accum, self.lr, self.l1, self.l2, grad, indices) 38 return out[0] 39 40 class NetWrapper(nn.Cell): 41 def __init__(self): 42 super(NetWrapper, self).__init__() 43 self.unq = P.Unique() 44 self.add = P.Add() 45 self.expand_dims = P.ExpandDims() 46 self.cast = P.Cast() 47 self.net = Net() 48 49 def construct(self, grad, inp): 50 ids, _ = self.unq(inp) 51 new_grad = self.expand_dims(ids, 1) 52 new_grad = self.cast(new_grad, mstype.float32) + grad 53 return self.net(new_grad, ids) 54 55 net = NetWrapper() 56 grad = Tensor(np.random.rand(1, 80).astype(np.float32)) 57 indices = Tensor(np.ones([7800]), mstype.int32) 58 net(grad, indices) 59 60 61def test_sparse_apply_ftrl(): 62 class SparseApplyFtrlNet(nn.Cell): 63 def __init__(self): 64 super(SparseApplyFtrlNet, self).__init__() 65 self.sparse_apply_ftrl = P.SparseApplyFtrl(lr=0.01, l1=0.0, l2=0.0, lr_power=-0.5) 66 self.var = Parameter(Tensor(np.random.rand(7800, 80).astype(np.float32)), name="var") 67 self.accum = Parameter(Tensor(np.random.rand(7800, 80).astype(np.float32)), name="accum") 68 self.linear = Parameter(Tensor(np.random.rand(7800, 80).astype(np.float32)), name="linear") 69 70 def construct(self, grad, indices): 71 out = self.sparse_apply_ftrl(self.var, self.accum, self.linear, grad, indices) 72 return out[0] 73 74 class NetWrapper(nn.Cell): 75 def __init__(self): 76 super(NetWrapper, self).__init__() 77 self.unq = P.Unique() 78 self.add = P.Add() 79 self.expand_dims = P.ExpandDims() 80 self.cast = P.Cast() 81 self.net = SparseApplyFtrlNet() 82 83 def construct(self, grad, inp): 84 ids, _ = self.unq(inp) 85 new_grad = self.expand_dims(ids, 1) 86 new_grad = self.cast(new_grad, mstype.float32) + grad 87 return self.net(new_grad, ids) 88 89 net = NetWrapper() 90 grad = Tensor(np.random.rand(1, 80).astype(np.float32)) 91 indices = Tensor(np.ones([7800]), mstype.int32) 92 net(grad, indices) 93 94 95def test_gatherv2(): 96 class Net(nn.Cell): 97 def __init__(self): 98 super(Net, self).__init__() 99 self.unq = P.Unique() 100 self.gather = P.Gather() 101 self.yy = Tensor(np.ones([8], dtype=np.int32)) 102 103 def construct(self, x, y): 104 shp = P.Shape()(self.yy) 105 y = P.Reshape()(y, shp) 106 u, _ = self.unq(y) 107 u_shp = P.DynamicShape()(u) 108 z = self.gather(x, u, 0) 109 return z, u_shp 110 111 x = Tensor(np.ones([20, 12], dtype=np.float32)) 112 y = Tensor(np.ones([2, 4], dtype=np.int32)) 113 net = Net() 114 net(x, y) 115 116 117def test_segmentsum(): 118 class Net(nn.Cell): 119 def __init__(self): 120 super(Net, self).__init__() 121 self.unq = P.Unique() 122 self.segment_ids = Tensor([0, 0, 1, 2, 1, 1, 1, 1], mstype.int32) 123 self.sum = P.UnsortedSegmentSum() 124 def construct(self, x): 125 u, _ = self.unq(x) 126 shp = P.DynamicShape()(u) 127 z = self.sum(x, self.segment_ids, shp[0]) 128 return z, shp[0] 129 130 x = Tensor(np.ones([8], dtype=np.int32)) 131 net = Net() 132 net(x) 133 134 135def test_addn(): 136 class Net(nn.Cell): 137 def __init__(self): 138 super(Net, self).__init__() 139 self.unq = P.Unique() 140 self.addn = P.AddN() 141 142 def construct(self, x): 143 u, _ = self.unq(x) 144 u = self.addn((u, u, u)) 145 z = self.addn([u, u]) 146 return z 147 148 y = Tensor(np.ones([8], dtype=np.int32)) 149 net = Net() 150 net(y) 151