See `init_from_checkpoint` for documentation.
(ckpt_dir_or_file, assignment_map, reset_version=False)
| 295 | |
| 296 | |
| 297 | def _init_from_checkpoint(ckpt_dir_or_file, assignment_map, reset_version=False): |
| 298 | """See `init_from_checkpoint` for documentation.""" |
| 299 | ckpt_file = _get_checkpoint_filename(ckpt_dir_or_file) |
| 300 | reader = load_checkpoint(ckpt_dir_or_file) |
| 301 | variable_map = reader.get_variable_to_shape_map() |
| 302 | for tensor_name_in_ckpt, current_var_or_name in sorted( |
| 303 | six.iteritems(assignment_map)): |
| 304 | var = None |
| 305 | # Check if this is Variable object or list of Variable objects (in case of |
| 306 | # partitioned variables). |
| 307 | if _is_variable(current_var_or_name) or ( |
| 308 | isinstance(current_var_or_name, list) |
| 309 | and all(_is_variable(v) for v in current_var_or_name)): |
| 310 | var = current_var_or_name |
| 311 | else: |
| 312 | store_vars = vs._get_default_variable_store()._vars # pylint:disable=protected-access |
| 313 | # Check if this variable is in var_store. |
| 314 | var = store_vars.get(current_var_or_name, None) |
| 315 | # Also check if variable is partitioned as list. |
| 316 | if var is None: |
| 317 | var = _collect_partitioned_variable(current_var_or_name, store_vars) |
| 318 | if var is not None: |
| 319 | # If 1 to 1 mapping was provided, find variable in the checkpoint. |
| 320 | if tensor_name_in_ckpt not in variable_map: |
| 321 | raise ValueError("Tensor %s is not found in %s checkpoint %s" % ( |
| 322 | tensor_name_in_ckpt, ckpt_dir_or_file, variable_map |
| 323 | )) |
| 324 | if _is_variable(var): |
| 325 | # Additional at-call-time checks. |
| 326 | if not var.get_shape().is_compatible_with( |
| 327 | variable_map[tensor_name_in_ckpt]): |
| 328 | raise ValueError( |
| 329 | "Shape of variable %s (%s) doesn't match with shape of " |
| 330 | "tensor %s (%s) from checkpoint reader." % ( |
| 331 | var.name, str(var.get_shape()), |
| 332 | tensor_name_in_ckpt, str(variable_map[tensor_name_in_ckpt]) |
| 333 | )) |
| 334 | var_name = var.name |
| 335 | else: |
| 336 | var_name = ",".join([v.name for v in var]) |
| 337 | _set_variable_or_list_initializer(var, ckpt_file, tensor_name_in_ckpt) |
| 338 | logging.debug("Initialize variable %s from checkpoint %s with %s", |
| 339 | var_name, ckpt_dir_or_file, tensor_name_in_ckpt) |
| 340 | else: |
| 341 | scopes = "" |
| 342 | # TODO(vihanjain): Support list of 'current_var_or_name' here. |
| 343 | if "/" in current_var_or_name: |
| 344 | scopes = current_var_or_name[:current_var_or_name.rindex("/")] |
| 345 | if not tensor_name_in_ckpt.endswith("/"): |
| 346 | raise ValueError( |
| 347 | "Assignment map with scope only name {} should map to scope only " |
| 348 | "{}. Should be 'scope/': 'other_scope/'.".format( |
| 349 | scopes, tensor_name_in_ckpt)) |
| 350 | # If scope to scope mapping was provided, find all variables in the scope |
| 351 | # and create variable to variable mapping. |
| 352 | scope_variables = set() |
| 353 | for var_name in store_vars: |
| 354 | var = store_vars.get(var_name, None) |
no test coverage detected