| 42 | : iterator_(iterator), ctx_(ctx) {} |
| 43 | |
| 44 | Status Dataset::FromGraph(Params params, const GraphDef& graph_def, |
| 45 | std::unique_ptr<Dataset>* result) { |
| 46 | Graph graph(OpRegistry::Global()); |
| 47 | TF_RETURN_IF_ERROR(ImportGraphDef({}, graph_def, &graph, nullptr)); |
| 48 | |
| 49 | // Instantiate enough of the TF runtime to run `graph` on a single CPU device. |
| 50 | auto device_mgr = absl::make_unique<DeviceMgr>(DeviceFactory::NewDevice( |
| 51 | "CPU", params.session_options, "/job:localhost/replica:0/task:0")); |
| 52 | Device* device = device_mgr->ListDevices()[0]; |
| 53 | // Clone the `FunctionLibraryDefinition` to extend its lifetime extends beyond |
| 54 | // the lifetime of `graph`. |
| 55 | auto flib_def = |
| 56 | absl::make_unique<FunctionLibraryDefinition>(graph.flib_def()); |
| 57 | auto pflr = absl::make_unique<ProcessFunctionLibraryRuntime>( |
| 58 | device_mgr.get(), Env::Default(), TF_GRAPH_DEF_VERSION, flib_def.get(), |
| 59 | OptimizerOptions{}, nullptr /* parent */); |
| 60 | |
| 61 | string fetch_node = ""; |
| 62 | for (auto node : graph_def.node()) { |
| 63 | if (node.op() == "_Retval") { |
| 64 | fetch_node = node.input(0); |
| 65 | } |
| 66 | } |
| 67 | if (fetch_node.empty()) { |
| 68 | return errors::NotFound("Failed to find a _Retval op in the given dataset"); |
| 69 | } |
| 70 | |
| 71 | // Run graph up to `output_node` and extract the `DatasetBase` stored in the |
| 72 | // DT_VARIANT output tensor. |
| 73 | data::DatasetBase* dataset; |
| 74 | { |
| 75 | std::vector<Tensor> outputs; |
| 76 | GraphRunner graph_runner(device); |
| 77 | TF_RETURN_IF_ERROR(graph_runner.Run(&graph, pflr->GetFLR("/device:CPU:0"), |
| 78 | {}, {fetch_node}, &outputs)); |
| 79 | TF_RETURN_IF_ERROR(GetDatasetFromVariantTensor(outputs[0], &dataset)); |
| 80 | // NOTE(mrry): The dataset is currently owned by `outputs[0]`, so acquire an |
| 81 | // additional reference. |
| 82 | dataset->Ref(); |
| 83 | } |
| 84 | |
| 85 | std::unique_ptr<thread::ThreadPool> pool( |
| 86 | NewThreadPoolFromSessionOptions(params.session_options)); |
| 87 | *result = |
| 88 | WrapUnique(new Dataset(dataset, device_mgr.release(), pflr.release(), |
| 89 | flib_def.release(), pool.release())); |
| 90 | return Status::OK(); |
| 91 | } // static |
| 92 | |
| 93 | Status Dataset::MakeIterator(std::unique_ptr<Iterator>* result) { |
| 94 | // Create an `IteratorContext`, which bundles together the necessary runtime |
nothing calls this directly
no test coverage detected