packed-sequence.h
1.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
// 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;
// data is a 2-D tensor of shape (sum(batch_sizes), channels)
Ort::Value data{nullptr};
// Return a shallow copy of data[start:start+size, :]
Ort::Value Get(int32_t start, int32_t size) {
auto shape = data.GetTensorTypeAndShapeInfo().GetShape();
std::array<int64_t, 2> ans_shape{size, shape[1]};
float *p = data.GetTensorMutableData<float>();
auto memory_info =
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
// a shallow copy
return Ort::Value::CreateTensor(memory_info, p + start * shape[1],
size * shape[1], ans_shape.data(),
ans_shape.size());
}
};
/** 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_