EngineX-Iluvatar/enginex-mr_series-sherpa-onnx
Archived
16
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Fix batch decoding for greedy search (#71)

Browse Source
This commit is contained in:
Fangjun Kuang
2023-03-01 12:18:20 +08:00
committed by GitHub
parent ebf5855290
commit e0b76655c8
1 changed files with 25 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
sherpa-onnx/csrc/online-transducer-greedy-search-decoder.cc
View File

@@ -10,47 +10,39 @@
#include <utility> #include <utility>
#include <vector> #include <vector>
#include "sherpa-onnx/csrc/macros.h"
#include "sherpa-onnx/csrc/onnx-utils.h" #include "sherpa-onnx/csrc/onnx-utils.h"
namespace sherpa_onnx { namespace sherpa_onnx {
static Ort::Value GetFrame(Ort::Value *encoder_out, int32_t t) { static Ort::Value GetFrame(OrtAllocator *allocator, Ort::Value *encoder_out,
int32_t t) {
std::vector<int64_t> encoder_out_shape = std::vector<int64_t> encoder_out_shape =
encoder_out->GetTensorTypeAndShapeInfo().GetShape(); encoder_out->GetTensorTypeAndShapeInfo().GetShape();
assert(encoder_out_shape[0] == 1);
int32_t encoder_out_dim = encoder_out_shape[2]; auto batch_size = encoder_out_shape[0];
auto num_frames = encoder_out_shape[1];
assert(t < num_frames);
auto encoder_out_dim = encoder_out_shape[2];
auto offset = num_frames * encoder_out_dim;
auto memory_info = auto memory_info =
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
std::array<int64_t, 2> shape{1, encoder_out_dim}; std::array<int64_t, 2> shape{batch_size, encoder_out_dim};
return Ort::Value::CreateTensor( Ort::Value ans =
memory_info, Ort::Value::CreateTensor<float>(allocator, shape.data(), shape.size());
encoder_out->GetTensorMutableData<float>() + t * encoder_out_dim,
encoder_out_dim, shape.data(), shape.size());
}
static Ort::Value Repeat(OrtAllocator *allocator, Ort::Value *cur_encoder_out,
int32_t n) {
if (n == 1) {
return std::move(*cur_encoder_out);
}
std::vector<int64_t> cur_encoder_out_shape =
cur_encoder_out->GetTensorTypeAndShapeInfo().GetShape();
std::array<int64_t, 2> ans_shape{n, cur_encoder_out_shape[1]};
Ort::Value ans = Ort::Value::CreateTensor<float>(allocator, ans_shape.data(),
ans_shape.size());
const float *src = cur_encoder_out->GetTensorData<float>();
float *dst = ans.GetTensorMutableData<float>(); float *dst = ans.GetTensorMutableData<float>();
for (int32_t i = 0; i != n; ++i) { const float *src = encoder_out->GetTensorData<float>();
std::copy(src, src + cur_encoder_out_shape[1], dst);
dst += cur_encoder_out_shape[1]; for (int32_t i = 0; i != batch_size; ++i) {
std::copy(src + t * encoder_out_dim, src + (t + 1) * encoder_out_dim, dst);
src += offset;
dst += encoder_out_dim;
} }
return ans; return ans;
@@ -83,10 +75,10 @@ void OnlineTransducerGreedySearchDecoder::Decode(
encoder_out.GetTensorTypeAndShapeInfo().GetShape(); encoder_out.GetTensorTypeAndShapeInfo().GetShape();
if (encoder_out_shape[0] != result->size()) { if (encoder_out_shape[0] != result->size()) {
fprintf(stderr, SHERPA_ONNX_LOGE(
"Size mismatch! encoder_out.size(0) %d, result.size(0): %d\n", "Size mismatch! encoder_out.size(0) %d, result.size(0): %d",
static_cast<int32_t>(encoder_out_shape[0]), static_cast<int32_t>(encoder_out_shape[0]),
static_cast<int32_t>(result->size())); static_cast<int32_t>(result->size()));
exit(-1); exit(-1);
} }
@@ -98,10 +90,10 @@ void OnlineTransducerGreedySearchDecoder::Decode(
Ort::Value decoder_out = model_->RunDecoder(std::move(decoder_input)); Ort::Value decoder_out = model_->RunDecoder(std::move(decoder_input));
for (int32_t t = 0; t != num_frames; ++t) { for (int32_t t = 0; t != num_frames; ++t) {
Ort::Value cur_encoder_out = GetFrame(&encoder_out, t); Ort::Value cur_encoder_out = GetFrame(model_->Allocator(), &encoder_out, t);
cur_encoder_out = Repeat(model_->Allocator(), &cur_encoder_out, batch_size);
Ort::Value logit = model_->RunJoiner( Ort::Value logit = model_->RunJoiner(
std::move(cur_encoder_out), Clone(model_->Allocator(), &decoder_out)); std::move(cur_encoder_out), Clone(model_->Allocator(), &decoder_out));
const float *p_logit = logit.GetTensorData<float>(); const float *p_logit = logit.GetTensorData<float>();
bool emitted = false; bool emitted = false;