Pass input maping to TF_GraphImportGraphDef for existing nodes

Author: Bart Schuurmans

tensorflow/src/TensorFlow/Internal/FFI.hs

@@ -19,12 +19,29 @@
 module TensorFlow.Internal.FFI
     ( TensorFlowException(..)
     , Raw.Session
-    , Raw.SessionOptions
     , withSession
+    , run
+
     , SessionAction
+
+    , Raw.SessionOptions
+
     , Raw.Graph
-    , importGraphDef
-    , run
+    , Raw.graphOperationByName
+    , resolveOutput
+    , resolveOperation
+    , graphImportGraphDef
+    , forGraphOperations_
+
+    , Raw.Operation
+    , Raw.operationName
+    , Raw.operationNumOutputs
+
+    , Raw.ImportGraphDefOptions
+    , Raw.importGraphDefOptionsAddInputMapping
+
+    , Raw.Output(..)
+
     , TensorData(..)
     , setSessionConfig
     , setSessionTarget
@@ -46,9 +63,8 @@ import Data.Int (Int64)
 import Data.Maybe (fromMaybe)
 import Data.Typeable (Typeable)
 import Data.Word (Word8)
-import Foreign (Ptr, FunPtr, nullPtr, castPtr)
-import Foreign.C.String (CString)
-import Foreign.ForeignPtr (newForeignPtr, newForeignPtr_, withForeignPtr)
+import Foreign (Ptr, FunPtr, nullPtr, castPtr, with)
+import Foreign.ForeignPtr (newForeignPtr_)
 import Foreign.Marshal.Alloc (free)
 import Foreign.Marshal.Array (withArrayLen, peekArray, mallocArray, copyArray)
 import System.IO.Unsafe (unsafePerformIO)
@@ -137,35 +153,50 @@ shutDownRunner r = do
     -- TODO(gnezdo): manage exceptions better than print.
     either print (const (return ())) =<< waitCatch r
 
-importGraphDef :: Raw.Graph -> GraphDef -> IO ()
-importGraphDef graph pb =
+graphImportGraphDef :: Raw.Graph
+                    -> GraphDef
+                    -> (Raw.ImportGraphDefOptions -> IO ())
+                    -> IO ()
+graphImportGraphDef graph pb optionSetter =
     useProtoAsBuffer pb $ \buffer ->
-        bracket Raw.newImportGraphDefOptions Raw.deleteImportGraphDefOptions $ \importGraphDefOptions ->
-            checkStatus $ Raw.importGraphDef graph buffer importGraphDefOptions
+        bracket Raw.newImportGraphDefOptions Raw.deleteImportGraphDefOptions $ \importGraphDefOptions -> do
+            optionSetter importGraphDefOptions
+            checkStatus $ Raw.graphImportGraphDef graph buffer importGraphDefOptions
+
+forGraphOperations_ :: Raw.Graph
+                    -> (Raw.Operation -> IO b)
+                    -> IO ()
+forGraphOperations_ graph f = with 0 go
+  where
+    go indexPtr = do
+        op <- Raw.graphNextOperation graph indexPtr
+        case op of
+          Raw.Operation ptr | ptr == nullPtr -> return ()
+          _ -> f op >> go indexPtr -- indexPtr is modified by Raw.graphNextOperation.
 
 run :: Raw.Session
     -> Raw.Graph
-    -> [(B.ByteString, TensorData)] -- ^ Inputs.
-    -> [B.ByteString]               -- ^ Outputs.
-    -> [B.ByteString]               -- ^ Target operations.
+    -> [(String, TensorData)] -- ^ Inputs.
+    -> [String]               -- ^ Outputs.
+    -> [String]               -- ^ Target operations.
     -> IO [TensorData]
 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_ $
         -- Inputs.
-        mapM (resolveOutput . fst) inputNamesData >>= \inputs ->
+        mapM (resolveOutput graph . fst) inputNamesData >>= \inputs ->
         withArrayLen inputs $ \nInputs cInputs ->
         mapM (createRawTensor . snd) inputNamesData >>= \inputTensors ->
         withArrayLen inputTensors $ \_ cInputTensors ->
         -- Outputs.
-        mapM resolveOutput outputNames >>= \outputs ->
+        mapM (resolveOutput graph) 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 nOutputs nullTensor) $ \_ cOutputTensors ->
         -- Target operations.
-        mapM resolveOperation targetNames >>= \targets ->
+        mapM (resolveOperation graph) targetNames >>= \targets ->
         withArrayLen targets $ \nTargets cTargets -> do
             checkStatus $ Raw.run
                 session
@@ -178,32 +209,33 @@ run session graph inputNamesData outputNames targetNames = do
             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 <- read idxStr
