tensorflow/cc/experimental/libtf/runtime/runtime.cc
// Import a module from a saved model.
//
// Returns a TaggedValue::Dict. All functions found on the root of the module
// will be attached as callables to this TaggedValue.
//
// `name` specifies the full path to the saved model.
//
// `ctx` should outlive the lifetime of the module.
Import a module from a saved model.

`name` specifies the full path to the saved model.

`ctx` should outlive the lifetime of the module.
Import a module from a saved model . ` name ` specifies the full path to the saved model . ` ctx ` should outlive the lifetime of the module .
static tensorflow::StatusOr<TaggedValue> ImportModule(String name,
                                                      AbstractContext* ctx) {
  tensorflow::StatusOr<TFPackage> tf_package = TFPackage::Load(name.get());
  if (!tf_package.status().ok()) {
    return tf_package.status();
  }
  TaggedValue module = TaggedValue::Dict();
  const RepeatedPtrField<FunctionDef> function_defs =
      tf_package->GetFunctionDefs();
  absl::flat_hash_map<std::string, AbstractFunctionPtr> traces;
  for (auto& fdef : function_defs) {
    AbstractFunctionPtr trace(new GraphFunction(fdef), false);
    traces[fdef.signature().name()] = trace;
  }
   For each child of the root, check what type it is.  If it's a
   SavedFunction, attach that function to the module as a Callable.
  const SavedObjectGraph object_graph = tf_package->GetObjectGraph();
  auto& nodes = object_graph.nodes();
   Get a map of the concrete functions to their input / output signatures.
  auto& concrete_functions = object_graph.concrete_functions();
  auto& root = nodes.at(0);
  for (auto& child : root.children()) {
     The child's name describes the name of the edge that connects to the
     parent object. This name will be the name of the object stored in the
     generated module.
    auto& child_node = nodes.at(child.node_id());
    auto child_name = child.local_name().c_str();
    if (child_node.kind_case() == tensorflow::SavedObject::kFunction) {
      Function tf_function;
      for (const std::string& fn_name :
           child_node.function().concrete_functions()) {
         Setup input signature.
         For now, we have to do a lot of manual digging through these and
         assume they are tensorspecs. Once TODO(b/190203981) is done, we
         should be able to pass along the `StructuredValue`s to an API in a
         much cleaner way.
         TODO(b/190206621): Implement API for inspecting signatures
        SavedConcreteFunction saved_concrete_function =
            concrete_functions.at(fn_name);
        TaggedValue input_signature = TaggedValue::Tuple();
        const RepeatedPtrField<StructuredValue>& args =
            saved_concrete_function.canonicalized_input_signature()
                .tuple_value()
                .values(0)
                .tuple_value()
                .values();
        for (const StructuredValue& arg : args) {
          PartialTensorShape shape = arg.tensor_spec_value().shape();
          DataType dtype = arg.tensor_spec_value().dtype();
          TaggedValue tensor_spec(shape, dtype);
          input_signature.tuple().emplace_back(tensor_spec);
        }
         Setup output signature.
        TensorSpecProto output_tensor_spec_proto =
            saved_concrete_function.output_signature().tensor_spec_value();
        PartialTensorShape output_shape = output_tensor_spec_proto.shape();
        DataType output_dtype = output_tensor_spec_proto.dtype();
        TaggedValue output_tensor_spec(output_shape, output_dtype);
         Register the function trace.
         This does *not* currently register the function with the runtime.
         Instead, we're registering JIT at call time. This is likely
         something that we'll change in TODO(b/190070277).
        auto& trace = traces[fn_name];
        Status status = tf_function.RegisterTrace(
            std::move(trace), input_signature, output_tensor_spec);
      }
      TaggedValue callable = MakeCallable(child_name, tf_function, ctx);
      module.dict()[TaggedValue(child_name)] = callable;
    }
  }
  return module;
}
static tensorflow :: StatusOr < TaggedValue > ImportModule ( String name , AbstractContext * ctx ) { tensorflow :: StatusOr < TFPackage > tf_package = TFPackage :: Load ( name . get ( ) ) ; if ( ! tf_package . status ( ) . ok ( ) ) { return tf_package . status ( ) ; } TaggedValue module = TaggedValue :: Dict ( ) ; const RepeatedPtrField < FunctionDef > function_defs = tf_package -> GetFunctionDefs ( ) ; absl :: flat_hash_map < std :: string , AbstractFunctionPtr > traces ; for ( auto & fdef : function_defs ) { AbstractFunctionPtr trace ( new GraphFunction ( fdef ) , false ) ; traces [ fdef . signature ( ) . name ( ) ] = trace ; } const SavedObjectGraph object_graph = tf_package -> GetObjectGraph ( ) ; auto & nodes = object_graph . nodes ( ) ; auto & concrete_functions = object_graph . concrete_functions ( ) ; auto & root = nodes . at ( 0 ) ; for ( auto & child : root . children ( ) ) { auto & child_node = nodes . at ( child . node_id ( ) ) ; auto child_name = child . local_name ( ) . c_str ( ) ; if ( child_node . kind_case ( ) == tensorflow :: SavedObject :: kFunction ) { Function tf_function ; for ( const std :: string & fn_name : child_node . function ( ) . concrete_functions ( ) ) { SavedConcreteFunction saved_concrete_function = concrete_functions . at ( fn_name ) ; TaggedValue input_signature = TaggedValue :: Tuple ( ) ; const RepeatedPtrField < StructuredValue > & args = saved_concrete_function . canonicalized_input_signature ( ) . tuple_value ( ) . values ( 0 ) . tuple_value ( ) . values ( ) ; for ( const StructuredValue & arg : args ) { PartialTensorShape shape = arg . tensor_spec_value ( ) . shape ( ) ; DataType dtype = arg . tensor_spec_value ( ) . dtype ( ) ; TaggedValue tensor_spec ( shape , dtype ) ; input_signature . tuple ( ) . emplace_back ( tensor_spec ) ; } TensorSpecProto output_tensor_spec_proto = saved_concrete_function . output_signature ( ) . tensor_spec_value ( ) ; PartialTensorShape output_shape = output_tensor_spec_proto . shape ( ) ; DataType output_dtype = output_tensor_spec_proto . dtype ( ) ; TaggedValue output_tensor_spec ( output_shape , output_dtype ) ; auto & trace = traces [ fn_name ] ; Status status = tf_function . RegisterTrace ( std :: move ( trace ) , input_signature , output_tensor_spec ) ; } TaggedValue callable = MakeCallable ( child_name , tf_function , ctx ) ; module . dict ( ) [ TaggedValue ( child_name ) ] = callable ; } } return module ; }
Import a module from a saved model.
Import a module from a saved model .
// Load the serialized model. // Initialize concrete function traces. /*add_ref=*/ // Setup polymorphic functions and wrap in Callables. // // For each child of the root, check what type it is.  If it's a // SavedFunction, attach that function to the module as a Callable. // Get a map of the concrete functions to their input / output signatures. // The child's name describes the name of the edge that connects to the // parent object. This name will be the name of the object stored in the // generated module. // Setup input signature. // // For now, we have to do a lot of manual digging through these and // assume they are tensorspecs. Once TODO(b/190203981) is done, we // should be able to pass along the `StructuredValue`s to an API in a // much cleaner way. // // TODO(b/190206621): Implement API for inspecting signatures // Setup output signature. // Register the function trace. // // This does *not* currently register the function with the runtime. // Instead, we're registering JIT at call time. This is likely // something that we'll change in TODO(b/190070277).