1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_ 17 #define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_ 18 19 #include <memory> 20 #include <string> 21 22 #include "absl/types/optional.h" 23 #include "tensorflow/c/eager/immediate_execution_tensor_handle.h" 24 #include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h" 25 #include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h" 26 27 namespace tensorflow { 28 29 // RestoredResource represents a TF2 "Resource" object loaded from a savedmodel, 30 // analogous to the Python _RestoredResource object: 31 // https://github.com/tensorflow/tensorflow/blob/fda326e542ca67534e8411edb180e8760a4828b7/tensorflow/python/saved_model/load.py#L481 32 // TF2 resource objects typically extend TrackableResource: 33 // https://github.com/tensorflow/tensorflow/blob/fda326e542ca67534e8411edb180e8760a4828b7/tensorflow/python/training/tracking/tracking.py#L285 34 // and are expected to implement "_create_resource", "_initialize", and 35 // "_destroy_resource" functions: 36 // https://github.com/tensorflow/tensorflow/blob/139ba9c5284799beafdd1d7f895127cf00e7c48f/tensorflow/python/training/tracking/tracking.py#L262-L281 37 class RestoredResource : TensorHandleConvertible { 38 public: 39 // Note(bmzhao): RestoredResource stores non-owning pointers to its associated 40 // functions because SavedModel internally owns all functions and objects in 41 // the RevivedObjects struct (which owns all functions). One alternative would 42 // be to have RevivedObjects store shared_ptr<TFConcreteFunction> instead, and 43 // change RestoredResource's constructor take shared_ptr<TFConcreteFunction>. 44 // To keep things simple, I've stuck to raw pointers for now. 45 // 46 // Params: 47 // device - The device string associated with the SavedResource 48 // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/saved_object_graph.proto#L182 49 // Conceptually, this is the same device used in CapturableResource: 50 // https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/python/training/tracking/tracking.py#L222-L225 51 // Implementation-wise, it is device used when invoking the 52 // create_resource function to produce the resource_handle 53 // associated with the object: 54 // https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/python/training/tracking/tracking.py#L246-L247 55 // create_resource - Non owning pointer to the create_resource function 56 // associated with this object. Must be NON-NULL. 57 // initialize - Non owning pointer to the initialize function associated with 58 // this object. Must be NON-NULL. 59 // destroy_resource - Non owning pointer to the destroy_resource function 60 // associated with this object. Ideally this should be 61 // NON-NULL, but in order to support models saved prior to 62 // https://github.com/tensorflow/tensorflow/commit/3c806101f57768e479f8646e7518bbdff1632ca3 63 // we allow null here. This will, however, leak resources. 64 RestoredResource(const std::string& device, 65 TFConcreteFunction* create_resource, 66 TFConcreteFunction* initialize, 67 TFConcreteFunction* destroy_resource, 68 ImmediateTensorHandlePtr resource_handle); 69 70 Status Initialize() const; 71 72 // RestoredResource is movable, but not copyable. 73 RestoredResource(RestoredResource&& other) = default; 74 RestoredResource& operator=(RestoredResource&& other) = default; 75 76 ~RestoredResource() override; 77 78 private: 79 std::string device_; 80 TFConcreteFunction* create_resource_; 81 TFConcreteFunction* initialize_; 82 TFConcreteFunction* destroy_resource_; 83 }; 84 85 } // namespace tensorflow 86 87 #endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_ 88