Work around TF problems in checkpoint path handling. Args: path: a user-input path Returns: str: the argument that can be passed to `tf.train.NewCheckpointReader`
(path)
| 189 | |
| 190 | |
| 191 | def get_checkpoint_path(path): |
| 192 | """ |
| 193 | Work around TF problems in checkpoint path handling. |
| 194 | |
| 195 | Args: |
| 196 | path: a user-input path |
| 197 | Returns: |
| 198 | str: the argument that can be passed to `tf.train.NewCheckpointReader` |
| 199 | """ |
| 200 | if os.path.basename(path) == path: |
| 201 | path = os.path.join('.', path) # avoid #4921 and #6142 |
| 202 | if os.path.basename(path) == 'checkpoint': |
| 203 | assert tfv1.gfile.Exists(path), path |
| 204 | path = tfv1.train.latest_checkpoint(os.path.dirname(path)) |
| 205 | # to be consistent with either v1 or v2 |
| 206 | |
| 207 | # fix paths if provided a wrong one |
| 208 | new_path = path |
| 209 | if '00000-of-00001' in path: |
| 210 | new_path = path.split('.data')[0] |
| 211 | elif path.endswith('.index'): |
| 212 | new_path = path.split('.index')[0] |
| 213 | if new_path != path: |
| 214 | logger.info( |
| 215 | "Checkpoint path {} is auto-corrected to {}.".format(path, new_path)) |
| 216 | path = new_path |
| 217 | assert tfv1.gfile.Exists(path) or tfv1.gfile.Exists(path + '.index'), path |
| 218 | return path |
| 219 | |
| 220 | |
| 221 | def get_all_checkpoints(dir: str, prefix: str = "model"): |
no test coverage detected