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 15import numpy as np 16 17import mindspore.nn as nn 18from mindspore import Tensor, context 19from mindspore.common.api import _cell_graph_executor 20from mindspore.ops import operations as P 21from ....train_step_wrap import train_step_with_loss_warp 22 23 24class DenseMutMulNet(nn.Cell): 25 def __init__(self): 26 super(DenseMutMulNet, self).__init__() 27 self.fc1 = nn.Dense(128, 768, activation='relu') 28 self.fc2 = nn.Dense(128, 768, activation='relu') 29 self.fc3 = nn.Dense(128, 768, activation='relu') 30 self.fc4 = nn.Dense(768, 768, activation='relu') 31 self.relu4 = nn.ReLU() 32 self.relu5 = nn.ReLU() 33 self.transpose = P.Transpose() 34 self.matmul1 = P.MatMul() 35 self.matmul2 = P.MatMul() 36 37 def construct(self, x): 38 q = self.fc1(x) 39 k = self.fc2(x) 40 v = self.fc3(x) 41 k = self.transpose(k, (1, 0)) 42 c = self.relu4(self.matmul1(q, k)) 43 s = self.relu5(self.matmul2(c, v)) 44 s = self.fc4(s) 45 return s 46 47 48def test_dmnet_train_step(): 49 context.reset_auto_parallel_context() 50 input_ = Tensor(np.ones([32, 128]).astype(np.float32) * 0.01) 51 label = Tensor(np.zeros([32, 768]).astype(np.float32)) 52 net = DenseMutMulNet() 53 net = train_step_with_loss_warp(DenseMutMulNet()) 54 net.set_train() 55 _cell_graph_executor.compile(net, input_, label) 56