def testSingleOutfeed(self): feed_name = next_feed_id() shape = [10, 10] with ops.device("/device:IPU:0"): a = array_ops.placeholder(np.float32, shape) b = array_ops.placeholder(np.float32, shape) add = math_ops.add(a, b) outfeed_op = gen_pop_datastream_ops.pop_datastream_outfeed_enqueue( [add], feed_id=feed_name, replication_factor=1, io_batch_size=1, output_shapes=[shape]) with ops.device('cpu'): outfeed = gen_pop_datastream_ops.pop_datastream_outfeed_dequeue( feed_id=feed_name, replication_factor=1, output_types=[np.float32], output_shapes=[shape]) with session_lib.Session() as sess: sess.run(variables.global_variables_initializer()) sess.run(outfeed_op, feed_dict={ a: np.ones(shape, np.float32), b: np.ones(shape, np.float32) }) outfed = sess.run(outfeed) self.assertEqual(len(outfed[0]), 1) self.assertAllClose(outfed[0][0], 2 * np.ones(shape, np.float32))
def testTupleOutfeedGetLast(self): feed_name = next_feed_id() shape_1 = [10, 10] shape_2 = [4, 4] with ops.device("/device:IPU:0"): a = array_ops.placeholder(np.float32, shape_1) b = array_ops.placeholder(np.float32, shape_1) c = array_ops.placeholder(np.float32, shape_2) d = array_ops.placeholder(np.float32, shape_2) add = math_ops.add(a, b) sub = math_ops.sub(c, d) outfeed_op = gen_pop_datastream_ops.pop_datastream_outfeed_enqueue( [add, sub], feed_id=feed_name, replication_factor=1, outfeed_mode='get_last') with ops.device('cpu'): outfeed = gen_pop_datastream_ops.pop_datastream_outfeed_dequeue( feed_id=feed_name, replication_factor=1, output_types=[np.float32, np.float32], output_shapes=[shape_1, shape_2]) def get_result(sess, result): result.append(sess.run(outfeed)) with session_lib.Session() as sess: sess.run(variables.global_variables_initializer()) result = [] sess.run(outfeed_op, feed_dict={ a: np.ones(shape_1, np.float32), b: np.ones(shape_1, np.float32), c: np.ones(shape_2, np.float32), d: np.ones(shape_2, np.float32) }) sess.run(outfeed_op, feed_dict={ a: 2 * np.ones(shape_1, np.float32), b: np.ones(shape_1, np.float32), c: 2 * np.ones(shape_2, np.float32), d: np.ones(shape_2, np.float32) }) outfed = sess.run(outfeed) self.assertTrue(len(outfed) == 2) self.assertEqual(outfed[0].shape, (10, 10)) self.assertEqual(outfed[1].shape, (4, 4)) self.assertAllClose(outfed[0], np.broadcast_to(3, [10, 10])) self.assertAllClose(outfed[1], np.broadcast_to(1, [4, 4]))
def testOutfeedGetLast(self): feed_name = next_feed_id() shape = [2, 2] with ops.device("/device:IPU:0"): a = array_ops.placeholder(np.float32, shape) b = array_ops.placeholder(np.float32, shape) add = math_ops.add(a, b) outfeed_op = gen_pop_datastream_ops.pop_datastream_outfeed_enqueue( [add], feed_id=feed_name, replication_factor=1, io_batch_size=1, outfeed_mode='get_last', output_shapes=[shape]) with ops.device('cpu'): outfeed_last = gen_pop_datastream_ops.pop_datastream_outfeed_dequeue( feed_id=feed_name, replication_factor=1, outfeed_mode='get_last', output_types=[np.float32], output_shapes=[shape]) with session_lib.Session() as sess: sess.run(variables.global_variables_initializer()) sess.run(outfeed_op, feed_dict={ a: np.ones(shape, np.float32), b: np.ones(shape, np.float32) }) sess.run(outfeed_op, feed_dict={ a: 3.1 * np.ones(shape, np.float32), b: 2 * np.ones(shape, np.float32) }) outfed = sess.run(outfeed_last) self.assertTrue(len(outfed) == 1) self.assertEqual(outfed[0].shape, (2, 2)) self.assertAllClose(outfed[0], (3.1 + 2) * np.ones(shape, np.float32))
def dequeue(self): """Generate host side operation to dequeue the outfeed values. The operation generated by this function will block if called prior to any enqueues. The return value of this operation depends on the enqueued tensors, replication factor and the execution mode. Examples: 1. Outfeed returning a single tensor: .. code-block:: python outfeed_queue = ipu_outfeed_queue.IPUOutfeedQueue(feed_name="outfeed", replication_factor=2) def body(input): output = input + 1 outfeed = outfeed_queue.enqueue(output) return (output, outfeed) def my_net(input): r = loops.repeat(20, body, (input)) return r with ipu.scopes.ipu_scope("/device:IPU:0"): res = ipu_compiler.compile(my_net, inputs=[v]) with ops.device('cpu'): v = tf.placeholder(np.float32, [4, 4]) outfeed = outfeed_queue.dequeue() with tf.Session() as sess: result = sess.run(res, {v:np.ones([4, 4], np.float32)}) outfed = sess.run(outfeed) In this example the tensor `output` is of shape [4, 4] and it's enqueued into the outfeed with replication_factor = 2. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.ALL`, then the shape of the resulting `outfed` tensor will be [20, 2, 4, 4], where the first dimension represents the number of times we have enqueued a tensor to the outfeed - in this example the loop is repeated 20 times, and therefore we get 20 values back from the outfeed. The second dimension is the replication_factor, which allows us to see the individual values from each replicated graph. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.LAST`, then the shape of the resulting `outfed` tensor will be [2, 4, 4], which represents the value of the output tensor the last time it was enqueued during execution for each of the replicated graphs. 2. Outfeed returning a tuple of tensors: .. code-block:: python outfeed_queue = ipu_outfeed_queue.IPUOutfeedQueue(feed_name="outfeed") def body(input): output = input + 1 sum = tf.reduce_sum(output) outfeed = outfeed_queue.enqueue((output, sum)) return (output, outfeed) def my_net(input): r = loops.repeat(20, body, (input)) return r with ipu.scopes.ipu_scope("/device:IPU:0"): res = ipu_compiler.compile(my_net, inputs=[v]) with ops.device('cpu'): v = tf.placeholder(np.float32, [4, 4]) outfeed = outfeed_queue.dequeue() with tf.Session() as sess: result = sess.run(res, {v:np.ones([4, 4], np.float32)}) outfed = sess.run(outfeed) In this example we outfeed a tuple of tensors, `output` and `sum`, where the former is of shape [4, 4] and latter [1]. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.ALL`, then the resulting outfed is a two-tuple of tensors with shapes ([20, 4, 4], [20, 1]), where the first dimension in each of the tensors represents the number of times we have enqueued these tensors to the outfeed - in this example the loop is repeated 20 times, and therefore we get 20 values back from the outfeed for each of the tensors in the tuple. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.LAST`, then the `outfed` is a two tuple of tensors with shapes ([4, 4], [1]), which represents the values of the `output` and `sum` tensors the last time they were enqueued during execution. Note that `replication_factor` here is the default (=1), which means that the extra replication dimension is not added. 3. Outfeed returning a dictionary of tensors: .. code-block:: python outfeed_queue = ipu_outfeed_queue.IPUOutfeedQueue(feed_name="outfeed", replication_factor=8) def body(input): output = input + 1 sum = tf.reduce_sum(output) outfeed = outfeed_queue.enqueue({"x": output, "y": sum}) return (output, outfeed) def my_net(input): r = loops.repeat(40, body, (input)) return r with ipu.scopes.ipu_scope("/device:IPU:0"): res = ipu_compiler.compile(my_net, inputs=[v]) with ops.device('cpu'): v = tf.placeholder(np.float32, [4, 4]) outfeed = outfeed_queue.dequeue() with tf.Session() as sess: result = sess.run(res, {v:np.ones([4, 4], np.float32)}) outfed = sess.run(outfeed) In this example we outfeed a dictionary of tensors, `output` and `sum`, where the former is of shape [4, 4] and latter [1]. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.ALL`, then the resulting outfed is a dictionary of tensors with shapes: {"x": [40, 8, 4, 4], "y": [40, 8, 1]}, where the first dimension in each of the tensors represents the number of times we have enqueued these tensors to the outfeed - in this example the loop is repeated 40 times, and therefore we get 40 values back from the outfeed for each of the tensors in the tuple. The second dimension is the replication_factor, which allows us to see the individual values from each replicated graph. If the `outfeed_mode` is `outfeed_mode == IPUOutfeedMode.LAST`, then the `outfed` is a dictionary of tensors with shapes: {"x": [8, 4, 4], "y": [8, 1]}, which represents the values of the `output` and `sum` tensors the last time they were enqueued during execution for each of the replicated graphs. """ if not self.enqueued: raise ValueError( "Trying to dequeue an outfeed which has not been enqueued.") with ops.device('cpu'): outfeed_dequeue = \ gen_pop_datastream_ops.pop_datastream_outfeed_dequeue( output_types=self._structure.flat_types, output_shapes=self._structure.flat_shapes, outfeed_mode=self._outfeed_mode.value, feed_id=self._feed_name, device_ordinal=self._device_ordinal, replication_factor=self._replication_factor) return self._structure.from_tensor_list(outfeed_dequeue)