1
0
Fork 0
mirror of https://github.com/tensorflow/haskell.git synced 2024-12-25 19:19:45 +01:00
tensorflow-haskell/tensorflow-ops/tests/BuildTest.hs
2016-10-24 19:26:42 +00:00

181 lines
6.2 KiB
Haskell

-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
-- http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Main where
import Control.Monad.IO.Class (liftIO)
import Data.Functor.Identity (runIdentity)
import Lens.Family2 ((^.))
import Data.List (sort)
import Proto.Tensorflow.Core.Framework.Graph
( node )
import Proto.Tensorflow.Core.Framework.NodeDef
( NodeDef
, device
, name
, op )
import TensorFlow.Build
( Build
, BuildT
, asGraphDef
, evalBuildT
, flushNodeBuffer
, hoistBuildT
, render
, withDevice
, colocateWith
, withNameScope
)
import TensorFlow.ControlFlow (named)
import TensorFlow.Nodes (unScalar)
import TensorFlow.Ops
( add
, assign
, constant
, initializedVariable
, variable
)
import TensorFlow.Output (Device(..))
import TensorFlow.Tensor (Tensor, Value, Ref)
import TensorFlow.Session
( build
, buildAnd
, run
, runSession
, run_
)
import Test.Framework.Providers.HUnit (testCase)
import Test.HUnit ((@=?))
import Google.Test (googleTest)
import qualified Data.Vector as V
-- | Test named behavior.
testNamed = testCase "testNamed" $ do
let graph = named "foo" <$> variable [] >>= render :: Build (Tensor Ref Float)
nodeDef :: NodeDef
nodeDef = head $ asGraphDef graph ^. node
"RefIdentity" @=? (nodeDef ^. op)
"foo" @=? (nodeDef ^. name)
-- | Test named deRef behavior.
testNamedDeRef = testCase "testNamedDeRef" $ do
let graph = named "foo" <$> do
v :: Tensor Ref Float <- variable []
assign v 5
-- TODO: Implement TensorFlow get_variable and test it.
runSession $ do
out <- buildAnd run graph
liftIO $ 5 @=? (unScalar out :: Float)
-- | Test that "run" will render and extend any pure ops that haven't already
-- been rendered.
testPureRender = testCase "testPureRender" $ runSession $ do
result <- run $ 2 `add` 2
liftIO $ 4 @=? (unScalar result :: Float)
-- | Test that "run" assigns any previously accumulated initializers.
testInitializedVariable =
testCase "testInitializedVariable" $ runSession $ do
(formula, reset) <- build $ do
v <- initializedVariable 42
r <- assign v 24
return (1 `add` v, r)
result <- run formula
liftIO $ 43 @=? (unScalar result :: Float)
run_ reset -- Updates v to a different value
rerunResult <- run formula
liftIO $ 25 @=? (unScalar rerunResult :: Float)
testInitializedVariableShape =
testCase "testInitializedVariableShape" $ runSession $ do
vector <- build $ initializedVariable (constant [1] [42 :: Float])
result <- run vector
liftIO $ [42] @=? (result :: V.Vector Float)
-- | Test nameScoped behavior.
testNameScoped = testCase "testNameScoped" $ do
let graph = withNameScope "foo" $ variable [] :: Build (Tensor Ref Float)
nodeDef :: NodeDef
[nodeDef] = asGraphDef graph ^. node
"foo/Variable_0" @=? (nodeDef ^. name) -- TODO: Check prefix.
"Variable" @=? (nodeDef ^. op)
-- | Test combined named and nameScoped behavior.
testNamedAndScoped = testCase "testNamedAndScoped" $ do
let graph :: Build (Tensor Ref Float)
graph = withNameScope "foo1" ((named "bar1" <$> variable []) >>= render)
nodeDef :: NodeDef
nodeDef = head $ asGraphDef graph ^. node
"RefIdentity" @=? (nodeDef ^. op)
"foo1/bar1" @=? (nodeDef ^. name)
-- | Lift a Build action into a context for HUnit to run.
liftBuild :: Build a -> BuildT IO a
liftBuild = hoistBuildT (return . runIdentity)
-- | Flush the node buffer and sort the nodes by name (for more stable tests).
flushed :: Ord a => (NodeDef -> a) -> BuildT IO [a]
flushed field = sort . map field <$> liftBuild flushNodeBuffer
-- | Test the interaction of rendering, CSE and scoping.
testRenderDedup = testCase "testRenderDedup" $ evalBuildT $ do
liftBuild renderNodes
names <- flushed (^. name)
liftIO $ ["Const_1", "Variable_0", "Variable_2"] @=? names
-- Render the nodes in a different scope, which should cause them
-- to be distinct from the previous ones.
liftBuild $ withNameScope "foo" renderNodes
scopedNames <- flushed (^. name)
liftIO $ ["foo/Const_4", "foo/Variable_3", "foo/Variable_5"] @=? scopedNames
where
renderNodes = do
-- A stateful op and a pure op.
_ :: Tensor Ref Float <- variable []
_ :: Tensor Value Float <- render 3
-- Another stateful op, and a pure op which should be
-- deduped with the previous one.
_ :: Tensor Ref Float <- variable []
_ :: Tensor Value Float <- render 3
return ()
-- | Test the interaction of rendering, CSE and scoping.
testDeviceColocation = testCase "testDeviceColocation" $ evalBuildT $ do
liftBuild renderNodes
devices <- flushed (\x -> (x ^. name, x ^. device))
liftIO $ [ ("Add_2","dev0")
, ("Const_1","dev0")
, ("Variable_0","dev0")] @=? devices
where
renderNodes = do
-- A stateful op and a pure op.
var :: Tensor Ref Float <- withDevice (Just $ Device "dev0") $ variable []
-- Uses render to cause the expression be added to the graph.
_ <- colocateWith var $ render $ 3 `add` var
return ()
main :: IO ()
main = googleTest [ testInitializedVariable
, testInitializedVariableShape
, testDeviceColocation
, testNamed
, testNamedDeRef
, testNameScoped
, testNamedAndScoped
, testPureRender
, testRenderDedup
]