Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`. Args: export_dir: Directory containing the SavedModel file. Returns: A `SavedModel` protocol buffer. Raises: IOError: If the file does not exist, or cannot be successfully parsed.
(export_dir)
| 40 | |
| 41 | |
| 42 | def parse_saved_model(export_dir): |
| 43 | """Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`. |
| 44 | |
| 45 | Args: |
| 46 | export_dir: Directory containing the SavedModel file. |
| 47 | |
| 48 | Returns: |
| 49 | A `SavedModel` protocol buffer. |
| 50 | |
| 51 | Raises: |
| 52 | IOError: If the file does not exist, or cannot be successfully parsed. |
| 53 | """ |
| 54 | # Build the path to the SavedModel in pbtxt format. |
| 55 | path_to_pbtxt = os.path.join( |
| 56 | compat.as_bytes(export_dir), |
| 57 | compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT)) |
| 58 | # Build the path to the SavedModel in pb format. |
| 59 | path_to_pb = os.path.join( |
| 60 | compat.as_bytes(export_dir), |
| 61 | compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB)) |
| 62 | |
| 63 | # Parse the SavedModel protocol buffer. |
| 64 | saved_model = saved_model_pb2.SavedModel() |
| 65 | if file_io.file_exists(path_to_pb): |
| 66 | try: |
| 67 | file_content = file_io.FileIO(path_to_pb, "rb").read() |
| 68 | saved_model.ParseFromString(file_content) |
| 69 | return saved_model |
| 70 | except message.DecodeError as e: |
| 71 | raise IOError("Cannot parse file %s: %s." % (path_to_pb, str(e))) |
| 72 | elif file_io.file_exists(path_to_pbtxt): |
| 73 | try: |
| 74 | file_content = file_io.FileIO(path_to_pbtxt, "rb").read() |
| 75 | text_format.Merge(file_content.decode("utf-8"), saved_model) |
| 76 | return saved_model |
| 77 | except text_format.ParseError as e: |
| 78 | raise IOError("Cannot parse file %s: %s." % (path_to_pbtxt, str(e))) |
| 79 | else: |
| 80 | raise IOError("SavedModel file does not exist at: %s/{%s|%s}" % |
| 81 | (export_dir, |
| 82 | constants.SAVED_MODEL_FILENAME_PBTXT, |
| 83 | constants.SAVED_MODEL_FILENAME_PB)) |
| 84 | |
| 85 | |
| 86 | # TODO(b/120594573): Make this symbol also available as private, so that |