Skip to content

Commit

Permalink
Removes tf.graph_util.set_cpu0() from the API.
Browse files Browse the repository at this point in the history
This function was only used internally.
Change: 140485370
  • Loading branch information
mrry authored and tensorflower-gardener committed Nov 29, 2016
1 parent df71c58 commit 3f8fc30
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 23 deletions.
2 changes: 0 additions & 2 deletions tensorflow/python/framework/graph_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
from tensorflow.python.framework.graph_util_impl import extract_sub_graph
from tensorflow.python.framework.graph_util_impl import must_run_on_cpu
from tensorflow.python.framework.graph_util_impl import remove_training_nodes
from tensorflow.python.framework.graph_util_impl import set_cpu0
from tensorflow.python.framework.graph_util_impl import tensor_shape_from_node_def_name
# pylint: enable=unused-import
from tensorflow.python.util.all_util import remove_undocumented
Expand All @@ -36,7 +35,6 @@
"convert_variables_to_constants",
"extract_sub_graph",
"must_run_on_cpu",
"set_cpu0",
"tensor_shape_from_node_def_name",
"remove_training_nodes",
]
Expand Down
18 changes: 0 additions & 18 deletions tensorflow/python/framework/graph_util_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
Expand All @@ -49,23 +48,6 @@ def _is_variable_op(op):
return op in _VARIABLE_OPS


def set_cpu0(device_string):
"""Creates a new device string based on `device_string' but using /CPU:0.
If the device is already on /CPU:0, this is a no-op.
Args:
device_string: A device string.
Returns:
A device string.
"""
parsed_device = pydev.DeviceSpec.from_string(device_string)
parsed_device.device_type = "CPU"
parsed_device.device_index = 0
return parsed_device.to_string()


def must_run_on_cpu(node, pin_variables_on_cpu=False):
"""Returns True if the given node_def must run on CPU, otherwise False.
Expand Down
21 changes: 18 additions & 3 deletions tensorflow/python/training/saver.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import errors
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.lib.io import file_io
Expand All @@ -62,6 +61,23 @@
"ResourceGather"])


def _set_cpu0(device_string):
"""Creates a new device string based on `device_string` but using /CPU:0.
If the device is already on /CPU:0, this is a no-op.
Args:
device_string: A device string.
Returns:
A device string.
"""
parsed_device = pydev.DeviceSpec.from_string(device_string)
parsed_device.device_type = "CPU"
parsed_device.device_index = 0
return parsed_device.to_string()


class BaseSaverBuilder(object):
"""Base class for Savers.
Expand Down Expand Up @@ -380,8 +396,7 @@ def _AddRestoreOps(self,
# available on the GPU.
# TODO(touts): Re-enable restore on GPU when we can support annotating
# string tensors as "HostMemory" inputs.
with ops.device(
graph_util.set_cpu0(saveable.device) if saveable.device else None):
with ops.device(_set_cpu0(saveable.device) if saveable.device else None):
with ops.control_dependencies(restore_control_inputs):
tensors = self.restore_op(filename_tensor, saveable, preferred_shard)
shapes = None
Expand Down

0 comments on commit 3f8fc30

Please sign in to comment.