• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Types for specifying saving and loading behavior."""
16from __future__ import absolute_import
17from __future__ import division
18from __future__ import print_function
19
20
21class SaveSpec(object):
22  """Class used to describe tensor slices that need to be saved."""
23
24  def __init__(self, tensor, slice_spec, name, dtype=None, device=None):
25    """Creates a `SaveSpec` object.
26
27    Args:
28      tensor: the tensor to save or callable that produces a tensor to save.
29        If the value is `None`, the `SaveSpec` is ignored.
30      slice_spec: the slice to be saved. See `Variable.SaveSliceInfo`.
31      name: the name to save the tensor under.
32      dtype: The data type of the Tensor. Required if `tensor` is callable.
33        Used for error checking in the restore op.
34      device: The device generating and consuming this tensor. Required if
35        `tensor` is callable. Used to group objects to save by device.
36    """
37    self._tensor = tensor
38    self.slice_spec = slice_spec
39    self.name = name
40    if callable(self._tensor):
41      if dtype is None or device is None:
42        raise AssertionError(
43            "When passing a callable `tensor` to a SaveSpec, an explicit "
44            "dtype and device must be provided.")
45      self.dtype = dtype
46      self.device = device
47    else:
48      self.dtype = tensor.dtype
49      if device is not None:
50        self.device = device
51      else:
52        self.device = tensor.device
53
54  @property
55  def tensor(self):
56    return self._tensor() if callable(self._tensor) else self._tensor
57
58
59class SaveableObject(object):
60  """Base class for saving and restoring saveable objects."""
61
62  def __init__(self, op, specs, name):
63    """Creates a `SaveableObject` object.
64
65    Args:
66      op: the "producer" object that this class wraps; it produces a list of
67        tensors to save.  E.g., a "Variable" object saving its backing tensor.
68      specs: a list of SaveSpec, each element of which describes one tensor to
69        save under this object. All Tensors must be on the same device.
70      name: the name to save the object under.
71    """
72    self.op = op
73    self.specs = specs
74    self.name = name
75
76  @property
77  def optional_restore(self):
78    """A hint to restore assertions that this object is optional."""
79    return False  # Default to required
80
81  @property
82  def device(self):
83    """The device for SaveSpec Tensors."""
84    return self.specs[0].device
85
86  def restore(self, restored_tensors, restored_shapes):
87    """Restores this object from 'restored_tensors'.
88
89    Args:
90      restored_tensors: the tensors that were loaded from a checkpoint
91      restored_shapes: the shapes this object should conform to after
92        restore, or None.
93
94    Returns:
95      An operation that restores the state of the object.
96
97    Raises:
98      ValueError: If the object cannot be restored using the provided
99        parameters.
100    """
101    # pylint: disable=unused-argument
102    raise ValueError("Calling an abstract method.")
103