146 lines
6.9 KiB
C++
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
|