-            -> (C.pack name, idx)
-          _ -> (name, 0)
 
     nullTensor = Raw.Tensor nullPtr
 
+resolveOutput :: Raw.Graph -> String -> IO Raw.Output
+resolveOutput graph name = do
+    let (opName, idx) = parseName name
+    op <- resolveOperation graph opName
+    pure $ Raw.Output op (safeConvert idx)
+  where
+    parseName :: String -> (String, Int)
+    parseName opName =
+        case break (== ':') opName of
+          (opName_, ':':idxStr) | idx <- read idxStr
+            -> (opName_, idx)
+          _ -> (opName, 0)
+
+resolveOperation :: Raw.Graph -> String -> IO Raw.Operation
+resolveOperation graph name = do
+    op <- Raw.graphOperationByName graph name
+    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
+
 
 -- Internal.
 
@@ -218,21 +250,6 @@ safeConvert x =
             show (fromIntegral x :: b)))
     (toIntegralSized x)
 
-
--- | Use a list of ByteString as a list of CString.
-withStringList :: [B.ByteString] -> ([CString] -> IO a) -> IO a
-withStringList strings fn = go strings []
-  where
-    go [] cs = fn (reverse cs)
-    -- TODO(fmayle): Is it worth using unsafeAsCString here?
-    go (x:xs) cs = B.useAsCString x $ \c -> go xs (c:cs)
-
-
--- | Use a list of ByteString as an array of CString.
-withStringArrayLen :: [B.ByteString] -> (Int -> Ptr CString -> IO a) -> IO a
-withStringArrayLen xs fn = withStringList xs (`withArrayLen` fn)
-
-
 -- | Create a Raw.Tensor from a TensorData.
 createRawTensor :: TensorData -> IO Raw.Tensor
 createRawTensor (TensorData dims dt byteVec) =

tensorflow/src/TensorFlow/Internal/Raw.chs

