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

Add PackPaddedSequence (#85)

Browse Source
This commit is contained in:
Fangjun Kuang
2023-03-08 14:12:20 +08:00
committed by GitHub
parent 3a79115884
commit 8c6a289e3d
7 changed files with 236 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sherpa-onnx/csrc/CMakeLists.txt
View File

@@ -16,6 +16,7 @@ set(sources
online-transducer-modified-beam-search-decoder.cc
online-zipformer-transducer-model.cc
onnx-utils.cc
packed-sequence.cc
pad-sequence.cc
parse-options.cc
resample.cc
@@ -123,6 +124,7 @@ endif()
if(SHERPA_ONNX_ENABLE_TESTS)
set(sherpa_onnx_test_srcs
cat-test.cc
packed-sequence-test.cc
pad-sequence-test.cc
slice-test.cc
transpose-test.cc

52
sherpa-onnx/csrc/packed-sequence-test.cc Normal file
View File

@@ -0,0 +1,52 @@
// sherpa-onnx/csrc/packed-sequence-test.cc
//
// Copyright (c) 2023 Xiaomi Corporation
#include "sherpa-onnx/csrc/packed-sequence.h"
#include <numeric>
#include "gtest/gtest.h"
#include "sherpa-onnx/csrc/onnx-utils.h"
namespace sherpa_onnx {
TEST(PackedSequence, Case1) {
Ort::AllocatorWithDefaultOptions allocator;
std::array<int64_t, 3> shape{5, 5, 4};
Ort::Value v =
Ort::Value::CreateTensor<float>(allocator, shape.data(), shape.size());
float *p = v.GetTensorMutableData<float>();
std::iota(p, p + shape[0] * shape[1] * shape[2], 0);
Ort::Value length =
Ort::Value::CreateTensor<int64_t>(allocator, shape.data(), 1);
int64_t *p_length = length.GetTensorMutableData<int64_t>();
p_length[0] = 1;
p_length[1] = 2;
p_length[2] = 3;
p_length[3] = 5;
p_length[4] = 2;
auto packed_seq = PackPaddedSequence(allocator, &v, &length);
fprintf(stderr, "sorted indexes: ");
for (auto i : packed_seq.sorted_indexes) {
fprintf(stderr, "%d ", static_cast<int32_t>(i));
}
fprintf(stderr, "\n");
// output index: 0 1 2 3 4
// sorted indexes: 3 2 1 4 0
// length: 5 3 2 2 1
Print3D(&v);
Print2D(&packed_seq.data);
fprintf(stderr, "batch sizes per time step: ");
for (auto i : packed_seq.batch_sizes) {
fprintf(stderr, "%d ", static_cast<int32_t>(i));
}
fprintf(stderr, "\n");
// TODO(fangjun): Check that the return value is correct
}
} // namespace sherpa_onnx

107
sherpa-onnx/csrc/packed-sequence.cc Normal file
View File

@@ -0,0 +1,107 @@
// sherpa-onnx/csrc/packed-sequence.cc
//
// Copyright (c) 2023 Xiaomi Corporation
#include "sherpa-onnx/csrc/packed-sequence.h"
#include <assert.h>
#include <algorithm>
#include <numeric>
#include <utility>
#include "sherpa-onnx/csrc/slice.h"
#include "sherpa-onnx/csrc/transpose.h"
namespace sherpa_onnx {
static Ort::Value IndexSelect(OrtAllocator *allocator, const Ort::Value *value,
const std::vector<int32_t> &sorted_indexes) {
auto shape = value->GetTensorTypeAndShapeInfo().GetShape();
assert(shape.size() == 3);
std::array<int64_t, 3> ans_shape{static_cast<int64_t>(sorted_indexes.size()),
shape[1], shape[2]};
Ort::Value ans = Ort::Value::CreateTensor<float>(allocator, ans_shape.data(),
ans_shape.size());
float *dst = ans.GetTensorMutableData<float>();
const float *src = value->GetTensorData<float>();
for (auto i : sorted_indexes) {
const float *start = src + i * shape[1] * shape[2];
std::copy(start, start + shape[1] * shape[2], dst);
dst += shape[1] * shape[2];
}
return ans;
}
PackedSequence PackPaddedSequence(OrtAllocator *allocator,
const Ort::Value *value, Ort::Value *length) {
std::vector<int64_t> v_shape = value->GetTensorTypeAndShapeInfo().GetShape();
std::vector<int64_t> l_shape = length->GetTensorTypeAndShapeInfo().GetShape();
assert(v_shape.size() == 3);
assert(l_shape.size() == 3);
assert(v_shape[0] == l_shape[0]);
std::vector<int32_t> indexes(v_shape[0]);
std::iota(indexes.begin(), indexes.end(), 0);
const int64_t *p_length = length->GetTensorData<int64_t>();
// sort in descending order
std::sort(indexes.begin(), indexes.end(), [p_length](int32_t i, int32_t j) {
return p_length[i] > p_length[j];
});
int32_t n = static_cast<int32_t>(v_shape[0]);
int64_t max_T = p_length[indexes[0]];
int32_t sum_T = std::accumulate(p_length, p_length + n, 0);
std::array<int64_t, 2> data_shape{sum_T, v_shape[2]};
Ort::Value data = Ort::Value::CreateTensor<float>(
allocator, data_shape.data(), data_shape.size());
float *dst = data.GetTensorMutableData<float>();
Ort::Value tensor = IndexSelect(allocator, value, indexes);
tensor = Transpose01(allocator, &tensor);
// batch size at each time step
std::vector<int32_t> batch_sizes;
batch_sizes.reserve(max_T);
int64_t prev_l = 0;
for (int32_t i = 0; i != n; ++i) {
auto cur_l = p_length[indexes[n - 1 - i]];
assert(cur_l >= prev_l);
if (cur_l == prev_l) {
continue;
}
auto cur_batch_size = n - i;
Ort::Value cur_batch =
Slice(allocator, &tensor, prev_l, cur_l, 0, cur_batch_size);
auto count = cur_batch.GetTensorTypeAndShapeInfo().GetElementCount();
const float *src = cur_batch.GetTensorData<float>();
std::copy(src, src + count, dst);
dst += count;
for (int32_t j = prev_l; j < cur_l; ++j) {
batch_sizes.push_back(cur_batch_size);
}
prev_l = cur_l;
}
PackedSequence packed_seq;
packed_seq.sorted_indexes = std::move(indexes);
packed_seq.data = std::move(data);
packed_seq.batch_sizes = std::move(batch_sizes);
return packed_seq;
}
} // namespace sherpa_onnx

33
sherpa-onnx/csrc/packed-sequence.h Normal file
View File

@@ -0,0 +1,33 @@
// sherpa-onnx/csrc/packed-sequence.h
//
// Copyright (c) 2023 Xiaomi Corporation
#ifndef SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_
#define SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_
#include <vector>
#include "onnxruntime_cxx_api.h" // NOLINT
namespace sherpa_onnx {
struct PackedSequence {
std::vector<int32_t> sorted_indexes;
std::vector<int32_t> batch_sizes;
Ort::Value data{nullptr};
};
/** Similar to torch.nn.utils.rnn.pad_sequence but it supports only
* batch_first=true.
*
* @param allocator
* @param value A 3-D tensor of shape (B, T, C). Its dtype is float.
* @param length A 1-D tensor of shape (B,). Its dtype is int64_t. Each
* element in it specifies the valid length of the corresponding
* entry in value before padding.
*/
PackedSequence PackPaddedSequence(OrtAllocator *allocator,
const Ort::Value *value, Ort::Value *length);
} // namespace sherpa_onnx
#endif // SHERPA_ONNX_CSRC_PACKED_SEQUENCE_H_

10
sherpa-onnx/csrc/slice-test.cc
View File

@@ -13,19 +13,19 @@ namespace sherpa_onnx {
TEST(Slice, Slice3D) {
Ort::AllocatorWithDefaultOptions allocator;
std::array<int64_t, 3> shape{3, 5, 4};
std::array<int64_t, 3> shape{5, 5, 4};
Ort::Value v =
Ort::Value::CreateTensor<float>(allocator, shape.data(), shape.size());
float *p = v.GetTensorMutableData<float>();
std::iota(p, p + shape[0] * shape[1] * shape[2], 0);
auto v1 = Slice(&v, 0, 2, 5);
auto v2 = Slice(&v, 1, 2, 4);
auto v1 = Slice(allocator, &v, 2, 4, 0, 2);
auto v2 = Slice(allocator, &v, 1, 3, 1, 3);
Print3D(&v);
Print2D(&v1);
Print2D(&v2);
Print3D(&v1);
Print3D(&v2);
// TODO(fangjun): Check that the results are correct
}

37
sherpa-onnx/csrc/slice.cc
View File

@@ -6,29 +6,48 @@
#include <assert.h>
#include <algorithm>
#include <vector>
namespace sherpa_onnx {
template <typename T /*=float*/>
Ort::Value Slice(const Ort::Value *v, int32_t dim0, int32_t dim1_start,
Ort::Value Slice(OrtAllocator *allocator, const Ort::Value *v,
int32_t dim0_start, int32_t dim0_end, int32_t dim1_start,
int32_t dim1_end) {
std::vector<int64_t> shape = v->GetTensorTypeAndShapeInfo().GetShape();
assert(shape.size() == 3);
auto memory_info =
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
assert(0 <= dim0_start);
assert(dim0_start < dim0_end);
assert(dim0_end <= shape[0]);
assert(0 <= dim1_start);
assert(dim1_start < dim1_end);
assert(dim1_end < shape[1]);
std::array<int64_t, 2> ans_shape{dim1_end - dim1_start, shape[2]};
const T *src = v->GetTensorData<T>();
src += dim0 * shape[1] * shape[2] + dim1_start * shape[2];
return Ort::Value::CreateTensor(memory_info, const_cast<T *>(src),
ans_shape[0] * ans_shape[1], ans_shape.data(),
ans_shape.size());
std::array<int64_t, 3> ans_shape{dim0_end - dim0_start, dim1_end - dim1_start,
shape[2]};
Ort::Value ans = Ort::Value::CreateTensor<T>(allocator, ans_shape.data(),
ans_shape.size());
T *dst = ans.GetTensorMutableData<T>();
for (int32_t i = dim0_start; i != dim0_end; ++i) {
const T *src = v->GetTensorData<T>() + i * shape[1] * shape[2];
const T *start = src + dim1_start * shape[2];
const T *end = src + dim1_end * shape[2];
std::copy(start, end, dst);
dst += ans_shape[1] * ans_shape[2];
}
return ans;
}
template Ort::Value Slice<float>(const Ort::Value *v, int32_t dim0,
template Ort::Value Slice<float>(OrtAllocator *allocator, const Ort::Value *v,
int32_t dim0_start, int32_t dim0_end,
int32_t dim1_start, int32_t dim1_end);
} // namespace sherpa_onnx

16
sherpa-onnx/csrc/slice.h
View File

@@ -8,21 +8,23 @@
namespace sherpa_onnx {
/** Get a shallow copy by slicing v.
/** Get a deep copy by slicing v.
*
* It returns v[dim0, dim1_start:dim1_end]
* It returns v[dim0_start:dim0_end, dim1_start:dim1_end]
*
* @param allocator
* @param v A 3-D tensor. Its data type is T.
* @param dim0 Start index of the first dimension..
* @param dim0_start Start index of the first dimension..
* @param dim0_end End index of the first dimension..
* @param dim1_start Start index of the second dimension.
* @param dim1_end End index of the second dimension.
*
* @return Return a 2-D tensor of shape (dim1_end-dim1_start, v.shape[2])
*
* @caution: The returned tensor is a shallow copy of `v`!
* @return Return a 3-D tensor of shape
* (dim0_end-dim0_start, dim1_end-dim1_start, v.shape[2])
*/
template <typename T = float>
Ort::Value Slice(const Ort::Value *v, int32_t dim0, int32_t dim1_start,
Ort::Value Slice(OrtAllocator *allocator, const Ort::Value *v,
int32_t dim0_start, int32_t dim0_end, int32_t dim1_start,
int32_t dim1_end);
} // namespace sherpa_onnx