add interface to copy batch tensors. (#2807)

* add interface to copy batch tensors.

* onnxruntime
This commit is contained in:
Ke Zhang 2020-01-09 16:52:34 -08:00 committed by GitHub
parent 7ef6570e27
commit b0019ac7fe
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 16 additions and 15 deletions

View file

@ -253,7 +253,7 @@ struct OrtApi {
// TODO: document the path separator convention? '/' vs '\'
// TODO: should specify the access characteristics of model_path. Is this read only during the
// execution of OrtCreateSession, or does the OrtSession retain a handle to the file/directory
// execution of CreateSession, or does the OrtSession retain a handle to the file/directory
// and continue to access throughout the OrtSession lifetime?
// What sort of access is needed to model_path : read or read/write?
OrtStatus*(ORT_API_CALL* CreateSession)(_In_ const OrtEnv* env, _In_ const ORTCHAR_T* model_path,
@ -432,22 +432,22 @@ struct OrtApi {
OrtStatus*(ORT_API_CALL* GetTensorMutableData)(_Inout_ OrtValue* value, _Outptr_ void** out)NO_EXCEPTION;
/**
* \param value A tensor created from OrtCreateTensor... function.
* \param value A tensor created from CreateTensor... function.
* \param s each A string array. Each string in this array must be null terminated.
* \param s_len length of s
*/
OrtStatus*(ORT_API_CALL* FillStringTensor)(_Inout_ OrtValue* value, _In_ const char* const* s, size_t s_len)NO_EXCEPTION;
/**
* \param value A tensor created from OrtCreateTensor... function.
* \param value A tensor created from CreateTensor... function.
* \param len total data length, not including the trailing '\0' chars.
*/
OrtStatus*(ORT_API_CALL* GetStringTensorDataLength)(_In_ const OrtValue* value, _Out_ size_t* len)NO_EXCEPTION;
/**
* \param s string contents. Each string is NOT null-terminated.
* \param value A tensor created from OrtCreateTensor... function.
* \param s_len total data length, get it from OrtGetStringTensorDataLength
* \param value A tensor created from CreateTensor... function.
* \param s_len total data length, get it from GetStringTensorDataLength
*/
OrtStatus*(ORT_API_CALL* GetStringTensorContent)(_In_ const OrtValue* value, _Out_ void* s, size_t s_len,
_Out_ size_t* offsets, size_t offsets_len)NO_EXCEPTION;
@ -470,7 +470,7 @@ struct OrtApi {
OrtStatus*(ORT_API_CALL* SetTensorElementType)(_Inout_ OrtTensorTypeAndShapeInfo*, enum ONNXTensorElementDataType type)NO_EXCEPTION;
/**
* \param info Created from OrtCreateTensorTypeAndShapeInfo() function
* \param info Created from CreateTensorTypeAndShapeInfo() function
* \param dim_values An array with length of `dim_count`. Its elements can contain negative values.
* \param dim_count length of dim_values
*/

View file

@ -9,6 +9,14 @@ common::Status IDataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const {
return CopyTensor(src, dst, 0);
}
common::Status IDataTransfer::CopyTensors(const Tensor* src, Tensor* dst, int size) const {
ORT_ENFORCE(nullptr != src && nullptr != dst);
for (int i = 0; i < size; ++i) {
ORT_RETURN_IF_ERROR(CopyTensor(src[i], dst[i], 0));
}
return Status::OK();
}
bool CPUDataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const {
return src_device.Type() == OrtDevice::CPU && dst_device.Type() == OrtDevice::CPU;
}

View file

@ -17,6 +17,7 @@ class IDataTransfer {
virtual common::Status CopyTensor(const Tensor& src, Tensor& dst) const;
virtual common::Status CopyTensor(const Tensor& src, Tensor& dst, int exec_queue_id) const = 0;
virtual common::Status CopyTensors(const Tensor* src, Tensor* dst, int size) const;
};
class CPUDataTransfer : public IDataTransfer {

View file

@ -807,16 +807,8 @@ ORT_API_STATUS_IMPL(OrtApis::AllocatorGetInfo, _In_ const OrtAllocator* ptr, _Ou
API_IMPL_END
}
///////////////////////////////////////////////////////////////////////////
// Code to handle non-tensor types
// OrtGetValueCount
// OrtGetVaue
// OrtCreateValue
///////////////////////////////////////////////////////////////////////////
const int NUM_MAP_INDICES = 2;
////////////////////
// OrtGetValueCount
template <typename T>
OrtStatus* OrtGetNumSequenceElements(const OrtValue* p_ml_value, size_t* out) {
auto& data = p_ml_value->Get<T>();
@ -867,7 +859,7 @@ ORT_API_STATUS_IMPL(OrtApis::GetValueCount, const OrtValue* value, size_t* out)
}
///////////////////
// OrtGetValue
// OrtGetValueImplSeqOfMap
template <typename T>
static OrtStatus* OrtGetValueImplSeqOfMap(const OrtValue* p_ml_value, int index, OrtValue** out) {
using TKey = typename T::value_type::key_type;