Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`. Args: saved_model_dir: Directory containing the SavedModel file. Returns: A `SavedModel` protocol buffer. Raises: IOError: If the file does not exist, or cannot be successfully parsed.
(saved_model_dir)
| 29 | |
| 30 | |
| 31 | def read_saved_model(saved_model_dir): |
| 32 | """Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`. |
| 33 | |
| 34 | Args: |
| 35 | saved_model_dir: Directory containing the SavedModel file. |
| 36 | |
| 37 | Returns: |
| 38 | A `SavedModel` protocol buffer. |
| 39 | |
| 40 | Raises: |
| 41 | IOError: If the file does not exist, or cannot be successfully parsed. |
| 42 | """ |
| 43 | # Build the path to the SavedModel in pbtxt format. |
| 44 | path_to_pbtxt = os.path.join( |
| 45 | compat.as_bytes(saved_model_dir), |
| 46 | compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT)) |
| 47 | # Build the path to the SavedModel in pb format. |
| 48 | path_to_pb = os.path.join( |
| 49 | compat.as_bytes(saved_model_dir), |
| 50 | compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB)) |
| 51 | |
| 52 | # Ensure that the SavedModel exists at either path. |
| 53 | if not file_io.file_exists(path_to_pbtxt) and not file_io.file_exists( |
| 54 | path_to_pb): |
| 55 | raise IOError("SavedModel file does not exist at: %s" % saved_model_dir) |
| 56 | |
| 57 | # Parse the SavedModel protocol buffer. |
| 58 | saved_model = saved_model_pb2.SavedModel() |
| 59 | if file_io.file_exists(path_to_pb): |
| 60 | try: |
| 61 | file_content = file_io.FileIO(path_to_pb, "rb").read() |
| 62 | saved_model.ParseFromString(file_content) |
| 63 | return saved_model |
| 64 | except message.DecodeError as e: |
| 65 | raise IOError("Cannot parse file %s: %s." % (path_to_pb, str(e))) |
| 66 | elif file_io.file_exists(path_to_pbtxt): |
| 67 | try: |
| 68 | file_content = file_io.FileIO(path_to_pbtxt, "rb").read() |
| 69 | text_format.Merge(file_content.decode("utf-8"), saved_model) |
| 70 | return saved_model |
| 71 | except text_format.ParseError as e: |
| 72 | raise IOError("Cannot parse file %s: %s." % (path_to_pbtxt, str(e))) |
| 73 | else: |
| 74 | raise IOError("SavedModel file does not exist at: %s/{%s|%s}" % |
| 75 | (saved_model_dir, constants.SAVED_MODEL_FILENAME_PBTXT, |
| 76 | constants.SAVED_MODEL_FILENAME_PB)) |
| 77 | |
| 78 | |
| 79 | def get_saved_model_tag_sets(saved_model_dir): |