Migrate from TF_DeprecatedSession to TF_Session

Author: Bart Schuurmans

tensorflow/src/TensorFlow/Internal/FFI.hs

@@ -19,8 +19,11 @@
 module TensorFlow.Internal.FFI
     ( TensorFlowException(..)
     , Raw.Session
+    , Raw.SessionOptions
     , withSession
-    , extendGraph
+    , SessionAction
+    , Raw.Graph
+    , importGraphDef
     , run
     , TensorData(..)
     , setSessionConfig
@@ -74,15 +77,26 @@ data TensorData = TensorData
     }
   deriving (Show, Eq)
 
+-- | The action can spawn concurrent tasks which will be canceled before
+-- withSession returns.
+type SessionAction m a = (IO () -> IO ()) -> Raw.Session -> Raw.Graph -> m a
+
 -- | Runs the given action after creating a session with options
 -- populated by the given optionSetter.
 withSession :: (MonadIO m, MonadMask m)
             => (Raw.SessionOptions -> IO ())
-            -> ((IO () -> IO ()) -> Raw.Session -> m a)
-            -- ^ The action can spawn concurrent tasks which will
-            -- be canceled before withSession returns.
+            -> SessionAction m a
             -> m a
-withSession optionSetter action = do
+withSession = withSession_ Raw.newSession
+
+withSession_ :: (MonadIO m, MonadMask m)
+             => (Raw.Graph -> Raw.SessionOptions -> Raw.Status -> IO Raw.Session)
+             -- ^ mkSession
+             -> (Raw.SessionOptions -> IO ())
+             -- ^ optionSetter
+             -> SessionAction m a
+             -> m a
+withSession_ mkSession optionSetter action = do
     drain <- liftIO $ newMVar []
     let cleanup s =
         -- Closes the session to nudge the pending run calls to fail and exit.
@@ -92,11 +106,12 @@ withSession optionSetter action = do
                 mapM_ shutDownRunner runners
                 checkStatus (Raw.deleteSession s)
     let bracketIO x y = bracket (liftIO x) (liftIO . y)
-    bracketIO Raw.newSessionOptions Raw.deleteSessionOptions $ \options -> do
-        bracketIO
-            (optionSetter options >> checkStatus (Raw.newSession options))
-            cleanup
-            (action (asyncCollector drain))
+    bracketIO Raw.newGraph Raw.deleteGraph $ \graph ->
+        bracketIO Raw.newSessionOptions Raw.deleteSessionOptions $ \options -> do
+            bracketIO
+                (optionSetter options >> checkStatus (mkSession graph options))
+                cleanup
+                (\session -> action (asyncCollector drain) session graph)
 
 asyncCollector :: MVar [Async ()] -> IO () -> IO ()
 asyncCollector drain runner = modifyMVarMasked_ drain launchAndRecord
@@ -109,43 +124,72 @@ shutDownRunner r = do
     -- TODO(gnezdo): manage exceptions better than print.
     either print (const (return ())) =<< waitCatch r
 
-extendGraph :: Raw.Session -> GraphDef -> IO ()
-extendGraph session pb =
-    useProtoAsVoidPtrLen pb $ \ptr len ->
-        checkStatus $ Raw.extendGraph session ptr len
-
+importGraphDef :: Raw.Graph -> GraphDef -> IO ()
+importGraphDef graph pb =
+    useProtoAsBuffer pb $ \buffer ->
+        bracket Raw.newImportGraphDefOptions Raw.deleteImportGraphDefOptions $ \importGraphDefOptions ->
+            checkStatus $ Raw.importGraphDef graph buffer importGraphDefOptions
 
 run :: Raw.Session
-    -> [(B.ByteString, TensorData)] -- ^ Feeds.
-    -> [B.ByteString]               -- ^ Fetches.
-    -> [B.ByteString]               -- ^ Targets.
+    -> Raw.Graph
+    -> [(B.ByteString, TensorData)] -- ^ Inputs.
+    -> [B.ByteString]               -- ^ Outputs.
+    -> [B.ByteString]               -- ^ Target operations.
     -> IO [TensorData]
