forked from mindspore-Ecosystem/mindspore
!49295 Add GenerateEodMask Operator AICPU
Merge pull request !49295 from huangxinjing/add_aicpu_mask_ms2
This commit is contained in:
Gitee
commit
6019ab4d7f
|
|
@ -0,0 +1,22 @@
|
|||
mindspore.ops.GenerateEodMask
|
||||
=============================
|
||||
|
||||
.. py:class:: mindspore.ops.GenerateEodMask(eod_token_id)
|
||||
|
||||
根据输入的 `inputs_ids`， 遇到 `eod_token_id` 时，会将输出的位置编码和注意力编码全部重置。
|
||||
即`position_id`从0开始重新计数,同时对应的掩码矩阵也会填充为0。
|
||||
|
||||
参数：
|
||||
- **eod_token_id** (int) - `eod_token_id` 的数值。在NLP场景中，这个值对应词表中的 `EodOfDocument` 的符号编码。
|
||||
|
||||
输入：
|
||||
- **inputs_ids** (Tensor) - 词序列。是一个二维Tensor，其shape为 :math:`(batch\_size, seq\_length)` 。
|
||||
|
||||
输出：
|
||||
- **position_id** (Tensor) - 位置编码矩阵。数据类型和shape与输入 `inputs_ids` 相同。
|
||||
- **attention_mask** (Tensor) - 注意力掩码矩阵。类型为float16，其shape为： :math:`(batch\_size, seq\_length)` 。
|
||||
|
||||
异常：
|
||||
- **TypeError** - 如果 `eod_token_id` 的数据类型不是int。
|
||||
- **TypeError** - 如果 `inputs_ids` 的数据类型不是整数类型。
|
||||
- **ValueError** - 如果 `inputs_ids` 不是二维的Tensor。
|
||||
|
|
@ -321,6 +321,7 @@ constexpr auto kEnvironGetOpName = "EnvironGet";
|
|||
constexpr auto kEnvironSetOpName = "EnvironSet";
|
||||
constexpr auto kEqualOpName = "Equal";
|
||||
constexpr auto kErfOpName = "Erf";
|
||||
constexpr auto kGenerateEodMaskOpName = "GenerateEodMask";
|
||||
constexpr auto kEuclideanNormOpName = "EuclideanNorm";
|
||||
constexpr auto kEuclideanNormDOpName = "EuclideanNormD";
|
||||
constexpr auto kExpandOpName = "Expand";
|
||||
|
|
|
|||
|
|
@ -0,0 +1,123 @@
|
|||
/**
|
||||
* Copyright 2023 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 "./generate_eod_mask_kernels.h"
|
||||
#include <Eigen/Dense>
|
||||
#include <map>
|
||||
#include <thread>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
#include <functional>
|
||||
#include "utils/kernel_util.h"
|
||||
#include "cpu_kernel_utils.h"
|
||||
|
||||
namespace {
|
||||
const char *kGenerateEodMask = "GenerateEodMask";
|
||||
constexpr auto kInputSize = 1;
|
||||
constexpr auto kOutputSize = 2;
|
||||
constexpr auto kInputIdsShape = 2;
|
||||
constexpr auto kAddressSize = 3;
|
||||
constexpr auto kDim0 = 0;
|
||||
constexpr auto kDim1 = 1;
|
||||
constexpr auto kDim2 = 2;
|
||||
constexpr auto kDim3 = 3;
|
||||
} // namespace
|
||||
namespace aicpu {
|
||||
uint32_t GenerateEodMaskCpuKernel::Compute(CpuKernelContext &ctx) {
|
||||
KERNEL_HANDLE_ERROR(NormalCheck(ctx, kInputSize, kOutputSize), "GenerateEodMaskCpu check input and output failed.");
|
||||
Tensor *input = ctx.Input(0);
|
||||
auto data_type_in = input->GetDataType();
|
||||
AttrValue *eod_token_value = ctx.GetAttr("eod_token_id");
|
||||
int64_t eod_token_id = (eod_token_value == nullptr) ? 0 : eod_token_value->GetInt();
|
||||
switch (data_type_in) {
|
||||
case DT_UINT16:
|
||||
return ComputeKernel<uint16_t>(ctx, eod_token_id);
|
||||
case DT_UINT32:
|
||||
return ComputeKernel<uint32_t>(ctx, eod_token_id);
|
||||
case DT_UINT64:
|
||||
return ComputeKernel<uint64_t>(ctx, eod_token_id);
|
||||
case DT_INT32:
|
||||
return ComputeKernel<int32_t>(ctx, eod_token_id);
|
||||
case DT_INT64:
|
||||
return ComputeKernel<int64_t>(ctx, eod_token_id);
|
||||
default:
|
||||
KERNEL_LOG_ERROR("GenerateEodMask kernel data type [%s] not support.", DTypeStr(data_type_in).c_str());
|
||||
return KERNEL_STATUS_PARAM_INVALID;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
uint32_t GenerateEodMaskCpuKernel::ComputeKernel(CpuKernelContext &ctx, const T &eod_token_id) {
|
||||
auto input_idsptr = reinterpret_cast<T *>(ctx.Input(0)->GetData());
|
||||
auto input_positionptr = reinterpret_cast<T *>(ctx.Output(0)->GetData());
|
||||
auto outputptr = reinterpret_cast<Eigen::half *>(ctx.Output(1)->GetData());
|
||||
auto output_shape = ctx.Output(1)->GetTensorShape()->GetDimSizes();
|
||||
size_t output_size =
|
||||
std::accumulate(output_shape.begin(), output_shape.end(), 1, std::multiplies<int64_t>()) * sizeof(Eigen::half);
|
||||
if (memset_s(outputptr, output_size, 0x0, output_size) != EOK) {
|
||||
KERNEL_LOG_ERROR("memset_s failed!");
|
||||
return KERNEL_STATUS_INNER_ERROR;
|
||||
}
|
||||
size_t batch_size = ctx.Input(0)->GetTensorShape()->GetDimSize(0);
|
||||
size_t seq_length = ctx.Input(0)->GetTensorShape()->GetDimSize(1);
|
||||
|
||||
auto shard_generate_tril = [&](size_t start, size_t end) {
|
||||
size_t x = seq_length * seq_length;
|
||||
for (size_t i = start; i < end; ++i) {
|
||||
for (size_t j = 0; j < seq_length; ++j) {
|
||||
for (size_t k = 0; k < j + 1; ++k) {
|
||||
outputptr[i * x + j * seq_length + k] = (Eigen::half)1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
auto shard_generate_eod_mask = [&](size_t start, size_t end) {
|
||||
for (size_t i = start; i < end; ++i) {
|
||||
T sub = 0;
|
||||
T pre_sub = 0;
|
||||
for (size_t index = 0; index < seq_length; ++index) {
|
||||
size_t sub_index = i * seq_length + index;
|
||||
if (input_idsptr[sub_index] == eod_token_id) {
|
||||
pre_sub = sub;
|
||||
sub = index + 1;
|
||||
size_t seq_n2 = seq_length * seq_length;
|
||||
for (size_t k = index + 1; k < seq_length; ++k) {
|
||||
for (size_t m = 0; m < index + 1; ++m) {
|
||||
outputptr[i * seq_n2 + k * seq_length + m] = (Eigen::half)0.0;
|
||||
}
|
||||
}
|
||||
input_positionptr[sub_index] = index - pre_sub;
|
||||
} else {
|
||||
input_positionptr[sub_index] = index - sub;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
auto get_per_unit_size = [&](int64_t data_size) -> int64_t {
|
||||
const int64_t max_core_num =
|
||||
std::max(static_cast<int64_t>(1), static_cast<int64_t>(aicpu::CpuKernelUtils::GetCPUNum(ctx) - 2));
|
||||
return data_size / std::min(max_core_num, data_size);
|
||||
};
|
||||
KERNEL_HANDLE_ERROR(CpuKernelUtils::ParallelFor(ctx, batch_size, get_per_unit_size(batch_size), shard_generate_tril),
|
||||
"GenerateEodMask kernel compute failed.");
|
||||
KERNEL_HANDLE_ERROR(
|
||||
CpuKernelUtils::ParallelFor(ctx, batch_size, get_per_unit_size(batch_size), shard_generate_eod_mask),
|
||||
"GenerateEodMask kernel compute failed.");
|
||||
return KERNEL_STATUS_OK;
|
||||
}
|
||||
REGISTER_CPU_KERNEL(kGenerateEodMask, GenerateEodMaskCpuKernel);
|
||||
} // namespace aicpu
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
/**
|
||||
* Copyright 2023 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 AICPU_KERNELS_GENERATEEODMASK_H_
|
||||
#define AICPU_KERNELS_GENERATEEODMASK_H_
|
||||
|
||||
#include <vector>
|
||||
#include "common/kernel_base.h"
|
||||
#include "cpu_ops_kernel.h"
|
||||
|
||||
namespace aicpu {
|
||||
class GenerateEodMaskCpuKernel : public CpuKernel {
|
||||
public:
|
||||
GenerateEodMaskCpuKernel() = default;
|
||||
~GenerateEodMaskCpuKernel() override = default;
|
||||
|
||||
protected:
|
||||
uint32_t Compute(CpuKernelContext &ctx) override;
|
||||
|
||||
private:
|
||||
template <typename T>
|
||||
uint32_t ComputeKernel(CpuKernelContext &ctx, const T &eod_token_id);
|
||||
};
|
||||
} // namespace aicpu
|
||||
#endif // AICPU_KERNELS_GENERATEEODMASK_H_
|
||||
|
|
@ -72,6 +72,7 @@ constexpr auto kConcatOffset = "ConcatOffset";
|
|||
constexpr auto kConcatOffsetV1 = "ConcatOffsetV1";
|
||||
constexpr auto kRandomChoiceWithMask = "RandomChoiceWithMask";
|
||||
constexpr auto kGatherDGradV2 = "GatherDGradV2";
|
||||
constexpr auto kGenerateEodMask = "GenerateEodMask";
|
||||
constexpr auto kResizeNearestNeighborV2 = "ResizeNearestNeighborV2";
|
||||
constexpr auto kResizeNearestNeighborV2Grad = "ResizeNearestNeighborV2Grad";
|
||||
constexpr auto kUpdateCache = "UpdateCache";
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ bool AICpuLibSelectPass::Process(const AnfNodePtr &node) const {
|
|||
mindspore::kComplexAbsOpName,
|
||||
mindspore::kConcatOpName,
|
||||
mindspore::kCosOpName,
|
||||
mindspore::kGenerateEodMaskOpName,
|
||||
mindspore::kCountNonZeroOpName,
|
||||
mindspore::kCumulativeLogsumexpOpName,
|
||||
mindspore::kCumProdOpName,
|
||||
|
|
|
|||
|
|
@ -765,6 +765,7 @@ GVAR_DEF(PrimitivePtr, kPrimNonZero, std::make_shared<Primitive>("NonZero"));
|
|||
GVAR_DEF(PrimitivePtr, kPrimNonZeroWithValue, std::make_shared<Primitive>("NonZeroWithValue"));
|
||||
GVAR_DEF(PrimitivePtr, kPrimNonZeroWithValueShape, std::make_shared<Primitive>("NonZeroWithValueShape"));
|
||||
GVAR_DEF(PrimitivePtr, kPrimNoRepeatNGram, std::make_shared<Primitive>("NoRepeatNGram"));
|
||||
GVAR_DEF(PrimitivePtr, kPrimGenerateEodMask, std::make_shared<Primitive>("GenerateEodMask"));
|
||||
GVAR_DEF(PrimitivePtr, kPrimRealInner, std::make_shared<Primitive>(kRealInner));
|
||||
GVAR_DEF(PrimitivePtr, kPrimReal, std::make_shared<Primitive>(kReal));
|
||||
GVAR_DEF(PrimitivePtr, kPrimImag, std::make_shared<Primitive>(kImag));
|
||||
|
|
|
|||
|
|
@ -0,0 +1,98 @@
|
|||
/**
|
||||
* Copyright 2023 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 "ops/generate_eod_mask.h"
|
||||
#include <string>
|
||||
#include <algorithm>
|
||||
#include <map>
|
||||
#include <set>
|
||||
#include <vector>
|
||||
|
||||
#include "ops/op_utils.h"
|
||||
#include "utils/check_convert_utils.h"
|
||||
#include "abstract/ops/primitive_infer_map.h"
|
||||
#include "mindapi/src/helper.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace ops {
|
||||
void GenerateEodMask::set_eod_token_id(const int64_t eod_token_id) {
|
||||
(void)this->AddAttr(kEodTokenId, api::MakeValue(eod_token_id));
|
||||
}
|
||||
/// \brief Get EodTokenId.
|
||||
///
|
||||
/// \return EodTokenId.
|
||||
int64_t GenerateEodMask::get_eod_token_id() const { return GetValue<int64_t>(GetAttr(kEodTokenId)); }
|
||||
|
||||
MIND_API_OPERATOR_IMPL(GenerateEodMask, BaseOperator);
|
||||
|
||||
// AG means auto generated
|
||||
class MIND_API AGGenerateEodMaskInfer : public abstract::OpInferBase {
|
||||
public:
|
||||
BaseShapePtr InferShape(const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
auto prim_name = primitive->name();
|
||||
for (const auto &item : input_args) {
|
||||
MS_EXCEPTION_IF_NULL(item);
|
||||
}
|
||||
const int64_t no_repeat_kShapeSize = 2;
|
||||
auto inputs_ids_shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape());
|
||||
auto inputs_ids_shape = inputs_ids_shape_map[kShape];
|
||||
|
||||
if (IsDynamicRank(inputs_ids_shape)) {
|
||||
auto any_shape = std::make_shared<abstract::Shape>(std::vector<int64_t>{abstract::Shape::kShapeRankAny});
|
||||
std::vector<BaseShapePtr> shapes_list = {any_shape, any_shape};
|
||||
return std::make_shared<abstract::TupleShape>(shapes_list);
|
||||
}
|
||||
|
||||
ShapeVector attention_mask_shape{inputs_ids_shape.begin(), inputs_ids_shape.end()};
|
||||
|
||||
attention_mask_shape.push_back(inputs_ids_shape.back());
|
||||
(void)CheckAndConvertUtils::CheckInteger("rank of inputs_ids", SizeToLong(inputs_ids_shape.size()), kEqual,
|
||||
no_repeat_kShapeSize, prim_name);
|
||||
|
||||
std::vector<BaseShapePtr> shapes_list = {};
|
||||
(void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(inputs_ids_shape));
|
||||
(void)shapes_list.emplace_back(std::make_shared<abstract::Shape>(attention_mask_shape));
|
||||
return std::make_shared<abstract::TupleShape>(shapes_list);
|
||||
}
|
||||
|
||||
TypePtr InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
for (const auto &item : input_args) {
|
||||
MS_EXCEPTION_IF_NULL(item);
|
||||
}
|
||||
std::map<std::string, TypePtr> input_types;
|
||||
auto input_ids_type = input_args[0]->BuildType();
|
||||
(void)input_types.emplace("inputs_ids", input_ids_type);
|
||||
std::set<TypePtr> valid_input_types = {kInt16, kInt32, kInt64, kUInt16, kUInt32, kUInt64};
|
||||
(void)CheckAndConvertUtils::CheckTensorTypeSame(input_types, valid_input_types, primitive->name());
|
||||
return std::make_shared<Tuple>(std::vector<TypePtr>{input_ids_type, kFloat16});
|
||||
}
|
||||
AbstractBasePtr InferShapeAndType(const abstract::AnalysisEnginePtr &engine, const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
const int64_t kInputsNum = 1;
|
||||
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, kInputsNum, primitive->name());
|
||||
auto type = InferType(primitive, input_args);
|
||||
auto shape = InferShape(primitive, input_args);
|
||||
return abstract::MakeAbstract(shape, type);
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_PRIMITIVE_OP_INFER_IMPL(GenerateEodMask, prim::kPrimGenerateEodMask, AGGenerateEodMaskInfer, false);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
/**
|
||||
* Copyright 2023 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_GENERATE_EOD_MASK_H
|
||||
#define MINDSPORE_GENERATE_EOD_MASK_H
|
||||
#include <map>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <memory>
|
||||
|
||||
#include "ops/base_operator.h"
|
||||
#include "mindapi/base/types.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace ops {
|
||||
constexpr auto kNameGenerateEodMask = "GenerateEodMask";
|
||||
/// \brief
|
||||
class MIND_API GenerateEodMask : public BaseOperator {
|
||||
public:
|
||||
MIND_API_BASE_MEMBER(GenerateEodMask);
|
||||
/// \brief Constructor.
|
||||
GenerateEodMask() : BaseOperator(kNameGenerateEodMask) {
|
||||
InitIOName({"inputs_ids"}, {"position_ids", "attention_mask"});
|
||||
}
|
||||
/// \brief Init.
|
||||
void Init() const {}
|
||||
/// \brief Set axis.
|
||||
void set_eod_token_id(const int64_t eod_token_id);
|
||||
/// \brief Get axis.
|
||||
///
|
||||
/// \return axis.
|
||||
int64_t get_eod_token_id() const;
|
||||
};
|
||||
|
||||
MIND_API abstract::AbstractBasePtr GenerateEodMaskInfer(const abstract::AnalysisEnginePtr &,
|
||||
const PrimitivePtr &primitive,
|
||||
const std::vector<abstract::AbstractBasePtr> &input_args);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_GENERATE_EOD_MASK_H
|
||||
|
|
@ -149,6 +149,7 @@ constexpr auto kNeginf = "neginf";
|
|||
constexpr auto kNesterov = "nesterov";
|
||||
constexpr auto kNewAxisMask = "new_axis_mask";
|
||||
constexpr auto kNgramSize = "ngram_size";
|
||||
constexpr auto kEodTokenId = "eod_token_id";
|
||||
constexpr auto kNGram = "ngram";
|
||||
constexpr auto kNmsThresh = "nms_thresh";
|
||||
constexpr auto kNormRegion = "norm_region";
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ from .init_data_set_queue import _init_data_set_queue_aicpu
|
|||
from .embedding_lookup import _embedding_lookup_aicpu
|
||||
from .padding import _padding_aicpu
|
||||
from .gather import _gather_aicpu
|
||||
from .generate_eod_mask import _generate_eod_mask_aicpu
|
||||
from .gather_grad import _gather_grad_aicpu
|
||||
from .gather_d_grad_v2 import _gather_d_grad_v2_aicpu
|
||||
from .gather_d import _gather_d_aicpu
|
||||
|
|
|
|||
|
|
@ -0,0 +1,36 @@
|
|||
# Copyright 2023 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.
|
||||
# ============================================================================
|
||||
|
||||
"""GenerateEodMask op"""
|
||||
from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
|
||||
|
||||
generate_eod_mask_op_info = AiCPURegOp("GenerateEodMask") \
|
||||
.fusion_type("OPAQUE") \
|
||||
.attr("eod_token_id", "int") \
|
||||
.input(0, "inputs_ids", "required") \
|
||||
.output(0, "position_ids", "required") \
|
||||
.output(1, "attention_mask", "required") \
|
||||
.dtype_format(DataType.U16_Default, DataType.U16_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.U32_Default, DataType.U32_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.U64_Default, DataType.U64_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.F16_Default) \
|
||||
.get_op_info()
|
||||
|
||||
|
||||
@op_info_register(generate_eod_mask_op_info)
|
||||
def _generate_eod_mask_aicpu():
|
||||
"""GenerateEodMask AiCPU register"""
|
||||
return
|
||||
|
|
@ -64,7 +64,7 @@ from .image_ops import (CropAndResize, NonMaxSuppressionV3, HSVToRGB, AdjustHue,
|
|||
CombinedNonMaxSuppression, RGBToHSV, ScaleAndTranslate, ResizeLinear1D, ResizeBicubic)
|
||||
from .inner_ops import (ScalarCast, Randperm, NoRepeatNGram, LambApplyOptimizerAssign, LambApplyWeightAssign,
|
||||
FusedWeightScaleApplyMomentum, FusedCastAdamWeightDecay, FusedAdaFactor,
|
||||
FusedAdaFactorWithGlobalNorm)
|
||||
FusedAdaFactorWithGlobalNorm, GenerateEodMask)
|
||||
from .linalg_ops import (Svd, Geqrf)
|
||||
from .math_ops import (Abs, ACos, Asin, Asinh, AddN, AccumulateNV2, AssignAdd, AssignSub, Atan2, BatchMatMul,
|
||||
BitwiseAnd, BitwiseOr, Ger, BitwiseXor, Inv, Invert, ApproximateEqual,
|
||||
|
|
@ -170,6 +170,7 @@ __all__ = [
|
|||
'Conv3D',
|
||||
'Conv2DTranspose',
|
||||
'Conv3DTranspose',
|
||||
"GenerateEodMask",
|
||||
'FillV2',
|
||||
'Flatten',
|
||||
'MaxPoolWithArgmax',
|
||||
|
|
|
|||
|
|
@ -655,6 +655,52 @@ class FusedAdaFactorWithGlobalNorm(FusedAdaFactor):
|
|||
return param_type
|
||||
|
||||
|
||||
class GenerateEodMask(Primitive):
|
||||
r"""
|
||||
Given the input `inputs_ids`, if found eod_token_id, the output position and attention mask matrix will be reset.
|
||||
This means the `position_id` will start counting from 0, and the corresponding mask matrix will be filled with 0.
|
||||
|
||||
Args:
|
||||
eod_token_id (int) - In the NLP scenario, this value corresponds to the id of
|
||||
the symbol of 'EodOfDocument' in the vocabulary.
|
||||
|
||||
Inputs:
|
||||
- **inputs_ids** (Tensor) - token id, a 2-D Tensor with shape :math:`(batch\_size, seq\_length)`.
|
||||
|
||||
Outputs:
|
||||
- **position_id** (Tensor) - position id matrix with same shape and type as original `inputs_ids`.
|
||||
- **attention_mask** (Tensor) - attention mask matrix with type
|
||||
float16 and shape :math:`(batch\_size, seq\_length)`.
|
||||
|
||||
Supported Platforms:
|
||||
``Ascend``
|
||||
|
||||
Examples:
|
||||
>>> op = ops.GenerateEodMask(eod_token_id=0)
|
||||
>>> position, mask = op(Tensor([[1, 0, 3], [1, 0, 0]], dtype=mindspore.int32))
|
||||
>>> print(position)
|
||||
[[0 1 0] [0 0 1]]
|
||||
>>> print(mask)
|
||||
[[[ 1. 0. 0.]
|
||||
[1. 1. 0.]
|
||||
[0. 0. 1.]]
|
||||
[[1. 0. 0.]
|
||||
[0. 1. 0.]
|
||||
[0. 1. 1.]]]
|
||||
|
||||
Raises:
|
||||
- **TypeError** - If `eod_token_id` is not int.
|
||||
- **TypeError** - If `inputs_ids` is not int.
|
||||
- **ValueError** - If `inputs_ids` is not a 2-D Tensor.
|
||||
"""
|
||||
@prim_attr_register
|
||||
def __init__(self, eod_token_id):
|
||||
"""Initialize GenerateEodMask"""
|
||||
validator.check_value_type("eod_token_id", eod_token_id, [int], self.name)
|
||||
self.init_prim_io_names(inputs=['inputs_ids'],
|
||||
outputs=['position_ids', 'attention_mask'])
|
||||
|
||||
|
||||
class ScaleGrad(PrimitiveWithInfer):
|
||||
"""
|
||||
Scale the input grad according to the loss scale.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,143 @@
|
|||
# Copyright 2023 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 pytest
|
||||
import numpy as np
|
||||
|
||||
import mindspore
|
||||
import mindspore.context as context
|
||||
import mindspore.nn as nn
|
||||
from mindspore import Tensor
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
|
||||
|
||||
|
||||
class Net(nn.Cell):
|
||||
def __init__(self, eod_token_id):
|
||||
super(Net, self).__init__()
|
||||
self.mask = P.GenerateEodMask(eod_token_id=eod_token_id)
|
||||
|
||||
def construct(self, tensor):
|
||||
return self.mask(tensor)
|
||||
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_generate_eod_mask():
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test eod_token_id 0
|
||||
Expectation: raise TypeError.
|
||||
"""
|
||||
x = np.array([[1, 0, 3, 4, 0, 6, 7, 8], [1, 0, 3, 0, 0, 6, 7, 0]])
|
||||
net = Net(0)
|
||||
position, mask = net(Tensor(x, dtype=mindspore.int32))
|
||||
assert position.shape == (2, 8)
|
||||
assert mask.shape == (2, 8, 8)
|
||||
assert np.all(mask.asnumpy() == np.array([[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]],
|
||||
[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]]]))
|
||||
|
||||
assert np.all(position.asnumpy() == np.array([[0, 1, 0, 1, 2, 0, 1, 2],
|
||||
[0, 1, 0, 1, 0, 0, 1, 2]]))
|
||||
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_generate_eod_mask_negative_value():
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test EodTokenId Negative
|
||||
Expectation: no errors
|
||||
"""
|
||||
x = np.array([[1, -1, 3, 4, -1, 6, 7, 8], [1, -1, 3, -1, -1, 6, 7, -1]])
|
||||
net = Net(eod_token_id=-1)
|
||||
position, mask = net(Tensor(x, dtype=mindspore.int32))
|
||||
assert position.shape == (2, 8)
|
||||
assert mask.shape == (2, 8, 8)
|
||||
assert np.all(mask.asnumpy() == np.array([[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]],
|
||||
[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]]]))
|
||||
|
||||
assert np.all(position.asnumpy() == np.array([[0, 1, 0, 1, 2, 0, 1, 2],
|
||||
[0, 1, 0, 1, 0, 0, 1, 2]]))
|
||||
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_generate_eod_mask_dynamic_inputs():
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test dynamic inputs
|
||||
Expectation: No error.
|
||||
"""
|
||||
x = np.array([[1, -1, 3, 4, -1, 6, 7, 8], [1, -1, 3, -1, -1, 6, 7, -1]])
|
||||
net = Net(eod_token_id=-1)
|
||||
dyn_x = Tensor(shape=(None, None), dtype=mindspore.int32)
|
||||
net.set_inputs(dyn_x)
|
||||
position, mask = net(Tensor(x, dtype=mindspore.int32))
|
||||
assert position.shape == (2, 8)
|
||||
assert mask.shape == (2, 8, 8)
|
||||
assert np.all(mask.asnumpy() == np.array([[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]],
|
||||
[[1, 0, 0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 0, 0, 0, 0, 0],
|
||||
[0, 0, 1, 1, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 1, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 0, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 0],
|
||||
[0, 0, 0, 0, 0, 1, 1, 1]]]))
|
||||
|
||||
assert np.all(position.asnumpy() == np.array([[0, 1, 0, 1, 2, 0, 1, 2],
|
||||
[0, 1, 0, 1, 0, 0, 1, 2]]))
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
# Copyright 2023 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 pytest
|
||||
import numpy as np
|
||||
|
||||
import mindspore
|
||||
import mindspore.context as context
|
||||
import mindspore.nn as nn
|
||||
from mindspore import Tensor
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
|
||||
|
||||
|
||||
class Net(nn.Cell):
|
||||
def __init__(self, eod_token_id):
|
||||
super(Net, self).__init__()
|
||||
self.mask = P.GenerateEodMask(eod_token_id=eod_token_id)
|
||||
|
||||
def construct(self, tensor):
|
||||
return self.mask(tensor)
|
||||
|
||||
|
||||
@pytest.mark.level1
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_generate_eod_mask_wrong_type():
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test float_token_id
|
||||
Expectation: raise TypeError.
|
||||
"""
|
||||
x = np.array([[1, 0, 3, 4, 0, 6, 7, 8], [1, 0, 3, 0, 0, 6, 7, 0]])
|
||||
net = Net(0)
|
||||
with pytest.raises(TypeError):
|
||||
net(Tensor(x, dtype=mindspore.float32))
|
||||
|
||||
|
||||
|
||||
@pytest.mark.level1
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
def test_generate_eod_mask_float_token_id():
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test float_token_id
|
||||
Expectation: raise TypeError.
|
||||
"""
|
||||
x = np.array([[1, 0, 3, 4, 0, 6, 7, 8], [1, 0, 3, 0, 0, 6, 7, 0]])
|
||||
with pytest.raises(TypeError):
|
||||
net = Net(-1.0)
|
||||
net(Tensor(x, dtype=mindspore.float32))
|
||||
|
||||
|
||||
@pytest.mark.level1
|
||||
@pytest.mark.platform_arm_ascend_training
|
||||
@pytest.mark.platform_x86_ascend_training
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('dtype', [mindspore.int32, mindspore.int64, mindspore.uint16,
|
||||
mindspore.uint32, mindspore.uint64])
|
||||
def test_generate_eod_mask_support_dtype(dtype):
|
||||
"""
|
||||
Feature: Test GenerateEodMask.
|
||||
Description: Test multi dtype inputs
|
||||
Expectation: Successful graph compilation.
|
||||
"""
|
||||
x = np.array([[1, 0, 3, 4, 0, 6, 7, 8], [1, 0, 3, 0, 0, 6, 7, 0]])
|
||||
net = Net(0)
|
||||
net(Tensor(x, dtype=dtype))
|
||||
Loading…
Reference in New Issue