def benchmark_create_1000_partitions_with_100_parameter_servers(self): workers, _ = test.create_local_cluster(num_workers=1, num_ps=100) worker_sessions = [session_lib.Session(w.target) for w in workers] worker = worker_sessions[0] partition_sizes = (1, 512, 1024 * 32, 1024 * 128) partitioned = [] for partition_size in partition_sizes: # max_shard_bytes is 4, shape is 1000*partition_size float32s which should # partition into 1000 shards, each containing partition_size float32s. print("Building partitioned variable with %d floats per partition" % partition_size) with ops.device(device_setter.replica_device_setter(ps_tasks=100)): partitioned_ix = variable_scope.get_variable( "partitioned_%d" % partition_size, shape=[1000 * partition_size], dtype=dtypes.float32, # Each partition to have exactly N float32s partitioner=partitioned_variables.variable_axis_size_partitioner( max_shard_bytes=4 * partition_size)) # Concatenates along axis 0 partitioned.append(ops.convert_to_tensor(partitioned_ix)) variables.global_variables_initializer().run(session=worker) for ix, partition_size in enumerate(partition_sizes): print("Running benchmark having partitions with %d floats" % partition_size) self.run_op_benchmark( worker, partitioned[ix], name=("read_concat_1000_partitions_from_" "100_parameter_servers_partsize_%d_floats" % partition_size))
def testTensorArrayGetsDeviceFromFirstWriteInWhileLoop(self): with ops.device("/job:worker/task:0/cpu:0"): ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2) def _body(i, ta_i): with ops.device("/job:worker/task:1/cpu:0"): return i + 1, ta_i.write(i, 0.0) _, ta_out = control_flow_ops.while_loop( lambda i, ta: i < 2, _body, loop_vars=[0, ta]) workers, _ = test.create_local_cluster(num_workers=3, num_ps=0) session = session_lib.Session(workers[0].target) run_options = config_pb2.RunOptions( trace_level=config_pb2.RunOptions.FULL_TRACE) run_metadata = config_pb2.RunMetadata() session.run(ta_out.flow, options=run_options, run_metadata=run_metadata) self.assertTrue(run_metadata.HasField("step_stats")) dev_stats = {d.device: d.node_stats for d in run_metadata.step_stats.dev_stats} for d in dev_stats: if "/task:1/" in d: self.assertTrue( [s for s in dev_stats[d] if "/TensorArray" in s.node_name]) else: self.assertFalse( [s for s in dev_stats[d] if "/TensorArray" in s.node_name])
def benchmarkCreateLocalCluster(self): deltas = [] iters = 5 for _ in range(iters): start_time = time.time() test.create_local_cluster(num_workers=1, num_ps=10) end_time = time.time() deltas.append(end_time - start_time) median_deltas = np.median(deltas) print("\n\nbenchmark_create_local_cluster_1_worker_10_ps. " "iterations: %d, median wall time: %g\n\n" % (iters, median_deltas)) self.report_benchmark( iters=iters, wall_time=median_deltas, name="benchmark_create_local_cluster_1_worker_10_ps")
def benchmarkCreateLocalCluster(self): deltas = [] iters = 5 for _ in range(iters): start_time = time.time() test.create_local_cluster(num_workers=1, num_ps=10) end_time = time.time() deltas.append(end_time - start_time) median_deltas = np.median(deltas) print("\n\nbenchmark_create_local_cluster_1_worker_10_ps. " "iterations: %d, median wall time: %g\n\n" % (iters, median_deltas)) self.report_benchmark( iters=iters, wall_time=median_deltas, name="benchmark_create_local_cluster_1_worker_10_ps")
def testBasics(self): (worker,), _ = test.create_local_cluster(num_workers=1, num_ps=0) self.assertTrue(worker.target.startswith("grpc://")) tf_context = xrt.get_tf_context(worker.target[len("grpc://"):], "worker") backend = xrt.XrtBackend(tf_context, "XLA_CPU") a = np.arange(10) b = np.arange(10) c = BuildAddAndScaleComputation( xla_client.Shape.from_pyval(a), xla_client.Shape.from_pyval(b)) executable = c.Compile(backend=backend) output = executable.ExecuteWithPythonValues((a, b)) self.assertAllEqual(output, (a + b) * 3)
def testTensorArrayGetsDeviceFromFirstWrite(self): with ops.device("/job:worker/task:0/cpu:0"): # this initial device will be ignored. ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2) with ops.device("/job:worker/task:1/cpu:0"): # the first write sets the op's device. ta = ta.write(0, 1.0) with ops.device("/job:worker/task:2/cpu:0"): # subsequent writes do not modify the op's device. ta = ta.write(1, 1.0) # The gradient TA will sit on the same device as the forward TA. ta_grad = ta.grad("grad") flows = [ta.flow, ta_grad.flow] # Similar tests for unpack and split with ops.device("/job:worker/task:0/cpu:0"): ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=3) with ops.device("/job:worker/task:1/cpu:0"): ta = ta.unstack([1.0, 2.0]) with ops.device("/job:worker/task:2/cpu:0"): ta = ta.write(2, 3.0) flows.append(ta.flow) with ops.device("/job:worker/task:0/cpu:0"): ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2) with ops.device("/job:worker/task:1/cpu:0"): ta = ta.split([1.0, 2.0], [1, 1]) flows.append(ta.flow) workers, _ = test.create_local_cluster(num_workers=3, num_ps=0) session = session_lib.Session(workers[0].target) run_options = config_pb2.RunOptions( trace_level=config_pb2.RunOptions.FULL_TRACE) run_metadata = config_pb2.RunMetadata() session.run(flows, options=run_options, run_metadata=run_metadata) self.assertTrue(run_metadata.HasField("step_stats")) dev_stats = {d.device: d.node_stats for d in run_metadata.step_stats.dev_stats} for d in dev_stats: if "/task:1/" in d: self.assertTrue( [s for s in dev_stats[d] if "/TensorArray" in s.node_name]) else: self.assertFalse( [s for s in dev_stats[d] if "/TensorArray" in s.node_name])
def testTuples(self): (worker,), _ = test.create_local_cluster(num_workers=1, num_ps=0) self.assertTrue(worker.target.startswith("grpc://")) tf_context = xrt.get_tf_context(worker.target[len("grpc://"):], "worker") backend = xrt.XrtBackend(tf_context, "XLA_CPU") a = np.random.randn(10) b = np.random.randn(15, 3) pieces = [ xla_client.Buffer.from_pyval(a, backend=backend), xla_client.Buffer.from_pyval(b, backend=backend) ] t = xla_client.Buffer.make_tuple(pieces, backend=backend) a_out, b_out = t.destructure() self.assertAllEqual(a, a_out.to_py()) self.assertAllEqual(b, b_out.to_py())
def testBasics(self): (worker,), _ = test.create_local_cluster(num_workers=1, num_ps=0) self.assertTrue(worker.target.startswith("grpc://")) tf_context = xrt.get_tf_context(worker.target[len("grpc://"):], "worker") backend = xrt.XrtBackend(tf_context, "XLA_CPU") a = np.arange(10) b = np.arange(10) c = BuildAddAndScaleComputation( xla_client.shape_from_pyval(a), xla_client.shape_from_pyval(b)) executable = c.Compile(backend=backend) output = xla_client.execute_with_python_values( executable, (a, b), backend=backend) self.assertAllEqual(output, (a + b) * 3)
def testTuples(self): (worker, ), _ = test.create_local_cluster(num_workers=1, num_ps=0) self.assertTrue(worker.target.startswith("grpc://")) tf_context = xrt.get_tf_context(worker.target[len("grpc://"):], "worker") backend = xrt.XrtBackend(tf_context, "XLA_CPU") a = np.random.randn(10) b = np.random.randn(15, 3) pieces = [ xla_client.Buffer.from_pyval(a, backend=backend), xla_client.Buffer.from_pyval(b, backend=backend) ] t = xla_client.Buffer.make_tuple(pieces, backend=backend) a_out, b_out = t.destructure() self.assertAllEqual(a, a_out.to_py()) self.assertAllEqual(b, b_out.to_py())
def testCreateLocalCluster(self): workers, _ = test.create_local_cluster(num_workers=2, num_ps=2) worker_sessions = [session_lib.Session(w.target) for w in workers] with ops.device("/job:ps/task:0"): var0 = variables.Variable(0.0) with ops.device("/job:ps/task:1"): var1 = variables.Variable(1.0) worker_sessions[0].run([var0.initializer, var1.initializer]) with ops.device("/job:ps/task:0"): var2 = variables.Variable(2.0) with ops.device("/job:ps/task:1"): var3 = variables.Variable(3.0) worker_sessions[1].run([var2.initializer, var3.initializer]) # Read values back in the opposite session self.assertAllEqual(0.0, var0.eval(session=worker_sessions[1])) self.assertAllEqual(1.0, var1.eval(session=worker_sessions[1])) self.assertAllEqual(2.0, var2.eval(session=worker_sessions[0])) self.assertAllEqual(3.0, var3.eval(session=worker_sessions[0]))
def testCreateLocalCluster(self): workers, _ = test.create_local_cluster(num_workers=2, num_ps=2) worker_sessions = [session_lib.Session(w.target) for w in workers] with ops.device("/job:ps/task:0"): var0 = variables.Variable(0.0) with ops.device("/job:ps/task:1"): var1 = variables.Variable(1.0) worker_sessions[0].run([var0.initializer, var1.initializer]) with ops.device("/job:ps/task:0"): var2 = variables.Variable(2.0) with ops.device("/job:ps/task:1"): var3 = variables.Variable(3.0) worker_sessions[1].run([var2.initializer, var3.initializer]) # Read values back in the opposite session self.assertAllEqual(0.0, var0.eval(session=worker_sessions[1])) self.assertAllEqual(1.0, var1.eval(session=worker_sessions[1])) self.assertAllEqual(2.0, var2.eval(session=worker_sessions[0])) self.assertAllEqual(3.0, var3.eval(session=worker_sessions[0]))
def benchmark_create_1000_partitions_with_100_parameter_servers(self): workers, _ = test.create_local_cluster(num_workers=1, num_ps=100) worker_sessions = [session_lib.Session(w.target) for w in workers] worker = worker_sessions[0] partition_sizes = (1, 512, 1024 * 32, 1024 * 128) partitioned = [] for partition_size in partition_sizes: # max_shard_bytes is 4, shape is 1000*partition_size float32s which should # partition into 1000 shards, each containing partition_size float32s. print( "Building partitioned variable with %d floats per partition" % partition_size) with ops.device(device_setter.replica_device_setter(ps_tasks=100)): partitioned_ix = variable_scope.get_variable( "partitioned_%d" % partition_size, shape=[1000 * partition_size], dtype=dtypes.float32, # Each partition to have exactly N float32s partitioner=partitioned_variables. variable_axis_size_partitioner(max_shard_bytes=4 * partition_size)) # Concatenates along axis 0 partitioned.append(ops.convert_to_tensor(partitioned_ix)) variables.global_variables_initializer().run(session=worker) for ix, partition_size in enumerate(partition_sizes): print("Running benchmark having partitions with %d floats" % partition_size) self.run_op_benchmark( worker, partitioned[ix], name=("read_concat_1000_partitions_from_" "100_parameter_servers_partsize_%d_floats" % partition_size))
def setUpClass(cls): super(TensorArrayTest, cls).setUpClass() cls._workers, _ = test.create_local_cluster(num_workers=3, num_ps=0)
def setUpClass(cls): super(TensorArrayTest, cls).setUpClass() cls._workers, _ = test.create_local_cluster(num_workers=3, num_ps=0)