!220 Add parallel operator for DropoutDoMask

Merge pull request !220 from yangzhenzhang/dropoutdomask

mindspore/ccsrc/parallel/auto_parallel/operator_costmodel.cc

@@ -613,5 +613,15 @@ double ReduceMeanCost::GetForwardComputationCost(const std::vector<TensorInfo>&
 
   return result;
 }
+
+double DropOutCost::GetForwardComputationCost(const std::vector<TensorInfo>& inputs, const std::vector<TensorInfo>&,
+                                              const int32_t&) const {
+  if (inputs.empty()) {
+    return 0.0;
+  }
+  TensorInfo input0 = inputs[0];
+  Shape input0_slice_shape = input0.slice_shape();
+  return ListProduct(input0_slice_shape) * static_cast<double>(inputs_type_lengths_[0]) * DROPOUT_COST_RATE;
+}
 }  // namespace parallel
 }  // namespace mindspore

mindspore/ccsrc/parallel/auto_parallel/operator_costmodel.h

@@ -26,6 +26,7 @@ namespace mindspore {
 namespace parallel {
 #define MAXIMUM_INPUT_NUMBER 100
 #define DEFAULT_DATA_TYPE_LENGTH 4
+#define DROPOUT_COST_RATE 1.125  // the DropoutGenMask need 12.5% memory
 
 class OperatorCost;
 using OperatorCostPtr = std::shared_ptr<OperatorCost>;
@@ -493,6 +494,37 @@ class GetNextCost : public OperatorCost {
   }
 };
 using GetNextCostPtr = std::shared_ptr<GetNextCost>;
+
+class DropOutCost : public OperatorCost {
+ public:
+  DropOutCost() = default;
+  ~DropOutCost() override = default;
+
+  double GetCommCost(const std::vector<TensorInfo>& inputs, const std::vector<TensorInfo>& outputs,
+                     const int32_t& stage_id) const override {
+    return GetForwardCommCost(inputs, outputs, stage_id) + GetBackwardCommCost(inputs, outputs, stage_id);
+  }
+  double GetForwardCommCost(const std::vector<TensorInfo>&, const std::vector<TensorInfo>&,
+                            const int32_t&) const override {
+    return 0.0;
+  }
+  double GetBackwardCommCost(const std::vector<TensorInfo>&, const std::vector<TensorInfo>&,
+                             const int32_t&) const override {
+    return 0.0;
+  }
+  double GetComputationCost(const std::vector<TensorInfo>& inputs, const std::vector<TensorInfo>& outputs,
+                            const int32_t& stage_id) const override {
+    return GetForwardComputationCost(inputs, outputs, stage_id) + GetBackwardComputationCost(inputs, outputs, stage_id);
+  }
+  double GetForwardComputationCost(const std::vector<TensorInfo>&, const std::vector<TensorInfo>&,
+                                   const int32_t&) const override;
+  double GetBackwardComputationCost(const std::vector<TensorInfo>&, const std::vector<TensorInfo>&,
+                                    const int32_t&) const override {
+    return 0.0;
+  }
+};
+
+using DropOutCostPtr = std::shared_ptr<DropOutCost>;
 }  // namespace parallel
 }  // namespace mindspore
 #endif  // PARALLEL_AUTO_PARALLEL_OPERATOR_COSTMODEL_H_

mindspore/ccsrc/parallel/dynamic_creator.h

@@ -111,7 +111,6 @@ REGISTER(ReduceMinInfo);
 REGISTER(TransposeInfo);
 REGISTER(PReLUInfo);
 REGISTER(DropoutDoMaskInfo);
-REGISTER(DropoutGenMaskInfo)
 REGISTER(ReshapeInfo);
 REGISTER(FloorDivInfo);
 REGISTER(MaximumInfo);

mindspore/ccsrc/parallel/node_check.cc

