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