gemm with empty input (#350)

* gemm with empty input

* check negative

* align the comments
This commit is contained in:
Tang, Cheng 2019-01-24 18:24:32 -08:00 committed by Ke Zhang
parent f922d427d3
commit 373177ddd3
3 changed files with 22 additions and 1 deletions

View file

@ -42,6 +42,9 @@ class Gemm : public OpKernel {
int64_t N = helper.N();
int64_t K = helper.K();
auto Y = context->Output(0, TensorShape({M, N}));
// if input is emtpy tensor, return directly as nothing need to be calculated.
if (M == 0 || N == 0)
return Status::OK();
T_Y* y_data = Y->template MutableData<T_Y>();
//bias

View file

@ -40,7 +40,8 @@ class GemmHelper {
if (!IsValidBroadcast(bias, M_, N_))
status_ = common::Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Gemm: Invalid bias shape for broadcast");
ORT_ENFORCE(M_ > 0 && N_ > 0 && K_ > 0);
// it is possible the input is empty tensor, for example the output of roipool in fast rcnn.
ORT_ENFORCE(M_ >= 0 && K_ > 0 && N_ >= 0);
}
int64_t M() const { return M_; }

View file

@ -161,5 +161,22 @@ TEST(MathOpTest, GemmFalseBroadcast) {
test.Run();
}
TEST(MathOpTest, GemmEmptyTensor) {
OpTester test("Gemm");
test.AddAttribute("transA", static_cast<int64_t>(0));
test.AddAttribute("transB", static_cast<int64_t>(0));
test.AddAttribute("alpha", 1.0f);
test.AddAttribute("beta", 1.0f);
test.AddInput<float>("A", {0, 4},
{});
test.AddInput<float>("B", {4, 3}, std::vector<float>(12, 1.0f));
test.AddInput<float>("C", {3}, std::vector<float>(3, 1.0f));
test.AddOutput<float>("Y", {0, 3},
{});
test.Run();
}
} // namespace test
} // namespace onnxruntime