diff --git a/src/TensorFlow/Logging.hs b/src/TensorFlow/Logging.hs --- a/src/TensorFlow/Logging.hs +++ b/src/TensorFlow/Logging.hs @@ -33,15 +33,19 @@ -- > TF.logSummary eventWriter step summary -- > else TF.run_ trainStep +{-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeOperators #-} + module TensorFlow.Logging ( EventWriter , withEventWriter , logEvent + , logGraph , logSummary , SummaryTensor , histogramSummary + , imageSummary , scalarSummary , mergeAllSummaries ) where @@ -54,24 +58,25 @@ import Control.Monad.IO.Class (MonadIO, liftIO) import Control.Monad.Trans.Resource (runResourceT) import Data.ByteString (ByteString) -import Data.Conduit ((=$=)) +import Data.Conduit ((.|)) import Data.Conduit.TQueue (sourceTBMQueue) import Data.Default (def) import Data.Int (Int64) +import Data.Word (Word8, Word16) import Data.ProtoLens (encodeMessage) import Data.Time.Clock (getCurrentTime) import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds) import Lens.Family2 ((.~), (&)) import Network.HostName (getHostName) import Proto.Tensorflow.Core.Framework.Summary (Summary) -import Proto.Tensorflow.Core.Util.Event (Event, fileVersion, step, summary, wallTime) +import Proto.Tensorflow.Core.Util.Event (Event, fileVersion, graphDef, step, summary, wallTime) import System.Directory (createDirectoryIfMissing) import System.FilePath (()) -import TensorFlow.Build (MonadBuild) +import TensorFlow.Build (MonadBuild, Build, asGraphDef) import TensorFlow.Ops (scalar) import TensorFlow.Records.Conduit (sinkTFRecords) import TensorFlow.Tensor (Tensor, render, SummaryTensor, addSummary, collectAllSummaries) -import TensorFlow.Types (TensorType, type(/=)) +import TensorFlow.Types (TensorType, type(/=), OneOf) import Text.Printf (printf) import qualified Data.ByteString.Lazy as L import qualified Data.Conduit as Conduit @@ -108,8 +113,8 @@ let writer = EventWriter q done consumeQueue = runResourceT $ Conduit.runConduit $ sourceTBMQueue q - =$= Conduit.map (L.fromStrict . encodeMessage) - =$= sinkTFRecords filename + .| Conduit.map (L.fromStrict . encodeMessage) + .| sinkTFRecords filename _ <- forkFinally consumeQueue (\_ -> putMVar done ()) logEvent writer $ def & wallTime .~ t & fileVersion .~ T.pack "brain.Event:2" @@ -123,6 +128,14 @@ logEvent :: MonadIO m => EventWriter -> Event -> m () logEvent (EventWriter q _) pb = liftIO (atomically (writeTBMQueue q pb)) +-- | Logs the graph for the given 'Build' action. +logGraph :: MonadIO m => EventWriter -> Build a -> m () +logGraph writer build = do + let graph = asGraphDef build + graphBytes = encodeMessage graph + graphEvent = (def :: Event) & graphDef .~ graphBytes + logEvent writer graphEvent + -- | Logs the given Summary event with an optional global step (use 0 if not -- applicable). logSummary :: MonadIO m => EventWriter -> Int64 -> Summary -> m () @@ -145,6 +158,16 @@ -- OneOf '[Int16, Int32, Int64, Int8, Word16, Word8, Double, Float] t) => ByteString -> Tensor v t -> m () histogramSummary tag = addSummary . CoreOps.histogramSummary (scalar tag) + +-- | Adds a 'CoreOps.imageSummary' node. The tag argument is intentionally +-- limited to a single value for simplicity. +imageSummary :: + (OneOf '[Word8, Word16, Float] t, MonadBuild m) + => ByteString + -> Tensor v t + -> m () + +imageSummary tag = addSummary . CoreOps.imageSummary (scalar tag) -- | Adds a 'CoreOps.scalarSummary' node. scalarSummary :: diff --git a/tensorflow-logging.cabal b/tensorflow-logging.cabal --- a/tensorflow-logging.cabal +++ b/tensorflow-logging.cabal @@ -1,5 +1,5 @@ name: tensorflow-logging -version: 0.1.0.0 +version: 0.2.0.0 synopsis: TensorBoard related functionality. description: Please see README.md homepage: https://github.com/tensorflow/haskell#readme @@ -29,11 +29,11 @@ , stm , stm-chans , stm-conduit - , tensorflow == 0.1.* - , tensorflow-core-ops == 0.1.* - , tensorflow-ops == 0.1.* - , tensorflow-proto == 0.1.* - , tensorflow-records-conduit + , tensorflow == 0.2.* + , tensorflow-core-ops == 0.2.* + , tensorflow-ops == 0.2.* + , tensorflow-proto == 0.2.* + , tensorflow-records-conduit == 0.1.* , text , time , transformers @@ -55,6 +55,7 @@ , proto-lens , resourcet , temporary + , tensorflow , tensorflow-logging , tensorflow-proto , tensorflow-records-conduit diff --git a/tests/LoggingTest.hs b/tests/LoggingTest.hs --- a/tests/LoggingTest.hs +++ b/tests/LoggingTest.hs @@ -16,16 +16,17 @@ module Main where import Control.Monad.Trans.Resource (runResourceT) -import Data.Conduit ((=$=)) +import Data.Conduit ((.|)) import Data.Default (def) import Data.List ((\\)) -import Data.ProtoLens (decodeMessageOrDie) +import Data.ProtoLens (encodeMessage, decodeMessageOrDie) import Lens.Family2 ((^.), (.~), (&)) -import Proto.Tensorflow.Core.Util.Event (Event, fileVersion, step) +import Proto.Tensorflow.Core.Util.Event (Event, graphDef, fileVersion, step) import System.Directory (getDirectoryContents) import System.FilePath (()) import System.IO.Temp (withSystemTempDirectory) -import TensorFlow.Logging (withEventWriter, logEvent) +import TensorFlow.Core (Build, ControlNode, asGraphDef, noOp) +import TensorFlow.Logging (withEventWriter, logEvent, logGraph) import TensorFlow.Records.Conduit (sourceTFRecords) import Test.Framework (defaultMain, Test) import Test.Framework.Providers.HUnit (testCase) @@ -52,13 +53,29 @@ files <- listDirectory dir assertEqual "One file exists after" 1 (length files) records <- runResourceT $ Conduit.runConduit $ - sourceTFRecords (dir head files) =$= Conduit.consume + sourceTFRecords (dir head files) .| Conduit.consume assertBool "File is not empty" (not (null records)) let (header:body) = decodeMessageOrDie . BL.toStrict <$> records assertEqual "Header has expected version" (T.pack "brain.Event:2") (header ^. fileVersion) assertEqual "Body has expected records" expected body +testLogGraph :: Test +testLogGraph = testCase "LogGraph" $ + withSystemTempDirectory "event_writer_logs" $ \dir -> do + let graphBuild = noOp :: Build ControlNode + expectedGraph = asGraphDef graphBuild + expectedGraphEvent = (def :: Event) & graphDef .~ (encodeMessage expectedGraph) + + withEventWriter dir $ \eventWriter -> + logGraph eventWriter graphBuild + files <- listDirectory dir + records <- runResourceT $ Conduit.runConduit $ + sourceTFRecords (dir head files) .| Conduit.consume + let (_:event:_) = decodeMessageOrDie . BL.toStrict <$> records + assertEqual "First record expected to be Event containing GraphDef" expectedGraphEvent event + main :: IO () main = defaultMain [ testEventWriter + , testLogGraph ]