@@ -71,6 +71,7 @@ const std::set<std::string> BLACK_LIST = {TUPLE_GETITEM,
                                           BROADCASTGRADIENTARGS,
                                           INVERTPERMUTATION,
                                           CONTROLDEPEND,
+                                          DROPOUT_GEN_MASK,
                                           EMBED,
                                           CREATINSTANCE,
                                           ZEROSLIKETENSOR,

mindspore/ccsrc/parallel/ops_info/dropout_do_mask_info.cc

 /**
- * Copyright 2019 Huawei Technologies Co., Ltd
+ * Copyright 2020 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "ir/value.h"
+#include "pipeline/resource.h"
 #include "parallel/auto_parallel/costmodel.h"
 #include "parallel/device_matrix.h"
 #include "parallel/strategy.h"
@@ -29,13 +30,32 @@
 
 namespace mindspore {
 namespace parallel {
+static int32_t SEED_NUM = 1;
+
 Status DropoutDoMaskInfo::CheckStrategy(const StrategyPtr& strategy) {
-  Shapes input_shape = {inputs_shape_.at(0)};
+  if (strategy == nullptr) {
+    MS_LOG(ERROR) << name_ << ": The strategy is null";
+    return FAILED;
+  }
+
+  std::vector<Dimensions> stra = strategy->GetInputDim();
+  if (stra.size() != 1) {
+    MS_LOG(ERROR) << name_ << ": Invalid strategy size " << stra.size() << ", it must be 1";
+    return FAILED;
+  }
+
+  if (inputs_shape_.empty()) {
+    MS_LOG(ERROR) << name_ << ": The inputs shape is empty";
+    return FAILED;
+  }
+
+  // only check the input[0]
+  Shapes input_shape = {inputs_shape_[0]};
   if (CheckStrategyValue(strategy, input_shape, is_auto_parallel_) != SUCCESS) {
     if (is_auto_parallel_) {
-      MS_LOG(DEBUG) << name_ << " : Invalid strategy.";
+      MS_LOG(DEBUG) << name_ << ": Invalid strategy";
     } else {
-      MS_LOG(ERROR) << name_ << " : Invalid strategy.";
+      MS_LOG(ERROR) << name_ << ": Invalid strategy";
     }
     return FAILED;
   }
@@ -43,68 +63,69 @@ Status DropoutDoMaskInfo::CheckStrategy(const StrategyPtr& strategy) {
 }
 
 Status DropoutDoMaskInfo::InferDevMatrixShape() {
-  std::vector<Dimensions> stra = strategy_->GetInputDim();
-  Dimensions input_strategy = stra.at(0);
+  if (strategy_ == nullptr) {
+    MS_LOG(ERROR) << name_ << ": The strategy is null";
+    return FAILED;
+  }
 
-  dev_matrix_shape_ = input_strategy;
+  std::vector<Dimensions> strategy = strategy_->GetInputDim();
+  if (strategy.empty()) {
+    MS_LOG(ERROR) << name_ << ": The strategy is empty";
+    return FAILED;
+  }
 
+  dev_matrix_shape_ = strategy[0];
   return SUCCESS;
 }
 
 Status DropoutDoMaskInfo::InferTensorMap() {
+  if (inputs_shape_.empty()) {
+    MS_LOG(ERROR) << name_ << ": The inputs shape is empty";
+    return FAILED;
+  }
+
   std::vector<int32_t> tensor_map_index;
-  size_t size = inputs_shape_.at(0).size();
-  // such as 4: tensor_map_index [3,2,1,0]
+  size_t size = inputs_shape_[0].size();
+  // if the dimension of input is 4, and tensor_map_index is [3, 2, 1, 0]
   for (size_t i = 0; i < size; ++i) {
-    tensor_map_index.push_back((int32_t)(LAST_INDEX(size) - i));
+    tensor_map_index.push_back(SizeToInt(size - i - 1));
   }
 
-  TensorMap input_b_tensor_map = {MAP_NONE};
-  inputs_tensor_map_.push_back(tensor_map_index);
-  inputs_tensor_map_.push_back(input_b_tensor_map);
-  outputs_tensor_map_.push_back(tensor_map_index);
+  // the input[1] do not need tensor map
+  inputs_tensor_map_.push_back(tensor_map_index);   // input_0
+  outputs_tensor_map_.push_back(tensor_map_index);  // output
   return SUCCESS;
 }
 
 Status DropoutDoMaskInfo::InferTensorInfo() {
-  // infer tensor shape
-  Shape input_a_shape = inputs_shape_.at(0);
-  Shape input_b_shape = inputs_shape_.at(1);
-  Shape output_shape = outputs_shape_.at(0);
-
-  // infer slice shape
-  Shapes inputs_slice_shape, outputs_slice_shape;
-  Strategys inputs_strategy = strategy_->GetInputDim();
-  Dimensions input_b_strategy = {1}, input_x_strategy = {};
-  inputs_strategy.emplace_back(input_b_strategy);
-  inputs_strategy.emplace_back(input_x_strategy);
-  Strategys outputs_strategy = {inputs_strategy.at(0)};
-  if (InferSliceShape(inputs_strategy, outputs_strategy, &inputs_slice_shape, &outputs_slice_shape) != SUCCESS) {
+  if (inputs_shape_.size() != 3) {
+    MS_LOG(ERROR) << name_ << ": Invalid inputs shape size " << inputs_shape_.size();
     return FAILED;
   }
-  Shape input_a_slice_shape = inputs_slice_shape.at(0);
-  Shape input_b_slice_shape = inputs_slice_shape.at(1);
-  Shape output_slice_shape = outputs_slice_shape.at(0);
 
-  TensorLayout input_a_tensor_layout, input_b_tensor_layout;
-  TensorLayout output_tensor_layout;
-  if (input_a_tensor_layout.InitFromVector(dev_matrix_shape_, inputs_tensor_map_[0], input_a_shape) != SUCCESS) {
+  if (strategy_ == nullptr) {
+    MS_LOG(ERROR) << name_ << ": The strategy is null";
     return FAILED;
   }
-  if (input_b_tensor_layout.InitFromVector(dev_matrix_shape_, inputs_tensor_map_[1], input_b_shape) != SUCCESS) {
+
+  Shape input_0_shape = inputs_shape_[0];
+
+  if (inputs_tensor_map_.empty()) {
+    MS_LOG(ERROR) << name_ << ": The inputs tensor map is empty";
     return FAILED;
   }
-  if (output_tensor_layout.InitFromVector(dev_matrix_shape_, outputs_tensor_map_[0], output_shape) != SUCCESS) {
+
+  TensorLayout input_0_tensor_layout;
+  if (input_0_tensor_layout.InitFromVector(dev_matrix_shape_, inputs_tensor_map_[0], input_0_shape) != SUCCESS) {
+    MS_LOG(ERROR) << name_ << ": Init tensor layout failed";
     return FAILED;
   }
-  TensorInfo input_a_tensor_info(input_a_tensor_layout, input_a_shape, input_a_slice_shape);
-  TensorInfo input_b_tensor_info(input_b_tensor_layout, input_b_shape, input_b_slice_shape);
-  TensorInfo output_tensor_info(output_tensor_layout, output_shape, output_slice_shape);
 
-  inputs_tensor_info_.push_back(input_a_tensor_info);
-  inputs_tensor_info_.push_back(input_b_tensor_info);
-  outputs_tensor_info_.push_back(output_tensor_info);
+  TensorInfo input_0_tensor_info(input_0_tensor_layout);
 
+  // input_1 do not need tensor info
+  inputs_tensor_info_.push_back(input_0_tensor_info);   // input_0
+  outputs_tensor_info_.push_back(input_0_tensor_info);  // output
   return SUCCESS;
 }
 
@@ -122,20 +143,29 @@ Status DropoutDoMaskInfo::SetCostUnderStrategy(const StrategyPtr& strategy) {
 }
 
 Status DropoutDoMaskInfo::GenerateStrategies(int32_t stage_id) {
-  CheckGlobalDeviceManager();
+  if (inputs_shape_.empty()) {
+    MS_LOG(ERROR) << name_ << ": The inputs shape is empty";
+    return FAILED;
+  }
+
   is_auto_parallel_ = true;
-  size_t dev_num = g_device_manager->GetDeviceListByStageId(stage_id).size();
-  Dimensions strategy(inputs_shape_[0].size() - 1, 1);
-  (void)strategy.insert(strategy.begin(), SizeToInt(dev_num));
-  std::vector<Dimensions> stra = {strategy};
-  StrategyPtr sp = std::make_shared<Strategy>(stage_id, stra);
-  if (SetCostUnderStrategy(sp) == SUCCESS) {
-    MS_LOG(INFO) << name_ << " : Successfully generated batch-parallel-strategy.";
-    PrintStrategy(sp);
-  } else {
-    MS_LOG(ERROR) << name_ << " : Generating batch-parallel-strategy failed.";
+  Shape input0_split(inputs_shape_[0].size(), 1);
+  Shapes splittable_inputs = {input0_split};
+  Shapes used_inputs_shape = {inputs_shape_[0]};
+
+  std::vector<StrategyPtr> sp_vector;
+  if (GenerateStrategiesForIndependentInputs(stage_id, used_inputs_shape, splittable_inputs, &sp_vector) != SUCCESS) {
+    MS_LOG(ERROR) << name_ << ": Generate strategies failed";
     return FAILED;
   }
+  size_t success = 0;
+  for (auto& sp : sp_vector) {
+    if (SetCostUnderStrategy(sp) == SUCCESS) {
+      success++;
+      MS_LOG(INFO) << name_ << ": Successfully generated " << success << " strategy";
+      PrintStrategy(sp);
+    }
+  }
   return SUCCESS;
 }
 
@@ -150,26 +180,105 @@ std::shared_ptr<std::vector<std::vector<int32_t>>> DropoutDoMaskInfo::GenerateBa
 
 Status DropoutDoMaskInfo::Init(const StrategyPtr& strategy) {
   if (InitWithAutoRepeatCalc(strategy) != SUCCESS) {
-    MS_LOG(ERROR) << name_ << " : Init failed.";
+    MS_LOG(ERROR) << name_ << ": Init failed.";
     return FAILED;
   }
 
-  MS_LOG(INFO) << name_ << " : Init success.";
+  MS_LOG(INFO) << name_ << ": Init success.";
   return SUCCESS;
 }
 
 Status DropoutDoMaskInfo::InitForCostModel(const StrategyPtr& strategy) {
   if (InitForCostModelWithAutoRepeatCalc(strategy) != SUCCESS) {
     if (is_auto_parallel_) {
-      MS_LOG(DEBUG) << name_ << " : Init for cost model failed.";
+      MS_LOG(DEBUG) << name_ << ": Init for cost model failed.";
     } else {
-      MS_LOG(ERROR) << name_ << " : Init for cost model failed.";
+      MS_LOG(ERROR) << name_ << ": Init for cost model failed.";
     }
     return FAILED;
   }
 
-  MS_LOG(INFO) << name_ << " : Init for cost model success.";
+  MS_LOG(INFO) << name_ << ": Init for cost model success.";
   return SUCCESS;
 }
+
+PrimitivePtr GetDropoutGenMaskPrim(const CNodePtr& cnode) {
+  MS_EXCEPTION_IF_NULL(cnode);
+  if (cnode->inputs().size() != DROPOUT_DO_MASK_CNODE_INPUT_SIZE) {
+    MS_LOG(EXCEPTION) << "The size of dropout do mask cnode's inputs must be " << DROPOUT_DO_MASK_CNODE_INPUT_SIZE;
+  }
+
+  AnfNodePtr dropout_gen_mask = cnode->input(DROPOUT_GEN_MASK_INDEX);
+  MS_EXCEPTION_IF_NULL(dropout_gen_mask);
+  if (!dropout_gen_mask->isa<CNode>()) {
+    MS_LOG(EXCEPTION) << "The dropout do mask cnode's input[" << DROPOUT_GEN_MASK_INDEX << "] must be a cnode";
+  }
+
+  auto dropout_gen_mask_cnode = dropout_gen_mask->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(dropout_gen_mask_cnode);
+  if (dropout_gen_mask_cnode->inputs().size() != DROPOUT_GEN_MASK_CNODE_INPUT_SIZE) {
+    MS_LOG(EXCEPTION) << "The size of dropout gen mask cnode's inputs must be " << DROPOUT_GEN_MASK_CNODE_INPUT_SIZE;
+  }
+  if (!IsValueNode<Primitive>(dropout_gen_mask_cnode->input(0))) {
+    MS_LOG(EXCEPTION) << "The input[0] of dropout gen mask cnode is not primitive";
+  }
+
+  ValueNodePtr value_node = dropout_gen_mask_cnode->input(0)->cast<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node);
+  PrimitivePtr prim = value_node->value()->cast<PrimitivePtr>();
+  MS_EXCEPTION_IF_NULL(prim);
+  if (prim->name() != DROPOUT_GEN_MASK) {
+    MS_LOG(EXCEPTION) << "The primitive name is not DropoutGenMask";
+  }
+  return prim;
+}
+
+// DropoutDoMask needs to be used together with DropoutGenMask. Only the first input tensor of DropoutGenMask is
+// split. Find the DropoutGenMask node in the anf graph according to DropoutDoMask node, and modify the input shape
+// of DropoutGenMask according to the strategy of DropoutDoMask. When the DropoutDoMask performs repeated calculation
+// and both seeds of DropoutGenMask are 0, two new seeds are automatically generated for DropoutGenMask.
+Operator DropoutDoMaskInfo::GetDropoutGenMaskReplaceOp(const CNodePtr& cnode) {
+  MS_EXCEPTION_IF_NULL(cnode);
+  PrimitivePtr prim = GetDropoutGenMaskPrim(cnode);
+  MS_EXCEPTION_IF_NULL(prim);
+
+  if (inputs_tensor_info_.empty()) {
+    MS_LOG(EXCEPTION) << "The tensor info of dropout do mask is empty";
+  }
+
+  if (cnode->inputs().size() != DROPOUT_DO_MASK_CNODE_INPUT_SIZE) {
+    MS_LOG(EXCEPTION) << "The size of dropout do mask cnode's inputs must be " << DROPOUT_DO_MASK_CNODE_INPUT_SIZE;
+  }
+
+  if (!cnode->input(DROPOUT_DO_MASK_KEEP_PROB_INDEX)->isa<ValueNode>()) {
+    MS_LOG(EXCEPTION) << "The keep prob of dropout do mask is not value node";
+  }
+
+  ValuePtr keep_prob = GetValueNode(cnode->input(DROPOUT_DO_MASK_KEEP_PROB_INDEX));
+  MS_EXCEPTION_IF_NULL(keep_prob);
+  auto attr = prim->attrs();
+  if ((attr.find(SEED0) == attr.end()) || (attr.find(SEED1) == attr.end())) {
+    MS_LOG(EXCEPTION) << "The attrs of dropout gen mask must be have seed0 and seed1";
+  }
+  int32_t seed_0 = GetValue<int32_t>(attr[SEED0]);
+  int32_t seed_1 = GetValue<int32_t>(attr[SEED1]);
+  if ((seed_0 == 0) && (seed_1 == 0) && (repeated_calc_num_ > 1)) {
+    seed_0 = SEED_NUM;
+    seed_1 = SEED_NUM;
+    SEED_NUM++;
+  }
+
+  Shape input_slice_shape = inputs_tensor_info_[0].slice_shape();
+  ValuePtr new_shape = MakeValue(input_slice_shape);
+  Attr attr_0 = std::make_pair(SEED0, MakeValue(seed_0));
+  Attr attr_1 = std::make_pair(SEED1, MakeValue(seed_1));
+  OperatorAttrs attrs = {attr_0, attr_1};
+  Attr param_0 = std::make_pair(SHAPE, new_shape);
+  Attr param_1 = std::make_pair(KEEP_PROB, keep_prob);
+  OperatorParams params = {std::make_pair(param_0, 1), std::make_pair(param_1, 2)};
+  OperatorArgs args = std::make_pair(attrs, params);
+  Operator replace_op = {std::make_pair(DROPOUT_GEN_MASK, args)};
+  return replace_op;
+}
 }  // namespace parallel
 }  // namespace mindspore

mindspore/ccsrc/parallel/ops_info/dropout_do_mask_info.h

@@ -33,7 +33,7 @@ class DropoutDoMaskInfo : public OperatorInfo {
  public:
   DropoutDoMaskInfo(const std::string& name, const Shapes& inputs_shape, const Shapes& outputs_shape,
                     const PrimitiveAttrs& attrs)
-      : OperatorInfo(name, inputs_shape, outputs_shape, attrs, std::make_shared<BatchParallelCost>()) {}
+      : OperatorInfo(name, inputs_shape, outputs_shape, attrs, std::make_shared<DropOutCost>()) {}
   ~DropoutDoMaskInfo() override = default;
 
   Status Init(const StrategyPtr& strategy) override;
@@ -41,6 +41,7 @@ class DropoutDoMaskInfo : public OperatorInfo {
   Status SetCostUnderStrategy(const StrategyPtr& strategy) override;
   Status InitForCostModel(const StrategyPtr& strategy) override;
   std::shared_ptr<std::vector<std::vector<int32_t>>> GenerateBatchStrategies() override;
+  Operator GetDropoutGenMaskReplaceOp(const CNodePtr& cnode);
 
  protected:
   Status CheckStrategy(const StrategyPtr& strategy) override;
@@ -51,6 +52,8 @@ class DropoutDoMaskInfo : public OperatorInfo {
   Status InferTensorInfo() override;
   Status InferDevMatrixShape() override;
 };
+
+using DropoutDoMaskInfoPtr = std::shared_ptr<DropoutDoMaskInfo>;
 }  // namespace parallel
 }  // namespace mindspore
 

mindspore/ccsrc/parallel/ops_info/generator_info.cc

-/**
- * Copyright 2019 Huawei Technologies Co., Ltd
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "parallel/ops_info/generator_info.h"
-
-#include <algorithm>
-#include <memory>
-#include <utility>
-#include <vector>
-
-#include "ir/value.h"
-#include "parallel/device_matrix.h"
-#include "parallel/strategy.h"
-#include "parallel/tensor_layout/tensor_redistribution.h"
-
-namespace mindspore {
-namespace parallel {
-Status GeneratorBase::InferTensorMap() {
-  TensorMap output_tensor_map = {MAP_NONE};
-  outputs_tensor_map_.push_back(output_tensor_map);
-  return SUCCESS;
-}
-
-Status GeneratorBase::InferTensorInfo() {
-  Shape output_shape = outputs_shape_.at(0);
-  Shape output_slice_shape = outputs_shape_.at(0);
-
-  TensorLayout output_tensor_layout;
-  if (output_tensor_layout.InitFromVector(dev_matrix_shape_, outputs_tensor_map_[0], output_shape) != SUCCESS) {
-    MS_LOG(ERROR) << name_ << " : Creat output tensor layout failed.";
-    return FAILED;
-  }
-  TensorInfo output_tensor_info(output_tensor_layout, output_shape, output_slice_shape);
-  outputs_tensor_info_.push_back(output_tensor_info);
-
-  return SUCCESS;
-}
-
-Status GeneratorBase::InferDevMatrixShape() {
-  std::vector<Dimensions> stra = strategy_->GetInputDim();
-  Dimensions input_strategy = stra.at(0);
-
-  dev_matrix_shape_ = input_strategy;
-
-  return SUCCESS;
-}
-
-Status GeneratorBase::SetCostUnderStrategy(const StrategyPtr &strategy) {
-  if (SetCostUnderStrategyBase(strategy) != SUCCESS) {
-    if (is_auto_parallel_) {
-      MS_LOG(DEBUG) << name_ << " : Set cost under strategy failed.";
-    } else {
-      MS_LOG(ERROR) << name_ << " : Set cost under strategy failed.";
-    }
-    return FAILED;
-  }
-
-  return SUCCESS;
-}
-
-Status DropoutGenMaskInfo::GenerateStrategies(int32_t stage_id) {
-  if (input_value_.empty()) {
-    MS_LOG(ERROR) << name_ << " : Input value is empty.";
-    return FAILED;
-  }
-  Shape param = GetValue<const std::vector<int>>(input_value_[0]);
-  if (param.empty()) {
-    MS_LOG(ERROR) << name_ << " : Input value [0] is empty.";
-    return FAILED;
-  }
-  // Now,only support batch parallel.
-  CheckGlobalDeviceManager();
-  is_auto_parallel_ = true;
-  size_t dev_num = g_device_manager->GetDeviceListByStageId(stage_id).size();
-  Dimensions strategy(param.size() - 1, 1);
-  (void)strategy.insert(strategy.begin(), SizeToInt(dev_num));
-  std::vector<Dimensions> stra = {strategy};
-  StrategyPtr sp = std::make_shared<Strategy>(stage_id, stra);
-  if (SetCostUnderStrategy(sp) == SUCCESS) {
-    MS_LOG(INFO) << name_ << " : Successfully generated batch-parallel-strategy.";
-    PrintStrategy(sp);
-  } else {
-    MS_LOG(ERROR) << name_ << " : Generating batch-parallel-strategy failed.";
-    return FAILED;
-  }
-  return SUCCESS;
-}
-
-Status DropoutGenMaskInfo::CheckStrategy(const StrategyPtr &strategy) {
-  if (strategy->GetInputNumber() != 1) {
-    if (is_auto_parallel_) {
-      MS_LOG(DEBUG) << name_ << " : The strategy is wrong.";
-    } else {
-      MS_LOG(ERROR) << name_ << " : The strategy is wrong.";
-    }
-    return FAILED;
-  }
-
-  return SUCCESS;
-}
-
-Status DropoutGenMaskInfo::InferReplaceOps(const StrategyPtr &strategy) {
-  Shape shape = GetValue<const std::vector<int>>(input_value_[0]);
-  Strategys stra = strategy->GetInputDim();
-  Dimensions input_strategy = stra.at(0);
-  int32_t dev_num = *(input_strategy.begin());
-  if (dev_num <= 0) {
-    MS_LOG(ERROR) << name_ << " : The number of devices should not be less than 0.";
-    return FAILED;
-  }
-  // Batch parallel
-  if (shape[0] % dev_num != 0) {
-    MS_LOG(ERROR) << name_ << " : The shape " << shape[0] << " can't be exact divided by device number " << dev_num;
-    return FAILED;
-  }
-  shape[0] = shape[0] / dev_num;
-  ValuePtr shape_ptr = MakeValue(shape);
-  Attr attr_0 = std::make_pair(SEED0, attrs_[SEED0]);
-  Attr attr_1 = std::make_pair(SEED1, attrs_[SEED1]);
-  OperatorAttrs attrs = {attr_0, attr_1};
-  Attr param_0 = std::make_pair(SHAPE, shape_ptr);
-  Attr param_1 = std::make_pair(KEEP_PROB, input_value_[1]);
-  OperatorParams params = {std::make_pair(param_0, 1), std::make_pair(param_1, 2)};
-  OperatorArgs args = std::make_pair(attrs, params);
-  replace_op_ = {std::make_pair(DROPOUT_GEN_MASK, args)};
-  return SUCCESS;
-}
-
-std::shared_ptr<std::vector<std::vector<int32_t>>> DropoutGenMaskInfo::GenerateBatchStrategies() {
-  if (input_value_.empty()) {
-    MS_LOG(EXCEPTION) << name_ << " : Input value is empty.";
-  }
-  Shape param = GetValue<const std::vector<int>>(input_value_[0]);
-  if (param.empty()) {
-    MS_LOG(EXCEPTION) << name_ << " : Input value [0] is empty.";
-  }
-  // Now,only support batch parallel.
-  CheckGlobalDeviceManager();
-  size_t dev_num = g_device_manager->GetDeviceListByStageId(0).size();
-  Dimensions strategy(param.size() - 1, 1);
-  (void)strategy.insert(strategy.begin(), SizeToInt(dev_num));
-  std::vector<Dimensions> strategy_v = {strategy};
-  return std::make_shared<std::vector<std::vector<int32_t>>>(strategy_v);
-}
-
-Status GeneratorBase::Init(const StrategyPtr &strategy) {
-  if (InitWithAutoRepeatCalc(strategy) != SUCCESS) {
-    MS_LOG(ERROR) << name_ << " : Init failed.";
-    return FAILED;
-  }
-
-  if (InferReplaceOps(strategy) != SUCCESS) {
-    MS_LOG(ERROR) << name_ << " : Infer replace ops failed.";
-    return FAILED;
-  }
-
-  MS_LOG(INFO) << name_ << " : Init success.";
-  return SUCCESS;
-}
-
-Status GeneratorBase::InitForCostModel(const StrategyPtr &strategy) {
-  if (InitForCostModelWithAutoRepeatCalc(strategy) != SUCCESS) {
-    if (is_auto_parallel_) {
-      MS_LOG(DEBUG) << name_ << " : Init for cost model failed.";
-    } else {
-      MS_LOG(ERROR) << name_ << " : Init for cost model failed.";
-    }
-    return FAILED;
-  }
-
-  MS_LOG(INFO) << name_ << " : Init for cost model success.";
-  return SUCCESS;
-}
-}  // namespace parallel
-}  // namespace mindspore

mindspore/ccsrc/parallel/ops_info/generator_info.h

-/**
- * Copyright 2019 Huawei Technologies Co., Ltd
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#ifndef MINDSPORE_CCSRC_PARALLEL_OPS_INFO_GENERATOR_INFO_H_
-#define MINDSPORE_CCSRC_PARALLEL_OPS_INFO_GENERATOR_INFO_H_
-
-#include <memory>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "parallel/auto_parallel/operator_costmodel.h"
-#include "parallel/ops_info/operator_info.h"
-#include "parallel/strategy.h"
-
-namespace mindspore {
-namespace parallel {
-class GeneratorBase : public OperatorInfo {
- public:
-  GeneratorBase(const std::string &operator_name, const Shapes &inputs_shape, const Shapes &outputs_shape,
-                const PrimitiveAttrs &attrs)
-      : OperatorInfo(operator_name, inputs_shape, outputs_shape, attrs, std::make_shared<GeneratorBaseCost>()) {}
-
-  ~GeneratorBase() override = default;
-
-  Status Init(const StrategyPtr &strategy) override;
-  Status SetCostUnderStrategy(const StrategyPtr &strategy) override;
-  Status InitForCostModel(const StrategyPtr &strategy) override;
-
- protected:
-  // For now, generator ops don't have attributes
-  Status GetAttrs() override { return Status::SUCCESS; }
-  Status InferTensorMap() override;
-  Status InferTensorInfo() override;
-  Status InferDevMatrixShape() override;
-  Status InferMirrorOps() override { return SUCCESS; }
-  Status InferForwardCommunication() override { return SUCCESS; }
-  virtual Status InferReplaceOps(const StrategyPtr &strategy) = 0;
-};
-
-class DropoutGenMaskInfo : public GeneratorBase {
- public:
-  DropoutGenMaskInfo(const std::string &name, const Shapes &inputs_shape, const Shapes &outputs_shape,
-                     const PrimitiveAttrs &attrs)
-      : GeneratorBase(name, inputs_shape, outputs_shape, attrs) {}
-  ~DropoutGenMaskInfo() override = default;
-  Status GenerateStrategies(int32_t stage_id) override;
-  std::shared_ptr<std::vector<std::vector<int32_t>>> GenerateBatchStrategies() override;
-
- protected:
-  Status CheckStrategy(const StrategyPtr &strategy) override;
-  Status InferReplaceOps(const StrategyPtr &strategy) override;
-};
-}  // namespace parallel
-}  // namespace mindspore
-
-#endif  // MINDSPORE_CCSRC_PARALLEL_OPS_INFO_GENERATOR_INFO_H_

mindspore/ccsrc/parallel/ops_info/ops_info_head_files.h

@@ -24,7 +24,6 @@
 #include "parallel/ops_info/comparison_function_info.h"
 #include "parallel/ops_info/dropout_do_mask_info.h"
 #include "parallel/ops_info/elementary_function_info.h"
-#include "parallel/ops_info/generator_info.h"
 #include "parallel/ops_info/get_next_info.h"
 #include "parallel/ops_info/l2_normalize_info.h"
 #include "parallel/ops_info/loss_info.h"

mindspore/ccsrc/parallel/ops_info/ops_utils.h

@@ -34,6 +34,10 @@ constexpr size_t SOFTMAX_ATTR_SIZE = 1;
 constexpr size_t ACTIVATION_INPUTS_SIZE = 1;
 constexpr size_t ACTIVATION_OUTPUTS_SIZE = 1;
 constexpr size_t EXPANDDIMS_INPUT_SIZE = 2;
+constexpr size_t DROPOUT_DO_MASK_CNODE_INPUT_SIZE = 4;
+constexpr size_t DROPOUT_GEN_MASK_CNODE_INPUT_SIZE = 3;
+constexpr size_t DROPOUT_GEN_MASK_INDEX = 2;
+constexpr size_t DROPOUT_DO_MASK_KEEP_PROB_INDEX = 3;
 constexpr size_t SoftmaxCrossEntropyWithLogitsAttrSize = 1;
 constexpr size_t SoftmaxCrossEntropyWithLogitsInputsSize = 2;
 constexpr size_t SoftmaxCrossEntropyWithLogitsOutputsSize = 2;

mindspore/ccsrc/parallel/step_auto_parallel.cc

@@ -69,7 +69,6 @@ std::vector<std::string> splittable_op_ = {MATMUL,
                                            RELU,
                                            ONEHOT,
                                            DROPOUT_DO_MASK,
-                                           DROPOUT_GEN_MASK,
                                            REDUCE_MAX,
                                            REDUCE_MIN,
                                            ARGMAXWITHVALUE,

mindspore/ccsrc/parallel/step_parallel.cc

@@ -484,8 +484,6 @@ void StepSplitTensor(const AnfNodePtr& node, const FuncGraphManagerPtr& manager)
     }
     if (IsParallelCareNode(use_cnode)) {
       SplitTensor(node, use_cnode, node_pair.second);
-    } else {
-      StepSplitTensor(use_cnode, manager);
     }
   }
 }
@@ -525,6 +523,26 @@ std::vector<AnfNodePtr> ReplaceOpInput(const Operator& replace_op, const std::st
   return replace_input;
 }
 
+void ReplaceOneOp(const Operator& replace_op, const CNodePtr& node) {
+  FuncGraphPtr func_graph = node->func_graph();
+  MS_EXCEPTION_IF_NULL(func_graph);
+  FuncGraphManagerPtr manager = func_graph->manager();
+  if (manager == nullptr) {
+    MS_LOG(EXCEPTION) << "Failure:AddNode error since manager is nullptr";
+  }
+  std::string instance_name = CreateInstanceName(node, 0);
+  std::vector<AnfNodePtr> replace_input;
+  replace_input = ReplaceOpInput(replace_op, instance_name, node);
+  CNodePtr replace_node = func_graph->NewCNode(replace_input);
+  MS_EXCEPTION_IF_NULL(replace_node);
+  ScopePtr scope = node->scope();
+  MS_EXCEPTION_IF_NULL(scope);
+  replace_node->set_scope(scope);
+  replace_node->set_in_forward_flag(true);
+  replace_input[0]->set_scope(scope);
+  (void)manager->Replace(node, replace_node);
+}
+
 void StepReplaceOp(OperatorVector replace_op, const CNodePtr& node) {
   // step1:get graph manager distribute_operator
   OperatorInfoPtr distribute_operator = node->operator_info();
@@ -1757,6 +1775,28 @@ void StepReplace(const OperatorInfoPtr& distribute_operator, const CNodePtr& cno
   }
 }
 
+void HandleDropoutNode(const OperatorInfoPtr& distribute_operator, const CNodePtr& cnode) {
+  MS_EXCEPTION_IF_NULL(distribute_operator);
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  std::string op_name = distribute_operator->name();
+  if (op_name.find(DROPOUT_DO_MASK) == std::string::npos) {
+    return;
+  }
+
+  DropoutDoMaskInfoPtr dropout_do_mask = std::dynamic_pointer_cast<DropoutDoMaskInfo>(distribute_operator);
+  MS_EXCEPTION_IF_NULL(dropout_do_mask);
+  Operator replace_op = dropout_do_mask->GetDropoutGenMaskReplaceOp(cnode);
+  if (cnode->inputs().size() != DROPOUT_DO_MASK_CNODE_INPUT_SIZE) {
+    MS_LOG(EXCEPTION) << "The size of drop out do mask cnode's input is not " << DROPOUT_DO_MASK_CNODE_INPUT_SIZE;
+  }
+  ReplaceOneOp(replace_op, cnode->input(DROPOUT_GEN_MASK_INDEX)->cast<CNodePtr>());
+}
+
+void HandleSpecialNode(const OperatorInfoPtr& distribute_operator, const CNodePtr& cnode) {
+  HandleDropoutNode(distribute_operator, cnode);
+}
+
 void ParallelCommunication(const FuncGraphPtr& root, const std::vector<AnfNodePtr>& all_nodes,
                            const FuncGraphManagerPtr& manager) {
   MS_EXCEPTION_IF_NULL(root);
@@ -1804,6 +1844,8 @@ void ParallelCommunication(const FuncGraphPtr& root, const std::vector<AnfNodePt
 
       // StepReplace
       StepReplace(distribute_operator, cnode);
+
+      HandleSpecialNode(distribute_operator, cnode);
     } else if (IsValueNode<Tensor>(node)) {
       StepSplitTensor(node, manager);
     }

tests/ut/cpp/parallel/ops_info/dropout_do_mask_info_test.cc

-/**
- * Copyright 2019 Huawei Technologies Co., Ltd
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <string>
-#include <list>
-#include <vector>
-#include "common/common_test.h"
-#include "parallel/strategy.h"
-#include "parallel/ops_info/dropout_do_mask_info.h"
-#include "parallel/device_manager.h"
-#include "parallel/step_parallel.h"
-
-namespace mindspore {
-namespace parallel {
-
-class DropoutDoMaskInfo;
-using DropoutDoMaskInfoPtr = std::shared_ptr<DropoutDoMaskInfo>;
-DropoutDoMaskInfoPtr do_mask;
-
-class TestDropoutDoMaskInfo : public UT::Common {
- public:
-  TestDropoutDoMaskInfo() {}
-  void SetUp();
-  void TearDown() {}
-};
-
-void TestDropoutDoMaskInfo::SetUp() {
-  std::vector<int32_t> dev_list;
-
-  for (int32_t i = 0; i < 34; i++) {
-    dev_list.push_back(i);
-  }
-
-  std::vector<int32_t> stage_map;
-  stage_map.push_back(32);
-  stage_map.push_back(2);
-
-  int32_t local_dev = 0;
-
-  // create a new g_device_manager
-  g_device_manager = std::make_shared<DeviceManager>();
-  g_device_manager->Init(dev_list, local_dev, stage_map, "hccl");
-
-  std::unordered_map<std::string, ValuePtr> attr;
-
-  Shapes inputs_shape = {{32, 128}, {64}, {}};
-  Shapes outputs_shape = {{32, 128}};
-  do_mask = std::make_shared<DropoutDoMaskInfo>("do_mask_info", inputs_shape, outputs_shape, attr);
-}
-
-TEST_F(TestDropoutDoMaskInfo, InferDevMatrixShape) {
-  std::vector<Dimensions> stra = {{4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  do_mask->Init(strategy);
-  std::vector<int32_t> dev_matrix_shape = do_mask->dev_matrix_shape();
-
-  std::vector<int32_t> expect = {4, 8};
-  ASSERT_EQ(dev_matrix_shape, expect);
-}
-
-TEST_F(TestDropoutDoMaskInfo, InferSliceShape) {
-  std::vector<Dimensions> stra = {{4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  do_mask->Init(strategy);
-  std::vector<TensorInfo> inputs = do_mask->inputs_tensor_info();
-  std::vector<TensorInfo> outputs = do_mask->outputs_tensor_info();
-
-  Shape input_a_slice_shape_expect = {8, 16};
-  Shape input_b_slice_shape_expect = {64};
-  Shape output_slice_shape_expect = {8, 16};
-
-  TensorInfo input_a_tensor_info = inputs.at(0);
-  TensorInfo input_b_tensor_info = inputs.at(1);
-  TensorInfo output_tensor_info = outputs.at(0);
-  Shape input_a_slice_shape = input_a_tensor_info.slice_shape();
-  Shape input_b_slice_shape = input_b_tensor_info.slice_shape();
-  Shape output_slice_shape = output_tensor_info.slice_shape();
-
-  ASSERT_EQ(input_a_slice_shape, input_a_slice_shape_expect);
-  ASSERT_EQ(input_b_slice_shape, input_b_slice_shape_expect);
-  ASSERT_EQ(output_slice_shape, output_slice_shape_expect);
-}
-
-TEST_F(TestDropoutDoMaskInfo, GetTensorLayout) {
-  std::vector<Dimensions> stra = {{4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  do_mask->Init(strategy);
-  std::vector<TensorInfo> inputs = do_mask->inputs_tensor_info();
-  std::vector<TensorInfo> outputs = do_mask->outputs_tensor_info();
-
-  TensorMap input_a_map_expect = {1, 0};
-  TensorMap input_b_map_expect = {-1};
-  TensorMap output_map_expect = {1, 0};
-
-  TensorInfo input_a_tensor_info = inputs.at(0);
-  TensorInfo input_b_tensor_info = inputs.at(1);
-  TensorInfo output_tensor_info = outputs.at(0);
-  Map input_a_tensor_map = input_a_tensor_info.tensor_layout().origin_tensor_map();
-  Map input_b_tensor_map = input_b_tensor_info.tensor_layout().origin_tensor_map();
-  Map output_tensor_map = output_tensor_info.tensor_layout().origin_tensor_map();
-
-  ASSERT_EQ(input_a_tensor_map.array(), input_a_map_expect);
-  ASSERT_EQ(input_b_tensor_map.array(), input_b_map_expect);
-  ASSERT_EQ(output_tensor_map.array(), output_map_expect);
-}
-
-TEST_F(TestDropoutDoMaskInfo, GetForwardOp) {
-  std::vector<Dimensions> stra = {{4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  do_mask->Init(strategy);
-  OperatorVector forward_op = do_mask->forward_op();
-  size_t size = forward_op.size();
-
-  ASSERT_EQ(size, 0);
-}
-
-TEST_F(TestDropoutDoMaskInfo, CheckStrategy1) {
-  std::vector<Dimensions> stra = {{4, 8, 2}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = do_mask->Init(strategy);
-  ASSERT_EQ(ret, FAILED);
-}
-
-TEST_F(TestDropoutDoMaskInfo, CheckStrategy2) {
-  std::vector<Dimensions> stra = {{8, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = do_mask->Init(strategy);
-  ASSERT_EQ(ret, FAILED);
-}
-
-TEST_F(TestDropoutDoMaskInfo, CheckStrategy3) {
-  std::vector<Dimensions> stra = {{4, 8}, {4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = do_mask->Init(strategy);
-  ASSERT_EQ(ret, FAILED);
-}
-
-TEST_F(TestDropoutDoMaskInfo, CheckStrategy4) {
-  std::vector<Dimensions> stra = {{4, 8}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = do_mask->Init(strategy);
-  ASSERT_EQ(ret, SUCCESS);
-}
-}  // namespace parallel
-}  // namespace mindspore

tests/ut/cpp/parallel/ops_info/generator_info_test.cc

-/**
- * Copyright 2019 Huawei Technologies Co., Ltd
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <string>
-#include <list>
-#include <vector>
-#include "common/common_test.h"
-#include "parallel/strategy.h"
-#include "parallel/ops_info/generator_info.h"
-#include "parallel/device_manager.h"
-#include "parallel/step_parallel.h"
-
-namespace mindspore {
-namespace parallel {
-
-class DropoutGenMaskInfo;
-using DropoutGenMaskInfoPtr = std::shared_ptr<DropoutGenMaskInfo>;
-DropoutGenMaskInfoPtr gen_mask;
-
-class TestDropoutGenMaskInfo : public UT::Common {
- public:
-  TestDropoutGenMaskInfo() {}
-  void SetUp();
-  void TearDown() {}
-};
-
-void TestDropoutGenMaskInfo::SetUp() {
-  std::vector<int32_t> dev_list;
-
-  for (int32_t i = 0; i < 10; i++) {
-    dev_list.push_back(i);
-  }
-
-  std::vector<int32_t> stage_map;
-  stage_map.push_back(8);
-  stage_map.push_back(2);
-
-  int32_t local_dev = 0;
-
-  // create a new g_device_manager
-  g_device_manager = std::make_shared<DeviceManager>();
-  g_device_manager->Init(dev_list, local_dev, stage_map, "hccl");
-
-  std::unordered_map<std::string, ValuePtr> attr;
-
-  Shapes inputs_shape;
-  Shapes outputs_shape = {{128}};
-  std::vector<int> shape = {32, 128};
-  ValuePtr val0 = MakeValue(shape);
-  ValuePtr val1;
-  std::vector<ValuePtr> val = {val0, val1};
-  gen_mask = std::make_shared<DropoutGenMaskInfo>("gen_mask_info", inputs_shape, outputs_shape, attr);
-  gen_mask->set_input_value(val);
-}
-
-TEST_F(TestDropoutGenMaskInfo, InferDevMatrixShape) {
-  std::vector<Dimensions> stra = {{8, 1}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  gen_mask->Init(strategy);
-  std::vector<int32_t> dev_matrix_shape = gen_mask->dev_matrix_shape();
-
-  std::vector<int32_t> expect = {8, 1};
-  ASSERT_EQ(dev_matrix_shape, expect);
-}
-
-TEST_F(TestDropoutGenMaskInfo, InferSliceShape) {
-  std::vector<Dimensions> stra = {{8, 1}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  gen_mask->Init(strategy);
-  std::vector<TensorInfo> outputs = gen_mask->outputs_tensor_info();
-
-  Shape output_slice_shape_expect = {128};
-
-  TensorInfo output_tensor_info = outputs.at(0);
-  Shape output_slice_shape = output_tensor_info.slice_shape();
-
-  ASSERT_EQ(output_slice_shape, output_slice_shape_expect);
-}
-
-TEST_F(TestDropoutGenMaskInfo, GetTensorLayout) {
-  std::vector<Dimensions> stra = {{8, 1}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  gen_mask->Init(strategy);
-  std::vector<TensorInfo> outputs = gen_mask->outputs_tensor_info();
-
-  TensorMap output_map_expect = {-1};
-
-  TensorInfo output_tensor_info = outputs.at(0);
-  Map output_tensor_map = output_tensor_info.tensor_layout().origin_tensor_map();
-
-  ASSERT_EQ(output_tensor_map.array(), output_map_expect);
-}
-
-TEST_F(TestDropoutGenMaskInfo, GetForwardOp) {
-  std::vector<Dimensions> stra = {{8, 1}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  gen_mask->Init(strategy);
-  OperatorVector forward_op = gen_mask->forward_op();
-  size_t size = forward_op.size();
-
-  ASSERT_EQ(size, 0);
-}
-
-TEST_F(TestDropoutGenMaskInfo, CheckStrategy1) {
-  std::vector<Dimensions> stra = {{4, 8, 2}, {2, 3}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = gen_mask->Init(strategy);
-  ASSERT_EQ(ret, FAILED);
-}
-
-TEST_F(TestDropoutGenMaskInfo, CheckStrategy2) {
-  std::vector<Dimensions> stra = {{8, 1}};
-  StrategyPtr strategy = NewStrategy(0, stra);
-
-  Status ret = gen_mask->Init(strategy);
-  ASSERT_EQ(ret, SUCCESS);
-}
-}  // namespace parallel
-}  // namespace mindspore

tests/ut/python/parallel/test_dropout_do_mask.py

+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import mindspore as ms
+from mindspore import context, Tensor, Parameter
+from mindspore.nn import Cell, TrainOneStepCell, Momentum
+from mindspore.ops import operations as P
+from mindspore.common.api import _executor
+
+
+class Net(Cell):
+    def __init__(self, mul_weight, strategy1=None, strategy2=None):
+        super().__init__()
+        self.mul = P.Mul().set_strategy(strategy1)
+        self.mul2 = P.Mul().set_strategy(strategy1)
+        self.dropout_do_mask = P.DropoutDoMask().set_strategy(strategy2)
+        self.dropout_gen_mask = P.DropoutGenMask()
+        self.get_shape = P.Shape()
+        self.cast = P.Cast()
+        self.mul_weight = Parameter(mul_weight, "w1")
+        self.mul_weight2 = Parameter(mul_weight, "w2")
+        self.keep_prob = Tensor(0.9)
+
+    def construct(self, x, b):
+        out = self.mul(x, self.mul_weight)
+        shape = self.get_shape(out)
+        dtype = P.DType()(out)
+        keep_prob = self.cast(self.keep_prob, dtype)
+        mask = self.dropout_gen_mask(shape, keep_prob)
+        out = self.dropout_do_mask(out, mask, keep_prob)
+        out = self.mul2(out, self.mul_weight2)
+        return out
+
+
+_x = Tensor(np.ones([128, 64]), dtype=ms.float32)
+_w1 = Tensor(np.ones([128, 64]), dtype=ms.float32)
+_b = Tensor(np.ones([128, 64]), dtype=ms.float32)
+
+
+def compile(net):
+    optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
+    train_net = TrainOneStepCell(net, optimizer)
+    _executor.compile(train_net, _x, _b)
+    context.reset_auto_parallel_context()
+
+
+def test_dropout_do_mask_data_parallel():
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
+    strategy1 = ((16, 1), (16, 1))
+    strategy2 = ((16, 1),)
+    net = Net(_w1, strategy1, strategy2)
+    compile(net)
+
+
+def test_dropout_do_mask_model_parallel():
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
+    strategy1 = ((1, 16), (1, 16))
+    strategy2 = ((1, 16),)
+    net = Net(_w1, strategy1, strategy2)
+    compile(net)
+
+
+def test_dropout_do_mask_hybrid_parallel():
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
+    strategy1 = ((4, 4), (4, 4))
+    strategy2 = ((4, 4),)
+    net = Net(_w1, strategy1, strategy2)
+    compile(net)
+
+
+def test_dropout_do_mask_auto_parallel():
+    context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=16, global_rank=0)
+    net = Net(_w1)
+    compile(net)
+
+
+def test_dropout_do_mask_repeat_calc():
+    context.set_auto_parallel_context(parallel_mode="semi_auto_parallel", device_num=16, global_rank=0)
+    strategy1 = ((4, 4), (4, 4))
+    strategy2 = ((2, 4),)
+    net = Net(_w1, strategy1, strategy2)
+    compile(net)