forked from huawei/mindspore2022
232 lines
8.3 KiB
C++
232 lines
8.3 KiB
C++
/**
|
|
* 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/loss_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 SoftmaxCrossEntropyWithLogitsInfo::CheckStrategy(const mindspore::parallel::StrategyPtr& strategy) {
|
|
if (CheckStrategyValue(strategy, inputs_shape_, is_auto_parallel_) != SUCCESS) {
|
|
if (is_auto_parallel_) {
|
|
MS_LOG(DEBUG) << name_ << " : Invalid strategy.";
|
|
} else {
|
|
MS_LOG(ERROR) << name_ << " : Invalid strategy.";
|
|
}
|
|
return FAILED;
|
|
}
|
|
|
|
std::vector<Dimensions> stra = strategy->GetInputDim();
|
|
Dimensions input_strategy = stra.at(0);
|
|
Dimensions label_strategy = stra.at(1);
|
|
if (input_strategy != label_strategy) {
|
|
MS_LOG(ERROR) << name_ << " : Strategies of relevant dimensions are not equal.";
|
|
return FAILED;
|
|
}
|
|
|
|
int32_t axis_index = axis_;
|
|
if (axis_ < 0) {
|
|
size_t input_dim = inputs_shape_.at(0).size();
|
|
axis_index = static_cast<int32_t>(input_dim) + axis_;
|
|
}
|
|
|
|
int32_t input_axis_strategy = input_strategy.at(IntToSize(axis_index));
|
|
int32_t label_axis_strategy = label_strategy.at(IntToSize(axis_index));
|
|
// Dimension corresponding to axis is un-splittable
|
|
if ((input_axis_strategy != MIN_SLICE_NUM) && (label_axis_strategy != MIN_SLICE_NUM)) {
|
|
if (is_auto_parallel_) {
|
|
MS_LOG(DEBUG) << name_
|
|
<< " : The strategy corresponding to axis dimension is not 1, input: " << input_axis_strategy
|
|
<< ", label: " << label_axis_strategy;
|
|
} else {
|
|
MS_LOG(ERROR) << name_
|
|
<< " : The strategy corresponding to axis dimension is not 1, input: " << input_axis_strategy
|
|
<< ", label: " << label_axis_strategy;
|
|
}
|
|
return FAILED;
|
|
}
|
|
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::GetAttrs() {
|
|
if ((inputs_shape_.size() != SoftmaxCrossEntropyWithLogitsInputsSize) ||
|
|
(outputs_shape_.size() != SoftmaxCrossEntropyWithLogitsOutputsSize)) {
|
|
MS_LOG(ERROR) << name_ << " : Inputs shape size or outputs shape size is wrong.";
|
|
return FAILED;
|
|
}
|
|
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::InferDevMatrixShape() {
|
|
std::vector<Dimensions> stra = strategy_->GetInputDim();
|
|
Dimensions input_strategy = stra.at(0);
|
|
dev_matrix_shape_ = input_strategy;
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::InferTensorMap() {
|
|
std::vector<int32_t> tensor_map_index;
|
|
size_t size = inputs_shape_[0].size();
|
|
// such as 4: tensor_map_index [3,2,1,0]
|
|
for (size_t i = 0; i < size; ++i) {
|
|
tensor_map_index.push_back((int32_t)(size - i - 1));
|
|
}
|
|
|
|
std::vector<int32_t> first_output_tensor_map = {tensor_map_index[0]};
|
|
inputs_tensor_map_.push_back(tensor_map_index); // input
|
|
inputs_tensor_map_.push_back(tensor_map_index); // label
|
|
outputs_tensor_map_.push_back(first_output_tensor_map); // output-0
|
|
outputs_tensor_map_.push_back(tensor_map_index); // output-1
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::InferTensorInfo() {
|
|
// infer tensor shape
|
|
Shape input_shape = inputs_shape_.at(0);
|
|
Shape first_output_shape = outputs_shape_.at(0);
|
|
|
|
// infer slice shape
|
|
Shapes inputs_slice_shape, outputs_slice_shape;
|
|
Strategys inputs_strategy = strategy_->GetInputDim();
|
|
Strategys outputs_strategy = {{inputs_strategy[0][0]}, inputs_strategy.at(0)};
|
|
if (InferSliceShape(inputs_strategy, outputs_strategy, &inputs_slice_shape, &outputs_slice_shape) != SUCCESS) {
|
|
return FAILED;
|
|
}
|
|
Shape input_slice_shape = inputs_slice_shape.at(0);
|
|
Shape first_output_slice_shape = outputs_slice_shape.at(0);
|
|
|
|
TensorMap input_tensor_map = inputs_tensor_map_.at(0);
|
|
TensorMap first_output_tensor_map = outputs_tensor_map_.at(0);
|
|
|
|
TensorLayout input_tensor_layout, first_output_tensor_layout;
|
|
if ((input_tensor_layout.InitFromVector(dev_matrix_shape_, input_tensor_map, input_shape) != SUCCESS) ||
|
|
(first_output_tensor_layout.InitFromVector(dev_matrix_shape_, first_output_tensor_map, first_output_shape) !=
|
|
SUCCESS)) {
|
|
return FAILED;
|
|
}
|
|
TensorInfo input_tensor_info(input_tensor_layout, input_shape, input_slice_shape);
|
|
TensorInfo first_output_tensor_info(first_output_tensor_layout, first_output_shape, first_output_slice_shape);
|
|
|
|
inputs_tensor_info_.push_back(input_tensor_info); // input
|
|
inputs_tensor_info_.push_back(input_tensor_info); // label
|
|
outputs_tensor_info_.push_back(first_output_tensor_info); // output-0
|
|
outputs_tensor_info_.push_back(input_tensor_info); // output-1
|
|
|
|
return SUCCESS;
|
|
}
|
|
|
|
// There are two outputs for SoftmaxCrossEntropyWithLogits, and outputs[1] is used for grad and overload the function.
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::InferAsLossDivisor() {
|
|
if (outputs_tensor_map_.size() != 2) {
|
|
MS_LOG(ERROR) << name_ << " : The size of outputs tensor map " << outputs_tensor_map_.size() << " is error.";
|
|
return FAILED;
|
|
}
|
|
as_loss_divisor_ = ComputeRepeatDeviceNumByTensorMap(dev_matrix_shape_, outputs_tensor_map_[1]);
|
|
MS_LOG(INFO) << name_ << " : The dev matrix shape is " << ShapeToString(dev_matrix_shape_)
|
|
<< ", the output tensor map is " << ShapeToString(outputs_tensor_map_[1]) << ", as_loss_divisor_ is "
|
|
<< as_loss_divisor_;
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::Init(const StrategyPtr& strategy) {
|
|
if (InitWithAutoRepeatCalc(strategy) != SUCCESS) {
|
|
MS_LOG(ERROR) << name_ << " : Init failed.";
|
|
return FAILED;
|
|
}
|
|
|
|
MS_LOG(INFO) << name_ << " : Init success.";
|
|
return SUCCESS;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::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;
|
|
}
|
|
|
|
void SoftmaxCrossEntropyWithLogitsInfo::ReComputeBatchSplitFlagList() {
|
|
for (size_t i = 0; i < inputs_shape_.size(); ++i) {
|
|
split_flag_list_[i] = true;
|
|
}
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::GenerateStrategies(int32_t stage_id) {
|
|
if (GetAttrs() != SUCCESS) {
|
|
MS_LOG(ERROR) << name_ << " : GetAttrs failed.";
|
|
return FAILED;
|
|
}
|
|
int32_t axis_index = axis_;
|
|
if (axis_ < 0) {
|
|
size_t input_dim = inputs_shape_[0].size();
|
|
axis_index = static_cast<int32_t>(input_dim) + axis_;
|
|
}
|
|
is_auto_parallel_ = true;
|
|
|
|
Shape input0_split(inputs_shape_[0].size(), 1);
|
|
input0_split[IntToSize(axis_index)] = 0;
|
|
Shapes splittable_inputs = {input0_split, input0_split};
|
|
std::vector<StrategyPtr> sp_vector;
|
|
if (GenerateStrategiesWithBroadcast(stage_id, 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;
|
|
}
|
|
|
|
Status SoftmaxCrossEntropyWithLogitsInfo::SetCostUnderStrategy(const StrategyPtr& strategy) {
|
|
PrintStrategy(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;
|
|
}
|
|
} // namespace parallel
|
|
} // namespace mindspore
|