| 49 | return PrefetchRunnerHook() |
| 50 | |
| 51 | def fill_prefetch_runner_options(options, |
| 52 | fetch_tensors, |
| 53 | cancel_fetching, |
| 54 | resume_fetching, |
| 55 | close_fetching, |
| 56 | feed_dict={}, |
| 57 | closed_exception_types=(errors.OUT_OF_RANGE,), |
| 58 | ignored_exception_types=(), |
| 59 | use_stage_subgraph_thread_pool=False, |
| 60 | stage_subgraph_thread_pool_id=0): |
| 61 | def _feed_fn(feed, feed_val): |
| 62 | for tensor_type, _, feed_fn, _ in _REGISTERED_EXPANSIONS: |
| 63 | if isinstance(feed, tensor_type): |
| 64 | return feed_fn(feed, feed_val) |
| 65 | raise TypeError('Feed argument %r has invalid type %r' % (feed, type(feed))) |
| 66 | |
| 67 | options.run_options.use_stage_subgraph_thread_pool = \ |
| 68 | use_stage_subgraph_thread_pool |
| 69 | options.run_options.stage_subgraph_thread_pool_id = \ |
| 70 | stage_subgraph_thread_pool_id |
| 71 | options.fetch_ops.extend([x.name for x in fetch_tensors]) |
| 72 | options.cancel_op = cancel_fetching.name |
| 73 | options.resume_op = resume_fetching.name |
| 74 | options.close_op = close_fetching.name |
| 75 | |
| 76 | feed_dict = nest.flatten_dict_items(feed_dict) |
| 77 | for feed, feed_val in feed_dict.items(): |
| 78 | for subfeed, subfeed_val in _feed_fn(feed, feed_val): |
| 79 | if not isinstance(subfeed_val, ops.Tensor): |
| 80 | raise TypeError('The value of a feed must be a tf.Tensor object. ' |
| 81 | 'but ' + str(feed) + ' was feed by ' |
| 82 | + str(type(feed_val))) |
| 83 | options.named_feed_input_tensors[subfeed.name]=subfeed_val.name |
| 84 | |
| 85 | for err_code in closed_exception_types: |
| 86 | options.closed_exceptions.append(err_code) |
| 87 | |
| 88 | for err_code in ignored_exception_types: |
| 89 | options.ignored_exceptions.append(err_code) |
| 90 | |
| 91 | @tf_export(v1=["staged"]) |
| 92 | def staged( |