experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add pruning steps to saved model importer paths

Browse Source 
PiperOrigin-RevId: 305147277
Change-Id: I0daaa9735063484f3064b77e392e555456453bc3
This commit is contained in:
Jaesung Chung 2020-04-06 16:51:05 -07:00 committed by TensorFlower Gardener
parent 47816940ea
commit b91970316f
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
View File

@ -2863,6 +2863,7 @@ StatusOr<mlir::OwningModuleRef> SavedModelObjectGraphImporter::Convert(
saved_model->debug_info() ? *saved_model->debug_info() : dummy_debug_info;
GraphImportConfig specs;
specs.prune_unused_nodes = true;
mlir::OwningModuleRef module =
mlir::ModuleOp::create(mlir::UnknownLoc::get(context));
std::unordered_map<std::string, std::string> tf_name_to_mlir_name;
@ -2888,6 +2889,8 @@ StatusOr<mlir::OwningModuleRef> SavedModelObjectGraphImporter::Convert(
module.get(), &tf_name_to_mlir_name,
&function_name_uniquifier);
TF_RETURN_IF_ERROR(importer.PrepareConvert(graph));
auto fn_names = graph.flib_def().ListFunctionNames();
for (const auto& fn_name : fn_names) {
TF_RETURN_IF_ERROR(importer.ConvertLibFunction(fn_name));
@ -3026,6 +3029,7 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
});
GraphImportConfig specs;
specs.prune_unused_nodes = true;
specs.inputs = ParseInputArrays(inputs);
for (auto& output : outputs) specs.outputs.push_back(output.second.name());