Files
deeplearning4j--deeplearning4j/libnd4j/include/ops/declarable/generic/decoder/ctc_beam_op.cpp
T
2026-07-13 12:47:05 +08:00

146 lines
6.9 KiB
C++

/*******************************************************************************
* Copyright (c) 2021 Deeplearning4j Contributors
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://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.
*
* SPDX-License-Identifier: Apache-2.0
*******************************************************************************/
//
// @author AbdelRauf
//
#include <ops/declarable/CustomOperations.h>
#include <ops/declarable/helpers/ctc.h>
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_ctc_beam)
namespace sd {
namespace ops {
//////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(ctc_beam, 2, 3, false, 0, -2) {
auto logit = INPUT_VARIABLE(0);
auto sequence_length = INPUT_VARIABLE(1);
auto result_sequences = OUTPUT_VARIABLE(0);
auto result_probs = OUTPUT_VARIABLE(1);
auto result_sequences_length = OUTPUT_VARIABLE(2);
auto arg_size = block.getIArguments()->size();
auto normalize_logits = block.numB() > 0 ? B_ARG(0) : false;
int blank_index = arg_size > 0 ? INT_ARG(0) : -1;
int beam_width = arg_size > 1 ? INT_ARG(1) : 25;
int nbest_len = arg_size > 2 ? INT_ARG(2) : 1;
REQUIRE_TRUE(
logit->rankOf() == 3, 0,
"Ctc Beam Search: logit Input fails to meet rank requirement {BATCH_LEN, MAX_FRAME_LEN, CLASS_LEN }: %i == 3 ",
logit->rankOf());
REQUIRE_TRUE(sequence_length->rankOf() == 1, 0,
"Ctc Beam Search: sequence frame length (sequence_length) Input fails to meet rank requirement "
"{BATCH_LEN}: %i == 1 ",
sequence_length->rankOf());
REQUIRE_TRUE(result_sequences->rankOf() == 3, 0,
"Ctc Beam Search: result_sequences Output fails to meet rank requirement {BATCH_LEN, NBEST_LEN, "
"MAX_FRAME_LEN }: %i == 3 ",
result_sequences->rankOf());
REQUIRE_TRUE(result_probs->rankOf() == 2, 0,
"Ctc Beam Search: result_probs Output fails to meet rank requirement {BATCH_LEN, NBEST_LEN}: %i == 2 ",
result_probs->rankOf());
REQUIRE_TRUE(
result_sequences_length->rankOf() == 2, 0,
"Ctc Beam Search: result_sequences_length Output fails to meet rank requirement {BATCH_LEN, NBEST_LEN}: %i == 2 ",
result_sequences_length->rankOf());
auto batchSize0 = logit->sizeAt(0);
auto batchSize1 = sequence_length->sizeAt(0);
auto batchSize2 = result_sequences->sizeAt(0);
auto batchSize3 = result_probs->sizeAt(0);
auto batchSize4 = result_sequences_length->sizeAt(0);
bool check_batches = (batchSize0 == batchSize1) && (batchSize2 == batchSize3);
check_batches = check_batches && (batchSize0 == batchSize4) && (batchSize0 == batchSize2);
REQUIRE_TRUE(nbest_len > 0 && nbest_len <= beam_width, 0, "Ctc Beam Search: nbest_len %i should be > 0 and <= %i",
nbest_len, beam_width);
REQUIRE_TRUE(check_batches, 0, "Ctc Beam Search: All batch sizes should be %i", batchSize0);
auto max_t = logit->sizeAt(1);
REQUIRE_TRUE(result_sequences->sizeAt(1) == nbest_len && result_sequences->sizeAt(2) == max_t, 0,
"Ctc Beam Search: shape of the result_sequences should be {%i, %i, %i} but got { %i, %i, %i}",
batchSize0, nbest_len, max_t, batchSize1, result_sequences->sizeAt(1), result_sequences->sizeAt(2));
REQUIRE_TRUE(result_probs->sizeAt(1) == nbest_len, 0,
"Ctc Beam Search: shape of the result_probs should be {%i, %i} but got { %i, %i}", batchSize0, nbest_len,
batchSize3, result_sequences->sizeAt(1));
REQUIRE_TRUE(result_sequences_length->sizeAt(1) == nbest_len, 0,
"Ctc Beam Search: shape of the result_sequences_length should be {%i, %i} but got { %i, %i}", batchSize0,
nbest_len, batchSize4, result_sequences_length->sizeAt(1));
REQUIRE_TRUE(result_sequences->ews() == 1 && result_sequences->ordering() == 'c', 0,
"Ctc Beam Search: result_sequences output should be ews()==1 and c order: %d == ews(1) %c == order(c) ",
result_sequences->ews(), result_sequences->ordering());
REQUIRE_TRUE(result_probs->ews() == 1 && result_probs->ordering() == 'c', 0,
"Ctc Beam Search: result_probs output should be ews()==1 and c order: %d == ews(1) %c == order(c) ",
result_probs->ews(), result_probs->ordering());
REQUIRE_TRUE(
result_sequences_length->ews() == 1 && result_sequences_length->ordering() == 'c', 0,
"Ctc Beam Search: result_sequences_length output should be ews()==1 and c order: %d == ews(1) %c == order(c) ",
result_sequences_length->ews(), result_sequences_length->ordering());
helpers::beamSearch(*logit, *sequence_length, *result_sequences, *result_probs, *result_sequences_length,
blank_index, beam_width, nbest_len, normalize_logits);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
DECLARE_TYPES(ctc_beam) {
getOpDescriptor()
->setAllowedInputTypes({ALL_FLOATS})
->setAllowedInputTypes(1, {ALL_INDICES})
->setAllowedOutputTypes(0, {ALL_INDICES})
->setAllowedOutputTypes(1, {ALL_FLOATS})
->setAllowedOutputTypes(2, {ALL_INDICES});
}
//////////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(ctc_beam) {
auto logitShapeInfo = inputShape->at(0);
auto sequenceShapeInfo = inputShape->at(1);
auto arg_size = block.getIArguments()->size();
auto nbest_len = arg_size > 2 ? INT_ARG(2) : 1;
REQUIRE_TRUE(
logitShapeInfo[0] == 3, 0,
"Ctc Beam Search: logit Input fails to meet rank requirement {BATCH_LEN, MAX_FRAME_LEN, CLASS_LEN }: %i == 3",
logitShapeInfo[0]);
auto batch_size = shape::shapeOf(logitShapeInfo)[0];
auto max_t = shape::shapeOf(logitShapeInfo)[1];
auto dtype_float = ArrayOptions::dataType(logitShapeInfo);
auto dtype_index = ArrayOptions::dataType(sequenceShapeInfo);
auto desc = ShapeBuilders::createShapeInfo(dtype_index, 'c', {batch_size, nbest_len, max_t});
auto output0 = ConstantShapeHelper::getInstance().bufferForShapeInfo(desc);
auto desc2 = ShapeBuilders::createShapeInfo(dtype_float, 'c', {batch_size, nbest_len});
auto output1 =
ConstantShapeHelper::getInstance().bufferForShapeInfo(desc2);
auto desc3 = ShapeBuilders::createShapeInfo(dtype_index, 'c', {batch_size, nbest_len});
auto output2 = ConstantShapeHelper::getInstance().bufferForShapeInfo(desc3);
return SHAPELIST(output0->primary(), output1->primary(), output2->primary());
}
} // namespace ops
} // namespace sd
#endif