Add support for loading a Session from a SavedModel

Author: Bart Schuurmans

tensorflow/src/TensorFlow/Core.hs

@@ -30,6 +30,9 @@ module TensorFlow.Core
     , sessionTracer
     , runSession
     , runSessionWithOptions
+    , SavedModelTag(..)
+    , runSavedModel
+    , runSavedModelWithOptions
       -- ** Building graphs
     , MonadBuild(..)
       -- ** Running graphs

tensorflow/src/TensorFlow/Internal/FFI.hs

@@ -21,6 +21,7 @@ module TensorFlow.Internal.FFI
     , Raw.Session
     , Raw.SessionOptions
     , withSession
+    , withSessionFromSavedModel
     , SessionAction
     , Raw.Graph
     , importGraphDef
@@ -45,12 +46,14 @@ import Data.Maybe (fromMaybe)
 import Data.Typeable (Typeable)
 import Data.Word (Word8)
 import Foreign (Ptr, FunPtr, nullPtr, castPtr)
+import Foreign.C (CInt)
 import Foreign.C.String (CString)
 import Foreign.ForeignPtr (newForeignPtr, newForeignPtr_, withForeignPtr)
 import Foreign.Marshal.Alloc (free)
 import Foreign.Marshal.Array (withArrayLen, peekArray, mallocArray, copyArray)
 import System.IO.Unsafe (unsafePerformIO)
 import qualified Data.ByteString as B
+import qualified Data.ByteString.Char8 as C
 import qualified Data.Text as T
 import qualified Data.Text.Encoding as T
 import qualified Data.Text.Encoding.Error as T
@@ -89,6 +92,30 @@ withSession :: (MonadIO m, MonadMask m)
             -> m a
 withSession = withSession_ Raw.newSession
 
+withSessionFromSavedModel :: (MonadIO m, MonadMask m)
+                          => B.ByteString
+                          -- ^ exportDir
+                          -> [B.ByteString]
+                          -- ^ Tags.
+                          -> (Raw.SessionOptions -> IO ())
+                          -- ^ optionSetter
+                          -> SessionAction m a
+                          -> m a
+withSessionFromSavedModel exportDir tags =
+    withSession_ $ \graph options status ->
+    B.useAsCString exportDir $ \cExportDir ->
+    withStringArrayLen tags $ \nTags cTags ->
+        Raw.loadSessionFromSavedModel options
+                                      runOptions
+                                      cExportDir
+                                      cTags (safeConvert nTags)
+                                      graph
+                                      metaGraphDef
+                                      status
+  where
+    runOptions = nullPtr
+    metaGraphDef = nullPtr
+
 withSession_ :: (MonadIO m, MonadMask m)
              => (Raw.Graph -> Raw.SessionOptions -> Raw.Status -> IO Raw.Session)
              -- ^ mkSession

tensorflow/src/TensorFlow/Internal/Raw.chs

@@ -177,6 +177,15 @@ deleteSessionOptions = {# call TF_DeleteSessionOptions as ^ #}
 newSession :: Graph -> SessionOptions -> Status -> IO Session
 newSession = {# call TF_NewSession as ^ #}
 
+loadSessionFromSavedModel :: SessionOptions
+                          -> BufferPtr           -- RunOptions proto.
+                          -> CString             -- Export directory.
+                          -> Ptr CString -> CInt -- Tags.
+                          -> Graph               -- Graph.
+                          -> BufferPtr           -- MetaGraphDef.
+                          -> Status
+                          -> IO Session
+loadSessionFromSavedModel = {# call TF_LoadSessionFromSavedModel as ^ #}
 
 closeSession :: Session -> Status -> IO ()
 closeSession = {# call TF_CloseSession as ^ #}

tensorflow/src/TensorFlow/Session.hs

@@ -27,13 +27,16 @@ module TensorFlow.Session (
     sessionTracer,
     runSession,
     runSessionWithOptions,
+    runSavedModel,
+    runSavedModelWithOptions,
     MonadBuild(..),
     addGraphDef,
     run,
     runWithFeeds,
     run_,
     runWithFeeds_,
     asyncProdNodes,
+    SavedModelTag(..),
     ) where
 
 import Data.ProtoLens.Message(defMessage)
@@ -57,6 +60,7 @@ import TensorFlow.Nodes
 import TensorFlow.Output (NodeName, unNodeName)
 import TensorFlow.Tensor
 
+import qualified Data.ByteString.Char8 as C
 import qualified Data.ByteString.Builder as Builder
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
@@ -92,6 +96,14 @@ data Options = Options
     , _sessionTracer :: Tracer
     }
 
+data SavedModelTag = GPU | TPU | Serve | Train
+
+savedModelTagValue :: SavedModelTag -> ByteString
+savedModelTagValue GPU = "gpu"
+savedModelTagValue TPU = "tpu"
+savedModelTagValue Serve = "serve"
+savedModelTagValue Train = "train"
+
 instance Default Options where
     def = Options
           { _sessionTarget = ""
@@ -122,6 +134,25 @@ runSessionWithOptions :: (MonadMask m, MonadIO m) => Options -> SessionT m a ->
 runSessionWithOptions options session =
     _runSessionWithOptions session options $ FFI.withSession
 
+runSavedModel :: (MonadMask m, MonadIO m)
+              => FilePath
+              -- ^ Export directory.
+              -> Set SavedModelTag
+              -> SessionT m a
+              -> m a
+runSavedModel exportDir tags = runSavedModelWithOptions exportDir tags def
+
+runSavedModelWithOptions :: (MonadMask m, MonadIO m)
+                         => FilePath
+                         -- ^ Export directory.
+                         -> Set SavedModelTag
+                         -> Options
+                         -> SessionT m a
+                         -> m a
+runSavedModelWithOptions exportDir tags options session =
+    _runSessionWithOptions session options $
+        FFI.withSessionFromSavedModel (C.pack exportDir) (map savedModelTagValue $ Set.toList tags)
+
 _runSessionWithOptions :: (MonadMask m, MonadIO m)
                        => SessionT m a
                        -> Options