1 /* 2 Copyright 2017 The TensorFlow Authors. All Rights Reserved. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package tensorflow 18 19 import ( 20 "runtime" 21 "unsafe" 22 ) 23 24 // #include <stdlib.h> 25 // #include "tensorflow/c/c_api.h" 26 import "C" 27 28 // SavedModel represents the contents of loaded SavedModel. 29 // TODO(jhseu): Add and document metagraphdef when we pregenerate protobufs. 30 type SavedModel struct { 31 Session *Session 32 Graph *Graph 33 } 34 35 // LoadSavedModel creates a new SavedModel from a model previously 36 // exported to a directory on disk. 37 // 38 // Exported models contain a set of graphs and, optionally, variable values. 39 // Tags in the model identify a single graph. LoadSavedModel initializes a 40 // session with the identified graph and with variables initialized to from the 41 // checkpoints on disk. 42 // 43 // The tensorflow package currently does not have the ability to export a model 44 // to a directory from Go. This function thus currently targets loading models 45 // exported in other languages, such as using tf.saved_model.builder in Python. 46 // See: 47 // https://www.tensorflow.org/code/tensorflow/python/saved_model/ 48 func LoadSavedModel(exportDir string, tags []string, options *SessionOptions) (*SavedModel, error) { 49 status := newStatus() 50 cOpt, doneOpt, err := options.c() 51 defer doneOpt() 52 if err != nil { 53 return nil, err 54 } 55 cExportDir := C.CString(exportDir) 56 cTags := make([]*C.char, len(tags)) 57 for i := range tags { 58 cTags[i] = C.CString(tags[i]) 59 } 60 graph := NewGraph() 61 // TODO(jhseu): Add support for run_options and meta_graph_def. 62 cSess := C.TF_LoadSessionFromSavedModel(cOpt, nil, cExportDir, (**C.char)(unsafe.Pointer(&cTags[0])), C.int(len(cTags)), graph.c, nil, status.c) 63 for i := range cTags { 64 C.free(unsafe.Pointer(cTags[i])) 65 } 66 C.free(unsafe.Pointer(cExportDir)) 67 68 if err := status.Err(); err != nil { 69 return nil, err 70 } 71 s := &Session{c: cSess} 72 runtime.SetFinalizer(s, func(s *Session) { s.Close() }) 73 return &SavedModel{Session: s, Graph: graph}, nil 74 } 75
