Build a graph containing a sequence of split operations. Args: device: string, the device to run on. input_shape: shape of the input tensor. output_sizes: size of each output along axis. axis: axis to be split along. Returns: An array of tensors to run()
(device, input_shape, output_sizes, axis)
| 32 | |
| 33 | |
| 34 | def build_graph(device, input_shape, output_sizes, axis): |
| 35 | """Build a graph containing a sequence of split operations. |
| 36 | |
| 37 | Args: |
| 38 | device: string, the device to run on. |
| 39 | input_shape: shape of the input tensor. |
| 40 | output_sizes: size of each output along axis. |
| 41 | axis: axis to be split along. |
| 42 | |
| 43 | Returns: |
| 44 | An array of tensors to run() |
| 45 | """ |
| 46 | with ops.device("/%s:0" % device): |
| 47 | inp = array_ops.zeros(input_shape) |
| 48 | |
| 49 | outputs = [] |
| 50 | for _ in range(100): |
| 51 | outputs.extend(array_ops.split(inp, output_sizes, axis)) |
| 52 | return control_flow_ops.group(*outputs) |
| 53 | |
| 54 | |
| 55 | class SplitBenchmark(test.Benchmark): |