new op type of Pangu-alpha supported in D-Rec

This commit is contained in:
haoran.wang 2021-11-29 09:48:36 +00:00
parent 8bdcc68bb7
commit e331f1d7ef
2 changed files with 10 additions and 5 deletions

View File

@ -51,6 +51,8 @@ enum OperatorType {
kRecReduce,
kRecPReLU,
kRecGatherV2,
kRecExpandDims,
kRecStridedSlice,
kRecArgWithValue,
kRecUnsortedSegmentOp
};

View File

@ -30,11 +30,11 @@
namespace mindspore {
namespace parallel {
static const std::set<OperatorType> ElementWiseOpType = {
OperatorType::kRecReLU, OperatorType::kRecLog, OperatorType::kRecExp, OperatorType::kRecAdd,
OperatorType::kRecElmWiseOp, OperatorType::kRecBiasAdd, OperatorType::kRecSub, OperatorType::kRecMul,
OperatorType::kRecDiv, OperatorType::kRecSqueeze, OperatorType::kRecReduce, OperatorType::kRecCast,
OperatorType::kRecReshape, OperatorType::kRecGatherV2, OperatorType::kRecArgWithValue, OperatorType::kRecSoftmax,
OperatorType::kRecOneHot};
OperatorType::kRecReLU, OperatorType::kRecLog, OperatorType::kRecExp, OperatorType::kRecAdd,
OperatorType::kRecElmWiseOp, OperatorType::kRecBiasAdd, OperatorType::kRecSub, OperatorType::kRecMul,
OperatorType::kRecDiv, OperatorType::kRecSqueeze, OperatorType::kRecReduce, OperatorType::kRecCast,
OperatorType::kRecReshape, OperatorType::kRecGatherV2, OperatorType::kRecArgWithValue, OperatorType::kRecSoftmax,
OperatorType::kRecOneHot, OperatorType::kRecExpandDims, OperatorType::kRecStridedSlice};
const std::map<std::string, OperatorType> DictOpType{
{MATMUL, OperatorType::kRecMatMul},
@ -57,6 +57,8 @@ const std::map<std::string, OperatorType> DictOpType{
{REDUCE_MIN, OperatorType::kRecReduce},
{REDUCE_MEAN, OperatorType::kRecReduce},
{GATHERV2, OperatorType::kRecGatherV2},
{EXPAND_DIMS, OperatorType::kRecExpandDims},
{STRIDEDSLICE, OperatorType::kRecStridedSlice},
{ARGMAXWITHVALUE, OperatorType::kRecArgWithValue},
{ARGMINWITHVALUE, OperatorType::kRecArgWithValue},
{UNSORTED_SEGMENT_SUM, OperatorType::kRecUnsortedSegmentOp},
@ -145,6 +147,7 @@ const std::map<std::string, OperatorType> DictOpType{
{ASSIGN_SUB, OperatorType::kRecElmWiseOp},
{"AssignAdd", OperatorType::kRecElmWiseOp},
{DROPOUT_DO_MASK, OperatorType::kRecElmWiseOp},
{DROPOUT, OperatorType::kRecElmWiseOp},
{STACK, OperatorType::kRecElmWiseOp}};
const TensorParam MakeTensor(int64_t n, int64_t c, int64_t h, int64_t w);