-run session feeds fetches targets = do
-    let nullTensor = Raw.Tensor nullPtr
+run session graph inputNamesData outputNames targetNames = do
     -- Use mask to avoid leaking input tensors before they are passed to 'run'
     -- and output tensors before they are passed to 'createTensorData'.
     mask_ $
-        -- Feeds
-        withStringArrayLen (fst <$> feeds) $ \feedsLen feedNames ->
-        mapM (createRawTensor . snd) feeds >>= \feedTensors ->
-        withArrayLen feedTensors $ \_ cFeedTensors ->
-        -- Fetches.
-        withStringArrayLen fetches $ \fetchesLen fetchNames ->
-        -- tensorOuts is an array of null Tensor pointers that will be filled
+        -- Inputs.
+        mapM (resolveOutput . fst) inputNamesData >>= \inputs ->
+        withArrayLen inputs $ \nInputs cInputs ->
+        mapM (createRawTensor . snd) inputNamesData >>= \inputTensors ->
+        withArrayLen inputTensors $ \_ cInputTensors ->
+        -- Outputs.
+        mapM resolveOutput outputNames >>= \outputs ->
+        withArrayLen outputs $ \nOutputs cOutputs ->
+        -- outputTensors is an array of null Tensor pointers that will be filled
         -- by the call to Raw.run.
-        withArrayLen (replicate fetchesLen nullTensor) $ \_ tensorOuts ->
-        -- Targets.
-        withStringArrayLen targets $ \targetsLen ctargets -> do
+        withArrayLen (replicate nOutputs nullTensor) $ \_ cOutputTensors ->
+        -- Target operations.
+        mapM resolveOperation targetNames >>= \targets ->
+        withArrayLen targets $ \nTargets cTargets -> do
             checkStatus $ Raw.run
                 session
-                nullPtr
-                feedNames cFeedTensors (safeConvert feedsLen)
-                fetchNames tensorOuts (safeConvert fetchesLen)
-                ctargets (safeConvert targetsLen)
-                nullPtr
-            mapM_ Raw.deleteTensor feedTensors
-            outTensors <- peekArray fetchesLen tensorOuts
+                nullPtr -- RunOptions proto.
+                cInputs  cInputTensors  (safeConvert nInputs)
+                cOutputs cOutputTensors (safeConvert nOutputs)
+                cTargets                (safeConvert nTargets)
+                nullPtr -- RunMetadata.
+            mapM_ Raw.deleteTensor inputTensors
+            outTensors <- peekArray nOutputs cOutputTensors
             mapM createTensorData outTensors
+  where
+    resolveOutput :: B.ByteString -> IO Raw.Output
+    resolveOutput name = do
+        let (opName, idx) = parseName name
+        op <- resolveOperation opName
+        pure $ Raw.Output op idx
+
+    resolveOperation :: B.ByteString -> IO Raw.Operation
+    resolveOperation name = do
+        op <- B.useAsCString name $ Raw.graphOperationByName graph
+        case op of
+          Raw.Operation ptr | ptr == nullPtr -> throwM exception
+          _ -> pure op
+      where
+        exception =
+            let msg = "Operation not found in graph: " <> (T.pack $ show name)
+            in TensorFlowException Raw.TF_INVALID_ARGUMENT msg
+
+    parseName :: B.ByteString -> (B.ByteString, CInt)
+    parseName name =
+        case break (== ':') (C.unpack name) of
+          (name, ':':idxStr) | [(idx, "" :: String)] <- read idxStr
+            -> (C.pack name, fromInteger idx)
+          _ -> (name, 0)
+
+    nullTensor = Raw.Tensor nullPtr
 
 
 -- Internal.
