saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make output 1 of ConcatTraining Optional and place on CPU (#7199)

Browse source 
* Optional input 1 on CPU ConcatTraining

* Rename output_1
This commit is contained in:
ashbhandare 2021-04-01 16:05:17 -07:00 committed by GitHub
parent 4543459984
commit 15c67ddbf0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 37 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -1332,7 +1332,8 @@ Example 4:
.Output(1, "per_input_length",
"Vector of length of each concatenated "
"input along the 'axis' dimension",
"Tint")
"Tint",
OpSchema::Optional)
.TypeConstraint(
"T",
OpSchema::all_tensor_types(),
@ -1368,9 +1369,11 @@ Example 4:
output_shape->add_dim();
}
ONNX_NAMESPACE::TensorShapeProto per_input_len_shape;
per_input_len_shape.add_dim()->set_dim_value(numInputs);
updateOutputShape(ctx, 1, per_input_len_shape);
if (ctx.getNumOutputs() > 1) {
ONNX_NAMESPACE::TensorShapeProto per_input_len_shape;
per_input_len_shape.add_dim()->set_dim_value(numInputs);
updateOutputShape(ctx, 1, per_input_len_shape);
}
for (size_t i = 0; i < numInputs; i++) {
const auto& shape = ctx.getInputType(i)->tensor_type().shape();

17
orttraining/orttraining/test/training_ops/cpu/tensor/concat_op_test.cc
View file

@ -67,5 +67,22 @@ TEST(ConcatTrainingOpTest, Concat3D_same_len) {
test.Run();
}
TEST(ConcatTrainingOpTest, Concat2D_optional_output1) {
OpTester test("ConcatTraining", 1, kMSDomain);
test.AddAttribute("axis", int64_t{1});
std::vector<int64_t> dims{4, 1};
test.AddInput<float>("input1", dims, {11.0f, 21.0f, 31.0f, 41.0f});
test.AddInput<float>("input2", {4, 2}, {12.0f, 13.0f, 22.0f, 23.0f, 32.0f, 33.0f, 42.0f, 43.0f});
test.AddInput<float>("input3", dims, {14.0f, 24.0f, 34.0f, 44.0f});
test.AddOutput<float>("concat_result", {4, 4},
{11.0f, 12.0f, 13.0f, 14.0f,
21.0f, 22.0f, 23.0f, 24.0f,
31.0f, 32.0f, 33.0f, 34.0f,
41.0f, 42.0f, 43.0f, 44.0f});
test.AddMissingOptionalOutput<int64_t>();
test.Run();
}
} // namespace test
} // namespace onnxruntime

15
orttraining/orttraining/training_ops/cpu/tensor/concat.cc
View file

@ -39,15 +39,14 @@ Status ConcatTraining::Compute(OpKernelContext* ctx) const {
if (p.output_num_elements == 0)
return Status::OK();
// Create output tensor for 'per_input_length'
std::vector<int64_t> per_input_length(input_count);
for (int i = 0; i < input_count; ++i) {
per_input_length[i] = input_tensors[i]->Shape()[p.axis];
// Create optional output tensor for 'per_input_length'
Tensor* per_input_length_tensor = ctx->Output(1, {input_count});
if (per_input_length_tensor) {
int64_t* per_input_length = per_input_length_tensor->template MutableData<int64_t>();
for (int i = 0; i < input_count; ++i) {
per_input_length[i] = input_tensors[i]->Shape()[p.axis];
}
}
Tensor* output_1_tensor = ctx->Output(1, {input_count});
int64_t* output_1_tensor_data = output_1_tensor->template MutableData<int64_t>();
std::copy(per_input_length.begin(), per_input_length.end(), output_1_tensor_data);
// Compute values to be placed in the output tensor
return ComputeImpl(p);
}

8
orttraining/orttraining/training_ops/cuda/tensor/concat.cc
View file

@ -11,6 +11,7 @@ ONNX_OPERATOR_KERNEL_EX(ConcatTraining,
1,
kCudaExecutionProvider,
KernelDefBuilder()
.OutputMemoryType<OrtMemTypeCPUInput>(1)
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()),
ConcatTraining);
@ -71,8 +72,11 @@ Status ConcatTraining::ComputeInternal(OpKernelContext* ctx) const {
input_ptr.GpuPtr(),
p.output_num_elements));
Tensor* output_1_tensor = ctx->Output(1, {input_count});
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_1_tensor->template MutableData<int64_t>(), concat_sizes_gpu.GpuPtr(), input_count * sizeof(int64_t), cudaMemcpyDeviceToDevice, Stream()));
// Create optional output tensor for 'per_input_length'
Tensor* per_input_length_tensor = ctx->Output(1, {input_count});
if (per_input_length_tensor) {
std::copy(concat_sizes.begin(), concat_sizes.end(), per_input_length_tensor->template MutableData<int64_t>());
}
return Status::OK();
}