def get_device_map_from(destinations): if isinstance(destinations, (value_lib.DistributedValues, value_lib.LogicalDeviceSpec)): return destinations.device_map, destinations.logical_device if isinstance(destinations, six.string_types): device = device_util.resolve(destinations) else: device = destinations.device return value_lib.SingleDeviceMap(device), 0
def __init__(self, container_strategy, device): super(OneDeviceExtended, self).__init__(container_strategy) self._device = device self._default_device = device self._input_device = device_util.canonicalize("/device:CPU:0") worker_device_pairs = [(self._input_device, [self._device])] device_map = values.SingleDeviceMap(device) self._input_workers = input_lib.InputWorkers( device_map, worker_device_pairs)
def testFunctionCanReturnPerReplica(self): f = def_function.function(lambda x: x) x = values.PerReplica(values.SingleDeviceMap("CPU"), (constant_op.constant(1.), )) y = f(x) self.assertIsNot(x, y) for a, b in zip(x._to_components(), y._to_components()): self.assertAllEqual(a, b) self.assertEqual(x._component_metadata(), y._component_metadata())
def testShapeInvariantToComponentsExplicitShape(self): v1 = constant_op.constant([1., 1., 1.]) v2 = constant_op.constant([2., 2., 2.]) per_replica = values.PerReplica(values.SingleDeviceMap("CPU"), (v1, v2)) shape = [None] self.assertEqual( per_replica._shape_invariant_to_components(shape=shape), (shape, shape))
def __init__(self, container_strategy, device): super(OneDeviceExtended, self).__init__(container_strategy) self._device = device_util.resolve(device) suffix_loc = self._device.rfind("/") self._input_device = self._device[:suffix_loc] + "/device:CPU:0" worker_device_pairs = [(self._input_device, [self._device])] device_map = values.SingleDeviceMap(device) self._input_workers = input_lib.InputWorkers(device_map, worker_device_pairs)
def testTypeSpec(self): device_map = values.SingleDeviceMap("CPU") vals = (constant_op.constant(1.),) per_replica = values.PerReplica(device_map, vals) spec = per_replica._type_spec self.assertEqual(spec._value_specs, (tensor_spec.TensorSpec([], dtypes.float32),)) self.assertEqual(spec._device_map, per_replica.device_map) self.assertEqual(spec._logical_device, per_replica.logical_device)
def testCondWithValuesNotConvertibleToTensor(self): device_map = values.SingleDeviceMap("CPU") per_replica_1 = values.PerReplica(device_map, (set(["a"]), )) per_replica_2 = values.PerReplica(device_map, (set(["b", "c"]), )) condition = array_ops.placeholder(dtypes.bool, []) with self.assertRaisesRegex(TypeError, "Could not build a TypeSpec for"): control_flow_ops.cond(condition, lambda: per_replica_1, lambda: per_replica_2)
def testFetchOnFrozenGraph(self): with context.graph_mode(): v = values.TPUMirroredVariable( strategy=None, device_map=values.SingleDeviceMap("/cpu:0"), values=[variables_lib.Variable(42.)], aggregation=None) self.evaluate(variables_lib.global_variables_initializer()) ops.get_default_graph().finalize() self.assertEqual(42., self.evaluate(v))
def testTypeSpecRoundTrip(self): device_map = values.SingleDeviceMap("CPU") vals = (constant_op.constant(1.),) per_replica = values.PerReplica(device_map, vals) spec = per_replica._type_spec tensor_list = spec._to_components(per_replica) reconstructed = spec._from_components(tensor_list) self.assertEqual(per_replica.device_map, reconstructed.device_map) self.assertEqual(per_replica.logical_device, reconstructed.logical_device) self.assertAllEqual(per_replica.values, reconstructed.values)
def testCondWithValuesConvertibleToTensor(self): device_map = values.SingleDeviceMap("CPU") per_replica_1 = values.PerReplica(device_map, ("a",)) per_replica_2 = values.PerReplica(device_map, ("b",)) condition = array_ops.placeholder_with_default(True, []) result = control_flow_ops.cond( condition, lambda: per_replica_1, lambda: per_replica_2) self.assertEqual(per_replica_1.device_map, result.device_map) self.assertEqual(per_replica_1.logical_device, result.logical_device) self.assertLen(result.values, 1) self.assertAllEqual(result.values[0], "a")
def _make_tensor_into_per_replica(input_tensor): """Converts a single tensor into a PerReplica object.""" if isinstance(input_tensor, (tuple, list)): raise ValueError("Cannot convert `input_tensor` to a `PerReplica` object, " "got %r but expected a object that is not a tuple or list." % (input_tensor,)) if isinstance(input_tensor, value_lib.PerReplica): return input_tensor try: device = input_tensor.device except AttributeError: raise ValueError("Cannot convert `input_tensor` to a `PerReplica` object " "because it doesn't have device set.") device_map = value_lib.SingleDeviceMap(device) return value_lib.PerReplica(device_map, (input_tensor,))
def test_supports_distributed_variables(self): device_map = distributed_values.SingleDeviceMap("/CPU:0") mirrored = distributed_values.MirroredVariable( None, device_map, [variables.Variable(1.)], variables.VariableAggregation.SUM) tpu = distributed_values.TPUMirroredVariable( strategy=None, device_map=device_map, values=[variables.Variable(42.)], aggregation=None) aggregating = distributed_values.AggregatingVariable( strategy=None, v=variables.Variable(1.), aggregation=None) m = module.Module() m.a = mirrored m.b = tpu m.c = aggregating self.assertEqual(m.variables, (mirrored, tpu, aggregating))
def testIsGraphTensor(self): per_replica = values.PerReplica(values.SingleDeviceMap("CPU"), (constant_op.constant(1.),)) self.assertEqual(per_replica._is_graph_tensor, not context.executing_eagerly())
def testShapeInvariantToComponents(self): v1 = constant_op.constant(1.) v2 = constant_op.constant(2.) per_replica = values.PerReplica(values.SingleDeviceMap("CPU"), (v1, v2)) self.assertEqual(per_replica._shape_invariant_to_components(), (v1.shape, v2.shape))