@@ -245,18 +289,26 @@ useProtoAsVoidPtrLen :: (Message msg, Integral c, Show c, Bits c) =>
 useProtoAsVoidPtrLen msg f = B.useAsCStringLen (encodeMessage msg) $
         \(bytes, len) -> f (castPtr bytes) (safeConvert len)
 
+-- | Serializes the given msg and provides it as BufferPtr argument
+-- to the given action.
+useProtoAsBuffer :: (Message msg) =>
+                    msg -> (Raw.BufferPtr -> IO a) -> IO a
+useProtoAsBuffer msg f =
+    B.useAsCStringLen (encodeMessage msg) $ \(bytes, len) ->
+        bracket (Raw.newBufferFromString (castPtr bytes) (safeConvert len))
+                Raw.deleteBuffer
+                f
+
 -- | Returns the serialized OpList of all OpDefs defined in this
 -- address space.
 getAllOpList :: IO B.ByteString
-getAllOpList = do
-    foreignPtr <-
-        mask_ (newForeignPtr Raw.deleteBuffer =<< checkCall)
-    -- Makes a copy because it is more reliable than eviscerating
-    -- Buffer to steal its memory (including custom deallocator).
-    withForeignPtr foreignPtr $
-        \ptr -> B.packCStringLen =<< (,)
-                <$> (castPtr <$> Raw.getBufferData ptr)
-                <*> (safeConvert <$> Raw.getBufferLength ptr)
+getAllOpList =
+    bracket checkCall Raw.deleteBuffer $ \buffer ->
+        -- Makes a copy because it is more reliable than eviscerating
+        -- Buffer to steal its memory (including custom deallocator).
+        B.packCStringLen =<< (,)
+            <$> (castPtr <$> Raw.getBufferData buffer)
+            <*> (safeConvert <$> Raw.getBufferLength buffer)
     where
       checkCall = do
           p <- Raw.getAllOpList

tensorflow/src/TensorFlow/Internal/Raw.chs