@@ -63,6 +63,8 @@ stringGetSize = {# call TF_StringGetSize as ^ #}
 
 -- Operation.
 {# pointer *TF_Operation as Operation newtype #}
+{# fun TF_OperationName as operationName { `Operation' } -> `String' #}
+{# fun TF_OperationNumOutputs as operationNumOutputs { `Operation' } -> `Int' #}
 
 instance Storable Operation where
     sizeOf (Operation t) = sizeOf t
@@ -74,18 +76,18 @@ instance Storable Operation where
 -- Output.
 data Output = Output
     { outputOperation :: Operation
-    , outputIndex     :: CInt
+    , outputIndex     :: Int
     }
 {# 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
+    peek p = Output <$> {# get TF_Output->oper #} p
+                    <*> (fromIntegral <$> {# get TF_Output->index #} p)
+    poke p (Output oper index) = do
+        {# set TF_Output->oper #} p oper
+        {# set TF_Output->index #} p $ fromIntegral index
 
 
 -- Buffer.
@@ -119,6 +121,8 @@ instance Storable Tensor where
 -- `CLLong`).
 type CInt64 = {#type int64_t #}
 
+{# pointer *size_t as CSizePtr -> CSize #}
+
 newTensor :: DataType
           -> Ptr CInt64   -- dimensions array
           -> CInt         -- num dimensions
@@ -149,27 +153,18 @@ tensorData = {# call TF_TensorData as ^ #}
 
 -- ImportGraphDefOptions.
 {# pointer *TF_ImportGraphDefOptions as ImportGraphDefOptions newtype #}
+{# fun TF_NewImportGraphDefOptions as newImportGraphDefOptions { } -> `ImportGraphDefOptions' #}
+{# fun TF_DeleteImportGraphDefOptions as deleteImportGraphDefOptions { `ImportGraphDefOptions' } -> `()' #}
+{# fun TF_ImportGraphDefOptionsAddInputMapping as importGraphDefOptionsAddInputMapping { `ImportGraphDefOptions', `String', `Int', %`OutputPtr' } -> `()' #}
 
-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 ^ #}
+{# fun TF_NewGraph as newGraph { } -> `Graph' #}
+{# fun TF_DeleteGraph as deleteGraph { `Graph' } -> `()' #}
+{# fun TF_GraphOperationByName as graphOperationByName { `Graph', `String' } -> `Operation' #}
+{# fun TF_GraphNextOperation as graphNextOperation { `Graph', `CSizePtr' } -> `Operation' #}
+{# fun TF_GraphImportGraphDef as graphImportGraphDef { `Graph', `BufferPtr', `ImportGraphDefOptions', `Status' } -> `()' #}
 
 
 -- Session Options.

tensorflow/src/TensorFlow/Session.hs

@@ -45,22 +45,24 @@ import Control.Monad.Trans.Class (MonadTrans, lift)
 import Control.Monad.Trans.Reader (ReaderT(..), ask, asks)
 import Data.ByteString (ByteString)
 import Data.Default (Default, def)
+import Data.Foldable (for_)
 import Data.ProtoLens (showMessage)
 import Data.Set (Set)
-import Data.Text.Encoding (encodeUtf8)
+import Foreign.Marshal.Utils (with)
 import Lens.Family2 (Lens', (^.), (&), (.~))
 import Lens.Family2.Unchecked (lens)
 import Proto.Tensorflow.Core.Framework.Graph (GraphDef)
 import Proto.Tensorflow.Core.Framework.Graph_Fields (node)
 import Proto.Tensorflow.Core.Protobuf.Config (ConfigProto)
 import TensorFlow.Build
 import TensorFlow.Nodes
-import TensorFlow.Output (NodeName, unNodeName)
+import TensorFlow.Output (NodeName(..), unNodeName)
 import TensorFlow.Tensor
 
 import qualified Data.ByteString.Builder as Builder
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
+import qualified Data.Text as T
 import qualified TensorFlow.Internal.FFI as FFI
 
 -- | An action for logging.
@@ -130,8 +132,8 @@ _runSessionWithOptions :: (MonadMask m, MonadIO m)
                        -> m a
 _runSessionWithOptions (Session m) options withSession =
     withSession applyOptions $
-        \asyncCollector rawSession rawGraph ->
-            let initState = SessionState rawSession rawGraph asyncCollector (options ^. sessionTracer)
+        \ac rSession rGraph ->
+            let initState = SessionState rSession rGraph ac (options ^. sessionTracer)
             in evalBuildT (runReaderT m initState)
   where
     applyOptions opt = do
@@ -154,7 +156,18 @@ extend = do
     unless (null nodesToExtend) $ liftIO $ do
         let graphDef = (defMessage :: GraphDef) & node .~ nodesToExtend
         trace ("Session.extend " <> Builder.string8 (showMessage graphDef))
-        FFI.importGraphDef graph graphDef
+        FFI.graphImportGraphDef graph graphDef $ \opts ->
+            -- All inputs of the nodes in the GraphDef should either refer to
+            -- other nodes in the GraphDef, or be mapped to nodes already in
+            -- the Graph by adding an input mapping.
+            -- We add an input mapping for all existing nodes in the Graph in
+            -- case they are referenced in the GraphDef.
+            FFI.forGraphOperations_ graph $ \op -> do
+                srcName <- FFI.operationName op
+                numOutputs <- FFI.operationNumOutputs op
+                for_ [0..numOutputs] $ \srcIndex -> do
+                    let dst = FFI.Output op srcIndex
+                    with dst $ FFI.importGraphDefOptionsAddInputMapping opts srcName srcIndex
     -- Now that all the nodes are created, run the initializers.
     initializers <- build flushInitializers
     unless (null initializers) $
@@ -181,7 +194,7 @@ runFetchWithFeeds :: MonadIO m => [Feed] -> Set NodeName -> Fetch a -> SessionT
 runFetchWithFeeds feeds target (Fetch fetch restore) = do
     extend
     let feeds' = fixFeeds feeds
-    let fetchNames = encodeUtf8 <$> Set.toList fetch
+    let fetchNames = T.unpack <$> Set.toList fetch
         targetNames = toNodeNames $ Set.toList target
     state <- Session ask
     runResult <- liftIO $ FFI.run (rawSession state)
@@ -192,8 +205,8 @@ runFetchWithFeeds feeds target (Fetch fetch restore) = do
     let resultTensorsMap = Map.fromList $ zip (Set.toList fetch) runResult
     return $ restore resultTensorsMap
 
-toNodeNames :: [NodeName] -> [ByteString]
-toNodeNames = map (encodeUtf8 . unNodeName)
+toNodeNames :: [NodeName] -> [String]
+toNodeNames = map (T.unpack . unNodeName)
 
 -- | Run a subgraph 't', rendering and extending any dependent nodes that aren't
 -- already rendered.  This behaves like 'run' except that it doesn't do any
@@ -210,8 +223,8 @@ runWithFeeds_ feeds t = do
     ns <- build $ getNodes t
     runFetchWithFeeds feeds ns (pure ())
 
-fixFeeds :: [Feed] -> [(ByteString, FFI.TensorData)]
-fixFeeds = map $ \(Feed o d) -> (encodeUtf8 $ encodeOutput o, d)
+fixFeeds :: [Feed] -> [(String, FFI.TensorData)]
+fixFeeds = map $ \(Feed o d) -> (T.unpack $ encodeOutput o, d)
 
 -- | Starts a concurrent thread which evaluates the given Nodes
 -- forever until runSession exits or an exception occurs. Graph