Fix regression introduced in cast transformer (#4658)

This commit is contained in:
Hariharan Seshadri 2020-07-29 23:32:01 -07:00 committed by GitHub
parent f90a2d46ae
commit 382f94c95c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 27 additions and 3 deletions

View file

@ -254,7 +254,10 @@ Status InsertCastTransformer::ApplyImpl(onnxruntime::Graph& graph, bool& modifie
ORT_ENFORCE(dtype_attribute->second.has_i(),
"InsertCastTransformer works on the assumption that `dtype` attribute holds an integer.");
dtype_attribute->second.set_i(TensorProto_DataType_FLOAT);
// Modify the dtype attribute (which defines the output type) to FLOAT if it is FLOAT16.
if (dtype_attribute->second.i() == TensorProto_DataType_FLOAT16) {
dtype_attribute->second.set_i(TensorProto_DataType_FLOAT);
}
}
}

View file

@ -134,8 +134,9 @@ TEST(TransformerTest, ThreeInARowRemoval) {
ASSERT_TRUE(op_to_count["Cast"] == 2);
}
// test a case where the ONNX inferred type (float16) is different from the type bound
// to the output NodeArg of the "RandomNormalLike" node because of the InsertCaseTransformer
// test a case where the ONNX inferred output type (float16) is different from the type bound
// to the output NodeArg of the "RandomNormalLike" node (input is float16) because of the InsertCaseTransformer
// Here the ONNX inferred output type (float16) must be made float because that is what the kernel produces
TEST(TransformerTest, RandomNormalLikeWithFloat16Inputs) {
auto model_uri = MODEL_FOLDER ORT_TSTR("random_normal_like_float16.onnx");
std::shared_ptr<Model> model;
@ -153,5 +154,25 @@ TEST(TransformerTest, RandomNormalLikeWithFloat16Inputs) {
EXPECT_TRUE(status.IsOK()) << status;
}
// A case where the ONNX inferred output type is int32 to a node that consumes float16 input
// Here the InsertCastTransformer must not change the ONNX inferred output type and keep it
// as is (int32)
TEST(TransformerTest, MultinomialWithFloat16Input) {
auto model_uri = MODEL_FOLDER ORT_TSTR("multinomial_float16.onnx");
std::shared_ptr<Model> model;
auto status = Model::Load(model_uri, model, nullptr, DefaultLoggingManager().DefaultLogger());
ASSERT_TRUE(status.IsOK()) << status;
Graph& graph = model->MainGraph();
InsertCastTransformer transformer("Test");
bool modified = false;
status = transformer.Apply(graph, modified, DefaultLoggingManager().DefaultLogger());
EXPECT_TRUE(status.IsOK()) << status;
EXPECT_TRUE(modified) << "Transformer should have added some Cast nodes";
status = graph.Resolve();
EXPECT_TRUE(status.IsOK()) << status;
}
} // namespace test
} // namespace onnxruntime

Binary file not shown.