@@ -44,6 +44,33 @@ message :: Status -> IO CString
 message = {# call TF_Message as ^ #}
 
 
+-- Operation.
+{# pointer *TF_Operation as Operation newtype #}
+
+instance Storable Operation where
+    sizeOf (Operation t) = sizeOf t
+    alignment (Operation t) = alignment t
+    peek p = fmap Operation (peek (castPtr p))
+    poke p (Operation t) = poke (castPtr p) t
+
+
+-- Output.
+data Output = Output
+    { outputOperation :: Operation
+    , outputIndex     :: CInt
+    }
+{# pointer *TF_Output as OutputPtr -> Output #}
+
+instance Storable Output where
+    sizeOf _ = {# sizeof TF_Output #}
+    alignment _ = {# alignof TF_Output #}
+    peek ptr = Output <$> {# get TF_Output->oper  #} ptr
+                      <*> {# get TF_Output->index #} ptr
+    poke ptr (Output oper index) = do
+        {# set TF_Output->oper  #} ptr oper
+        {# set TF_Output->index #} ptr index
+
+
 -- Buffer.
 data Buffer
 {# pointer *TF_Buffer as BufferPtr -> Buffer #}
@@ -54,6 +81,12 @@ getBufferData = {# get TF_Buffer->data #}
 getBufferLength :: BufferPtr -> IO CULong
 getBufferLength = {# get TF_Buffer->length #}
 
+newBufferFromString :: Ptr () -> CULong -> IO BufferPtr
+newBufferFromString = {# call TF_NewBufferFromString as ^ #}
+
+deleteBuffer :: BufferPtr -> IO ()
+deleteBuffer = {# call TF_DeleteBuffer as ^ #}
+
 -- Tensor.
 {# pointer *TF_Tensor as Tensor newtype #}
 
@@ -97,6 +130,30 @@ tensorByteSize = {# call TF_TensorByteSize as ^ #}
 tensorData :: Tensor -> IO (Ptr ())
 tensorData = {# call TF_TensorData as ^ #}
 
+-- ImportGraphDefOptions.
+{# pointer *TF_ImportGraphDefOptions as ImportGraphDefOptions newtype #}
+
+newImportGraphDefOptions :: IO ImportGraphDefOptions
+newImportGraphDefOptions = {# call TF_NewImportGraphDefOptions as ^ #}
+
+deleteImportGraphDefOptions :: ImportGraphDefOptions -> IO ()
+deleteImportGraphDefOptions = {# call TF_DeleteImportGraphDefOptions as ^ #}
+
+-- Graph.
+{# pointer *TF_Graph as Graph newtype #}
+
+newGraph :: IO Graph
+newGraph = {# call TF_NewGraph as ^ #}
+
+deleteGraph :: Graph -> IO ()
+deleteGraph = {# call TF_DeleteGraph as ^ #}
+
+graphOperationByName :: Graph -> CString -> IO Operation
+graphOperationByName = {# call TF_GraphOperationByName as ^ #}
+
+importGraphDef :: Graph -> BufferPtr -> ImportGraphDefOptions -> Status -> IO ()
+importGraphDef = {# call TF_GraphImportGraphDef as ^ #}
+
 
 -- Session Options.
 {# pointer *TF_SessionOptions as SessionOptions newtype #}
@@ -115,29 +172,27 @@ deleteSessionOptions = {# call TF_DeleteSessionOptions as ^ #}
 
 
 -- Session.
-{# pointer *TF_DeprecatedSession as Session newtype #}
+{# pointer *TF_Session as Session newtype #}
+
+newSession :: Graph -> SessionOptions -> Status -> IO Session
+newSession = {# call TF_NewSession as ^ #}
 
-newSession :: SessionOptions -> Status -> IO Session
-newSession = {# call TF_NewDeprecatedSession as ^ #}
 
 closeSession :: Session -> Status -> IO ()
-closeSession = {# call TF_CloseDeprecatedSession as ^ #}
+closeSession = {# call TF_CloseSession as ^ #}
 
 deleteSession :: Session -> Status -> IO ()
-deleteSession = {# call TF_DeleteDeprecatedSession as ^ #}
-
-extendGraph :: Session -> Ptr () -> CULong -> Status -> IO ()
-extendGraph = {# call TF_ExtendGraph as ^ #}
+deleteSession = {# call TF_DeleteSession as ^ #}
 
 run :: Session
-    -> BufferPtr                          -- RunOptions proto.
-    -> Ptr CString -> Ptr Tensor -> CInt  -- Input (names, tensors, count).
-    -> Ptr CString -> Ptr Tensor -> CInt  -- Output (names, tensors, count).
-    -> Ptr CString -> CInt                -- Target nodes (names, count).
-    -> BufferPtr                          -- RunMetadata proto.
+    -> BufferPtr                       -- RunOptions proto.
+    -> OutputPtr -> Ptr Tensor -> CInt -- Input (names, tensors, count).
+    -> OutputPtr -> Ptr Tensor -> CInt -- Output (names, tensors, count).
+    -> Ptr Operation           -> CInt -- Target operations (names, count).
+    -> BufferPtr                       -- RunMetadata proto.
     -> Status
     -> IO ()
-run = {# call TF_Run as ^ #}
+run = {# call TF_SessionRun as ^ #}
 
 -- FFI helpers.
 type TensorDeallocFn = Ptr () -> CULong -> Ptr () -> IO ()
@@ -153,6 +208,3 @@ foreign import ccall "wrapper"
 -- in this address space.
 getAllOpList :: IO BufferPtr
 getAllOpList = {# call TF_GetAllOpList as ^ #}
-
-foreign import ccall "&TF_DeleteBuffer"
-  deleteBuffer :: FunPtr (BufferPtr -> IO ())