diff --git a/src/TensorFlow/Gradient.hs b/src/TensorFlow/Gradient.hs --- a/src/TensorFlow/Gradient.hs +++ b/src/TensorFlow/Gradient.hs @@ -20,6 +20,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TypeApplications #-} module TensorFlow.Gradient ( GradientCompatible @@ -30,11 +31,12 @@ import Control.Monad.State.Strict (State, evalState, gets, modify) import Data.ByteString (ByteString) import Data.Complex (Complex) -import Data.Default (def) +import Data.ProtoLens.Default(def) import Data.Int (Int32, Int64) import Data.Foldable (foldlM) import Data.List (foldl', sortBy) import Data.Map.Strict (Map) +import qualified Data.IntSet as IntSet import Data.Maybe (fromMaybe, maybeToList, mapMaybe) import Data.Ord (comparing) import Data.ProtoLens.TextFormat (showMessage) @@ -45,7 +47,7 @@ import Lens.Family2.State.Strict (uses) import Lens.Family2.Stock (at, intAt) import Lens.Family2.Unchecked (lens, iso) -import Prelude hiding (sum) +import Prelude hiding (sum, tanh) import Text.Printf (printf) import qualified Data.Graph.Inductive.Basic as FGL import qualified Data.Graph.Inductive.Graph as FGL @@ -76,11 +78,15 @@ , matMul' , reducedShape , reluGrad + , tanh + , tanhGrad , reshape , scalar , shape , softmaxCrossEntropyWithLogits , sum + , sigmoid + , sigmoidGrad , scalarize , vector , zerosLike @@ -103,8 +109,9 @@ , ToTensor(..) ) import TensorFlow.Types (Attribute, OneOf, TensorType, attrLens) -import Proto.Tensorflow.Core.Framework.NodeDef - (NodeDef, attr, input, op, name) +import Proto.Tensorflow.Core.Framework.NodeDef (NodeDef) +import Proto.Tensorflow.Core.Framework.NodeDef_Fields + ( attr, input, op, name) type GradientCompatible a = -- TODO(fmayle): MaxPoolGrad doesn't support Double for some reason. @@ -161,6 +168,11 @@ (\f x -> fromMaybe (error $ "no NodeDef found for " ++ show x) (f x)) . flip Map.lookup let (gr, nodeMap) = createGraph yName nodeDefLookup + xnodes = mapMaybe (\x -> nodeMap ^. (at . outputNodeName . renderedOutput $ x)) xs + -- make a set of the nodes reachable from the xnodes + -- The xnodes are not part of this set (unless reachable from another xnode) + reachableSet = computeReachableSet xnodes gr + -- Set gradient of y to one. -- TODO: nicer let initPending :: Map.Map FGL.Node (PendingGradients a) @@ -171,7 +183,7 @@ .~ [yOne] ) -- Calculate the gradients of y w.r.t. each node in the graph. - gradientMap <- graphGrads gr initPending + gradientMap <- graphGrads gr reachableSet initPending -- Lookup the gradients for each x. forM xs $ \x -> let Output i xName = renderedOutput x @@ -179,6 +191,13 @@ n <- nodeMap ^. at xName gradientMap ^. at n . nonEmpty . outputIxAt i +-- | Compute a set of nodes reachable from the start nodes +-- +-- the start nodes are excluded, unless reachable from another start node +computeReachableSet :: [FGL.Node] -> Graph -> IntSet.IntSet +computeReachableSet vs g = + IntSet.fromList $ concatMap (drop 1 . FGL.preorder) (FGL.dff vs g) + outputIxAt :: OutputIx -> Lens' (IntMap.IntMap v) (Maybe v) outputIxAt = intAt . unOutputIx @@ -241,16 +260,15 @@ -- | Calculate the gradients for every node in a graph. graphGrads :: forall a. GradientCompatible a => Graph + -> IntSet.IntSet -> Map FGL.Node (PendingGradients a) -- ^ Initial gradients (usually just 1 for the node of interest). -> Build (Map FGL.Node (Gradients a)) -graphGrads gr initPending = view gradientsResult <$> foldlM go initState nodeOrder +graphGrads gr reachableSet initPending = view gradientsResult <$> foldlM go initState nodeOrder where initState = GradientsState initPending Map.empty -- Reverse topological sort. - -- TODO(fmayle): Filter out nodes that are not successors of any x in xs to - -- avoid calculating gradients that won't be used. - nodeOrder = FGL.topsort $ FGL.grev gr + nodeOrder = FGL.topsort . FGL.grev $ gr go :: GradientsState a -> Int -> Build (GradientsState a) go state node = do -- Aggregate the accumulated gradients for this node. @@ -259,11 +277,17 @@ if null outputGrads then pure state else do - let ctx = FGL.context gr node - inputGrads <- calculateInputGrads ctx outputGrads gr - -- Calculate the gradients for each of the node's inputs. let nextState = state & gradientsResult %~ Map.insert node outputGrads - pure $ updatePendingGradients ctx inputGrads nextState + -- Only consider nodes that are reachable from the inputs to + -- avoid calculating gradients that won't be used. + if node `IntSet.member` reachableSet + then do + let ctx = FGL.context gr node + inputGrads <- calculateInputGrads ctx outputGrads gr + -- Calculate the gradients for each of the node's inputs. + pure $ updatePendingGradients ctx inputGrads nextState + else + pure nextState -- | Reduce accumulated gradients for each output to one Tensor. sumPendingGradient :: GradientCompatible a @@ -458,6 +482,8 @@ opGrad "Neg" _ [_] [dz] = [Just $ negate $ expr dz] opGrad "Relu" _ [toT -> x] [dz] = [Just $ reluGrad dz x] opGrad "ReluGrad" _ [_, toT -> x ] [dz] = [Just $ reluGrad dz x, Just $ CoreOps.zerosLike x] +opGrad "Tanh" _ [toT -> x] [dz] = [Just $ tanhGrad (tanh x) dz] +opGrad "Sigmoid" _ [toT -> x] [dz] = [Just $ sigmoidGrad (sigmoid x) dz] opGrad "Concat" _ _ix [dy] -- Concat concatenates input tensors @@ -551,7 +577,7 @@ grad = reshape dz outputShapeKeptDims opGrad "Mean" u v@[toT -> x, _] w = - [Just $ dz `CoreOps.div` CoreOps.cast factor, Nothing] + [Just $ dz `CoreOps.div` (CoreOps.stopGradient $ CoreOps.cast $ factor), Nothing] where [Just dz, Nothing] = opGrad "Sum" u v w inputShape = shape (x :: Tensor Build a) @@ -624,6 +650,25 @@ [ Just $ matMul' (transAttrs True True) y dz , Just $ matMul' (transAttrs True True) dz x] +opGrad "BatchMatMul" nodeDef [toT -> x, toT -> y] [dz] = + let adjX = lookupAttr nodeDef "adj_x" + adjY = lookupAttr nodeDef "adj_y" + adjAttrs a b = + (opAttr "adj_x" .~ a) . (opAttr "adj_y" .~ b) + in case (adjX, adjY) of + (False, False) -> + [ Just $ CoreOps.batchMatMul' (adjAttrs False True) dz y + , Just $ CoreOps.batchMatMul' (adjAttrs True False) x dz] + (False, True) -> + [ Just $ CoreOps.batchMatMul dz y + , Just $ CoreOps.batchMatMul' (adjAttrs True False) dz x] + (True, False) -> + [ Just $ CoreOps.batchMatMul' (adjAttrs False True) y dz + , Just $ CoreOps.batchMatMul x dz] + (True, True) -> + [ Just $ CoreOps.batchMatMul' (adjAttrs True True) y dz + , Just $ CoreOps.batchMatMul' (adjAttrs True True) dz x] + opGrad "Transpose" _ [_, toT -> p] [dz] = [ Just $ CoreOps.transpose dz (CoreOps.invertPermutation p :: Tensor Build Int32) @@ -671,6 +716,41 @@ useCudnnOnGpu = lookupAttr nodeDef "use_cudnn_on_gpu" :: Bool dataFormat = lookupAttr nodeDef "data_format" :: ByteString +opGrad "DepthwiseConv2dNative" nodeDef [toT -> x, toT -> y] [dz] = + [ Just $ CoreOps.depthwiseConv2dNativeBackpropInput' + ((opAttr "strides" .~ strides) + . (opAttr "padding" .~ padding) + . (opAttr "data_format" .~ dataFormat)) + (shape x) y dz + , Just $ CoreOps.depthwiseConv2dNativeBackpropFilter' + ((opAttr "strides" .~ strides) + . (opAttr "padding" .~ padding) + . (opAttr "data_format" .~ dataFormat)) + x (shape y) dz + ] + where + strides = lookupAttr nodeDef "strides" :: [Int64] + padding = lookupAttr nodeDef "padding" :: ByteString + dataFormat = lookupAttr nodeDef "data_format" :: ByteString + +opGrad "DepthwiseConv2dNativeBackpropInput" nodeDef [_, toT -> x, toT -> y] [dz] = + [ Nothing + , Just $ CoreOps.depthwiseConv2dNativeBackpropFilter' + ((opAttr "strides" .~ strides) + . (opAttr "padding" .~ padding) + . (opAttr "data_format" .~ dataFormat)) + dz (shape x) y + , Just $ CoreOps.depthwiseConv2dNative' + ((opAttr "strides" .~ strides) + . (opAttr "padding" .~ padding) + . (opAttr "data_format" .~ dataFormat)) + dz x + ] + where + strides = lookupAttr nodeDef "strides" :: [Int64] + padding = lookupAttr nodeDef "padding" :: ByteString + dataFormat = lookupAttr nodeDef "data_format" :: ByteString + opGrad "MaxPool" nodeDef [toT -> x] [dz] = [ Just $ CoreOps.maxPoolGrad' ((opAttr "ksize" .~ ksize) @@ -687,9 +767,53 @@ padding = lookupAttr nodeDef "padding" :: ByteString dataFormat = lookupAttr nodeDef "data_format" :: ByteString -opGrad "Reshape" _ [toT -> x, _] [dz] = - [Just $ reshape dz $ shape (x :: Tensor Build a), Nothing] +opGrad "Reshape" _ [toT -> x, _] [dz] = [Just $ reshape dz $ shape (x :: Tensor Build a), Nothing] +opGrad "ExpandDims" n xs@[toT -> _, _] dzs@[_] = opGrad "Reshape" n xs dzs +opGrad "Squeeze" _ [toT -> x] [dz] = [Just $ reshape dz $ shape (x :: Tensor Build a)] +opGrad "Pad" _ [toT -> x, toT -> padPattern] [dz] = + [Just $ CoreOps.slice dz gradientSliceBegin gradientSliceSize, Nothing] + where + v1 = vector [1] + -- For some reason rankx' has an empty shape + rankx' = CoreOps.rank (x :: Tensor Build Float) + rankx = CoreOps.reshape rankx' v1 + -- Size of column that is sliced from pad pattern + padPatternSliceSize = CoreOps.concat 0 [rankx, v1] + padPatternSliceBegin = vector [0, 0] + padPatternSliced :: Tensor Build Int32 = CoreOps.slice padPattern padPatternSliceBegin padPatternSliceSize + -- The slice of the pad pattern has the same rank as the pad pattern itself + gradientSliceBegin = CoreOps.reshape padPatternSliced rankx + gradientSliceSize = shape (x :: Tensor Build Float) +-- Gradient for Slice +-- Create an Nx2 padding where N is the rank of (grad of) Slice and the first +-- column represents how many zeros are to be prepended for each dimension, and the second +-- column indicates how many zeros are appended. +-- The number of zeros to prepend is the shape of the beginvec. +-- The number of zeros to append is the shape of the inputvec +-- elementwise-subtracted by both the beginvec and sizevec. +-- Some more reshaping is needed to assemble this tensor with the +-- right dimensions. +opGrad "Slice" _ [toT -> inputvec, toT -> beginvec, _] [dz] = + [Just $ CoreOps.pad dz paddings, Nothing, Nothing] + where + v1 = vector [1 :: Int32] + inputRank' = CoreOps.rank (inputvec :: Tensor Build Float) + -- For some reason inputRank' has an empty shape + inputRank = CoreOps.reshape inputRank' v1 + padShape = CoreOps.concat 0 [inputRank, v1] + beforePad = CoreOps.reshape beginvec padShape + afterPad = CoreOps.reshape (shape inputvec - shape dz - beginvec) padShape + paddings = CoreOps.concat 1 [beforePad, afterPad] + +-- TODO: This could be either Int32 or Int64. +opGrad "BatchToSpaceND" _ [_, toT @Int32 -> blockShape, toT @Int32 -> crops] [dz] = + [Just $ CoreOps.spaceToBatchND dz blockShape crops, Nothing, Nothing] + +-- TODO: This could be either Int32 or Int64. +opGrad "SpaceToBatchND" _ [_, toT @Int32 -> blockShape, toT @Int32 -> paddings] [dz] = + [Just $ CoreOps.batchToSpaceND dz blockShape paddings, Nothing, Nothing] + opGrad "OneHot" _ _ _ = [Nothing, Nothing, Nothing, Nothing] opGrad "TruncatedNormal" _ _ _ = [Nothing] @@ -768,6 +892,17 @@ axes = CoreOps.range 0 (CoreOps.size splitShape) (2 :: Tensor Build Int32) reshapedDz = CoreOps.reshape dz splitShape +opGrad "ResizeBilinear" nodeDef [toT -> x, _] [dz] = + [ Just $ CoreOps.resizeBilinearGrad' + (opAttr "align_corners" .~ align) + (CoreOps.cast dz) + x + + , Nothing + ] + where + align = lookupAttr nodeDef "align_corners" :: Bool + opGrad "ZerosLike" _ _ _ = [Nothing] opGrad "Fill" _ _ [dz] = [Nothing, Just $ sum dz rx] where @@ -779,12 +914,14 @@ -- through each read. opGrad "ReadVariableOp" _ _ [dz] = [Just $ expr dz] --- TODO(fmayle): These can go away if we properly prune the graph. opGrad "Const" _ _ _ = [Nothing, Nothing] -opGrad "Placeholder" _ _ _ = [] +opGrad "StopGradient" _ _ _ = [Nothing] opGrad "VarHandleOp" _ _ _ = [] -opGrad "Variable" _ _ _ = [] +opGrad "Sqrt" _ [toT -> x] [dz] = [Just $ sq' `CoreOps.mul` dz] + where + sq' = scalar 1 `CoreOps.div` (scalar 2 `CoreOps.mul` CoreOps.sqrt x) + opGrad n nodeDef ins grads = error $ "no gradient implemented for " ++ show (n, length ins, length grads, showMessage nodeDef, ins) @@ -796,16 +933,21 @@ "Abs" -> 1 "Add" -> 1 "AddN" -> 1 + "BatchToSpaceND" -> 1 + "BatchMatMul" -> 1 "Cast" -> 1 "Const" -> 1 "Concat" -> 1 "Conv2D" -> 1 "Conv2DBackpropInput" -> 1 + "DepthwiseConv2dNative" -> 1 + "DepthwiseConv2dNativeBackpropInput" -> 1 "Div" -> 1 "DynamicStitch" -> 1 "DynamicPartition" -> fromIntegral (lookupAttr o "num_partitions" :: Int64) "Exp" -> 1 + "ExpandDims" -> 1 "Gather" -> 1 "LabelClasses" -> 1 "LabelWeights" -> 1 @@ -818,7 +960,9 @@ "Min" -> 1 "Mul" -> 1 "Neg" -> 1 + "Pad" -> 1 "Placeholder" -> 1 + "StopGradient" -> 1 "OneHot" -> 1 "ReadVariableOp" -> 1 "RefIdentity" -> 1 @@ -826,13 +970,20 @@ "ReluGrad" -> 1 "Reshape" -> 1 "Select" -> 1 + "Sigmoid" -> 1 "Size" -> 1 + "Slice" -> 1 "SoftmaxCrossEntropyWithLogits" -> 2 - "Square" -> 1 + "SpaceToBatchND" -> 1 "SparseSegmentSum" -> 1 + "Square" -> 1 + "Squeeze" -> 1 + "Sqrt" -> 1 "Sub" -> 1 "Sum" -> 1 + "Tanh" -> 1 "Tile" -> 1 + "ResizeBilinear" -> 1 "Transpose" -> 1 "TruncatedNormal" -> 1 "VarHandleOp" -> 1 diff --git a/src/TensorFlow/Ops.hs b/src/TensorFlow/Ops.hs --- a/src/TensorFlow/Ops.hs +++ b/src/TensorFlow/Ops.hs @@ -112,6 +112,8 @@ , CoreOps.relu' , CoreOps.reluGrad , CoreOps.reluGrad' + , CoreOps.tanh + , CoreOps.tanhGrad , CoreOps.reshape , CoreOps.reshape' , restore @@ -121,6 +123,8 @@ , scalar' , shape , shape' + , CoreOps.sigmoid + , CoreOps.sigmoidGrad , CoreOps.sign , CoreOps.sign' , CoreOps.size @@ -156,17 +160,18 @@ import Data.Int (Int32, Int64) import Data.Word (Word16) import Prelude hiding (abs, sum, concat) -import Data.ProtoLens (def) +import Data.ProtoLens.Default(def) import Data.Text.Encoding (encodeUtf8) import Lens.Family2 ((.~), (&)) import Text.Printf (printf) -import Proto.Tensorflow.Core.Framework.Tensor - ( TensorProto - , dtype +import Proto.Tensorflow.Core.Framework.Tensor (TensorProto) +import Proto.Tensorflow.Core.Framework.Tensor_Fields + ( dtype , tensorShape ) -import qualified Proto.Tensorflow.Core.Framework.TensorShape +import qualified Proto.Tensorflow.Core.Framework.TensorShape_Fields as TensorShape + import TensorFlow.Build import TensorFlow.BuildOp import TensorFlow.ControlFlow (group) diff --git a/src/TensorFlow/Variable.hs b/src/TensorFlow/Variable.hs --- a/src/TensorFlow/Variable.hs +++ b/src/TensorFlow/Variable.hs @@ -11,6 +11,7 @@ {-# LANGUAGE RecursiveDo #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE NoMonadFailDesugaring #-} module TensorFlow.Variable ( Variable , variable diff --git a/tensorflow-ops.cabal b/tensorflow-ops.cabal --- a/tensorflow-ops.cabal +++ b/tensorflow-ops.cabal @@ -1,5 +1,5 @@ name: tensorflow-ops -version: 0.2.0.0 +version: 0.2.0.1 synopsis: Friendly layer around TensorFlow bindings. description: Please see README.md homepage: https://github.com/tensorflow/haskell#readme @@ -21,7 +21,7 @@ , TensorFlow.NN , TensorFlow.Queue , TensorFlow.Variable - build-depends: proto-lens == 0.2.* + build-depends: proto-lens >= 0.4.0 && < 0.6.0 , base >= 4.7 && < 5 , bytestring , fgl diff --git a/tests/BuildTest.hs b/tests/BuildTest.hs --- a/tests/BuildTest.hs +++ b/tests/BuildTest.hs @@ -21,11 +21,12 @@ import Control.Monad.IO.Class (liftIO) import Lens.Family2 ((^.), (.~)) import Data.List (sort) -import Proto.Tensorflow.Core.Framework.Graph +import Proto.Tensorflow.Core.Framework.Graph_Fields ( node ) import Proto.Tensorflow.Core.Framework.NodeDef - ( NodeDef - , device + ( NodeDef ) +import Proto.Tensorflow.Core.Framework.NodeDef_Fields + ( device , name , op ) import TensorFlow.Build diff --git a/tests/EmbeddingOpsTest.hs b/tests/EmbeddingOpsTest.hs --- a/tests/EmbeddingOpsTest.hs +++ b/tests/EmbeddingOpsTest.hs @@ -15,6 +15,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NoMonadFailDesugaring #-} -- | Tests for EmbeddingOps. module Main where diff --git a/tests/GradientTest.hs b/tests/GradientTest.hs --- a/tests/GradientTest.hs +++ b/tests/GradientTest.hs @@ -16,6 +16,7 @@ {-# LANGUAGE NoMonomorphismRestriction #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE NoMonadFailDesugaring #-} import Data.Int (Int32, Int64) import Data.List (sort) @@ -32,17 +33,18 @@ import Control.Monad.IO.Class (liftIO) import qualified TensorFlow.Core as TF -import qualified TensorFlow.GenOps.Core as TF (conv2DBackpropInput', max, maximum, tile) +import qualified TensorFlow.GenOps.Core as TF (conv2DBackpropInput', max, maximum, resizeBilinear', tile, pad, batchToSpaceND, spaceToBatchND, squeeze, sqrt, slice, shape, diag, depthwiseConv2dNative', depthwiseConv2dNativeBackpropInput', batchMatMul, batchMatMul', sum, conjugateTranspose) import qualified TensorFlow.Gradient as TF -import qualified TensorFlow.Ops as TF hiding (zeroInitializedVariable) +import qualified TensorFlow.Ops as TF hiding (zeroInitializedVariable, shape) import qualified TensorFlow.Output as TF import qualified TensorFlow.Types as TF import qualified TensorFlow.Variable as TF -import Proto.Tensorflow.Core.Framework.Graph (node) -import Proto.Tensorflow.Core.Framework.NodeDef (op) +import Proto.Tensorflow.Core.Framework.Graph_Fields (node) +import Proto.Tensorflow.Core.Framework.NodeDef_Fields (op) import qualified Data.ByteString.Char8 as BS +import TensorFlow.Session (SessionT) testGradientSimple :: Test testGradientSimple = testCase "testGradientSimple" $ do @@ -122,7 +124,66 @@ ] sort expected @=? sort ops +testGradientIncidental :: Test +testGradientIncidental = testCase "testGradientIncidental" $ do + let grads = do + x <- TF.render $ TF.scalar (3 :: Float) + b <- TF.render $ TF.scalar (4 :: Float) + w <- TF.render $ TF.diag $ TF.vector [ 1.0 :: Float ] + let incidental = b `TF.mul` w + let y = (x `TF.mul` b) `TF.add` incidental + TF.gradients y [x] + -- Assert that the gradients are right. + [dx] <- TF.runSession $ grads >>= TF.run + 4 @=? TF.unScalar dx + -- Assert that the graph has the expected ops. + let graphDef = TF.asGraphDef grads + putStrLn $ showMessage graphDef + let ops = graphDef ^.. node . traverse . op + expected = [ "Add" + , "BroadcastGradientArgs" + , "BroadcastGradientArgs" + , "Const" + , "Const" + , "Const" + , "Const" + , "Diag" + , "Fill" + , "Mul" + , "Mul" + , "Mul" + , "Mul" + , "Reshape" + , "Reshape" + , "Reshape" + , "Reshape" + , "Shape" + , "Shape" + , "Shape" + , "Shape" + , "Shape" + , "Sum" + , "Sum" + , "Sum" + , "Sum" + ] + sort expected @=? sort ops + +testGradientPruning :: Test +testGradientPruning = testCase "testGradientPruning" $ do + let grads = do + x <- TF.render $ TF.scalar (3 :: Float) + b <- TF.render $ TF.scalar (4 :: Float) + bx <- TF.render $ b `TF.mul` x + let y = bx `TF.add` b + TF.gradients y [x, bx] + + -- Assert that the gradients are right. + [dx, dxb] <- TF.runSession $ grads >>= TF.run + 4 @=? TF.unScalar dx + 1 @=? TF.unScalar dxb + -- Test that identical "stateful" ops work with createGraph. testCreateGraphStateful :: Test testCreateGraphStateful = testCase "testCreateGraphStateful" $ do @@ -169,7 +230,24 @@ TF.gradients y [x] >>= TF.run V.fromList [2, 2, 2 :: Float] @=? dx +testMeanGradient :: Test +testMeanGradient = testCase "testMeanGradient" $ do + [dx] <- TF.runSession $ do + x <- TF.render $ TF.vector [1, 2, 0 :: Float] + let y = TF.mean x (TF.vector [0 :: Int32]) + TF.gradients y [x] >>= TF.run + V.fromList [1, 1, 1 :: Float] @=? dx +testMeanGradGrad :: Test +testMeanGradGrad = testCase "testMeanGradGrad" $ do + [ddx] <- TF.runSession $ do + x <- TF.render $ TF.vector [1, 2, 0 :: Float] + let y = TF.mean x (TF.vector [0 :: Int32]) + [dx] <- TF.gradients y [x] + TF.gradients dx [x] >>= TF.run + + V.fromList [0, 0, 0 :: Float] @=? ddx + testMaxGradient :: Test testMaxGradient = testCase "testMaxGradient" $ do [dx] <- TF.runSession $ do @@ -282,6 +360,127 @@ TF.gradients y' [x] >>= TF.run V.fromList [0] @=? dx +testTanhGrad :: Test +testTanhGrad = testCase "testTanhGrad" $ do + [dx] <- TF.runSession $ do + x <- TF.render $ TF.vector [0 :: Float] + let y = TF.tanh x + TF.gradients y [x] >>= TF.run + V.fromList [1] @=? dx + +testSigmoidGrad :: Test +testSigmoidGrad = testCase "testSigmoidGrad" $ do + [dx] <- TF.runSession $ do + x <- TF.render $ TF.vector [0 :: Float] + let y = TF.sigmoid x + TF.gradients y [x] >>= TF.run + V.fromList [0.25] @=? dx + +testExpandDims :: Test +testExpandDims = + testCase "testExpandDims" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [1, 2, 3 :: Int64] + let y = TF.expandDims x $ TF.constant (TF.Shape [1]) [0 :: Int32] + calculateGradWithShape y x + V.fromList [1, 1, 1, 1, 1, 1] @=? dx + V.fromList [1, 2, 3] @=? s + +testReshape :: Test +testReshape = + testCase "testReshape" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [2, 2 :: Int64] + let y = TF.reshape x $ TF.constant (TF.Shape [2]) [1, 4 :: Int32] + calculateGradWithShape y x + V.fromList [1, 1, 1, 1] @=? dx + V.fromList [2, 2] @=? s + +testPad :: Test +testPad = + testCase "testPad" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [2, 2, 3 :: Int64] + let y = TF.pad x $ TF.constant (TF.Shape [3, 2]) [1, 4, 1, 1, 2, 3 :: Int32] + calculateGradWithShape y x + V.fromList [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] @=? dx + V.fromList [2, 2, 3] @=? s + + +testSqrt :: Test +testSqrt = testCase "testSqrt" $ do + [dx] <- TF.runSession $ do + x <- TF.render $ TF.vector [0.0625 :: Float] + let y = TF.sqrt x + TF.gradients y [x] >>= TF.run + V.fromList [2] @=? dx + +testSlice :: Test +testSlice = + testCase "testSlice" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [2, 3, 4 :: Int64] + (z :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [1, 2, 2 :: Int64] + let y = TF.slice x (TF.constant (TF.Shape [3]) [1, 1, 1 :: Int32]) (TF.shape z) + calculateGradWithShape y x + let expected = + [0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 0, 0, 0, + 0, 1, 1, 0, + 0, 1, 1, 0] + V.fromList expected @=? dx + V.fromList [2, 3, 4] @=? s + +testBatchToSpaceND :: Test +testBatchToSpaceND = + testCase "testBatchToSpaceND" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.constant (TF.Shape [4, 1, 1, 1 :: Int64]) [1, 2, 3, 4] + shape <- TF.render $ TF.vector [2, 2 :: Int32] + crops <- TF.render $ TF.constant (TF.Shape [2, 2]) [0, 0, 0, 0 :: Int32] + let y = TF.batchToSpaceND x shape crops + calculateGradWithShape y x + V.fromList [1, 1, 1, 1] @=? dx + V.fromList [4, 1, 1, 1] @=? s + +testSpaceToBatchND :: Test +testSpaceToBatchND = + testCase "testSpaceToBatchND" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.constant (TF.Shape [1, 2, 2, 1 :: Int64]) [1, 2, 3, 4] + shape <- TF.render $ TF.vector [2, 2 :: Int32] + paddings <- TF.render $ TF.constant (TF.Shape [2, 2]) [0, 0, 0, 0 :: Int32] + let y = TF.spaceToBatchND x shape paddings + calculateGradWithShape y x + V.fromList [1, 1, 1, 1] @=? dx + V.fromList [1, 2, 2, 1] @=? s + +testSqueeze :: Test +testSqueeze = + testCase "testSqueeze" $ do + ([dx], [s]) <- + TF.runSession $ do + (x :: TF.Tensor TF.Value Float) <- TF.render $ TF.zeros $ TF.Shape [1, 2, 3 :: Int64] + let y = TF.squeeze x + calculateGradWithShape y x + V.fromList [1, 1, 1, 1, 1, 1] @=? dx + V.fromList [1, 2, 3] @=? s + +calculateGradWithShape :: TF.Tensor TF.Build Float -> TF.Tensor TF.Value Float -> SessionT IO ([V.Vector Float], [V.Vector Int32]) +calculateGradWithShape y x = do + gs <- TF.gradients y [x] + xs <- TF.run gs + (shapes :: [V.Vector Int32]) <- mapM (TF.run . TF.shape) gs + return (xs, shapes) + testFillGrad :: Test testFillGrad = testCase "testFillGrad" $ do [dx] <- TF.runSession $ do @@ -315,6 +514,22 @@ shapeX @=? (shapeDX :: V.Vector Int32) V.fromList [6, 6, 6, 6, 6, 6::Float] @=? (dx :: V.Vector Float) +testResizeBilinearGrad :: Test +testResizeBilinearGrad = testCase "testResizeBilinearGrad" $ do + (dx, shapeDX, shapeX) <- TF.runSession $ do + let shape = TF.vector [1, 2, 2, 1 :: Int32] + x <- TF.render $ TF.fill shape (TF.scalar (1 :: Float)) + let outSize = TF.vector [4, 4 :: Int32] + align = TF.opAttr "align_corners" .~ True + y = TF.resizeBilinear' align x outSize + + [dx] <- TF.gradients y [x] + TF.run (dx, TF.shape dx, TF.shape x) + shapeX @=? (shapeDX :: V.Vector Int32) + let expect = V.fromList [4, 4, 4, 4 :: Float] + near = 0.00001 > (V.sum $ V.zipWith (-) expect (dx :: V.Vector Float)) + near @=? True + matMulGradient :: Test matMulGradient = testCase "matMulGradients" $ do @@ -389,6 +604,84 @@ transAttrs a b = (TF.opAttr "transpose_a" .~ a) . (TF.opAttr "transpose_b" .~ b) + +batchMatMulGradient :: Test +batchMatMulGradient = testCase "batchMatMulGradients" $ do + + let dfBuild = do + x <- TF.render $ TF.zeros $ TF.Shape [2,3, 1 :: Int64] + w <- TF.zeroInitializedVariable $ TF.Shape [2,1, 2 :: Int64] + let f = x `TF.batchMatMul` TF.readValue w :: TF.Tensor TF.Build Float + dfs <- TF.gradients f [x] + return (x, dfs) + + (xShape, dxShape) <- TF.runSession $ do + (x, [dx]) <- TF.build dfBuild + TF.run (TF.shape x, TF.shape dx) + + assertEqual "Shape of gradient must match shape of input" xShape (dxShape :: V.Vector Int32) + + +-- test that gradient of batchMatMul can be taken gradient of +batchMatMulGradGrad :: Test +batchMatMulGradGrad = testCase "batchMatMulGradGrad" $ do + let width = 2 :: Int64 + height = 3 :: Int64 + batch = 4 :: Int64 + + let tower = do + x <- TF.render $ TF.zeros $ TF.Shape [batch, height, 1] + w <- TF.zeroInitializedVariable $ TF.Shape [batch, 1, width] + let f = x `TF.batchMatMul` TF.readValue w + [dfdx] <- TF.gradients f [x] + let f'x = TF.sum dfdx (TF.vector [1, 2 :: Int32]) + [dfdw] <- TF.gradients f'x [w] -- take gradient again (this time over w) + return [TF.readValue w, TF.expr dfdw] + + TF.runSession $ do + [w, dfdw] <- TF.build tower + (wShape, dfdwShape) <- TF.run (TF.shape w, TF.shape dfdw) + liftIO $ assertEqual "Shape of gradient must match input" wShape (dfdwShape :: V.Vector Int32) + + let step = w `TF.add` dfdw + w0 <- TF.run step + liftIO $ V.fromList [3.0,3.0,3.0,3.0,3.0,3.0,3.0,3.0 :: Float] @=? w0 + + +-- test that gradient of batchMatMul deals correctly with adj_x and adj_y +batchMatMulAdjointGradient :: (Bool, Bool) -> Test +batchMatMulAdjointGradient axw = testCase ("batchMatMulAdjointGradients " ++ show axw) $ do + let (adjX, adjW) = axw + + let dfBuild = do + let xShape = TF.Shape [2, 3, 1 :: Int64] + let xZeros = TF.zeros xShape + x <- TF.render $ if adjX then TF.conjugateTranspose xZeros (TF.vector [0, 2, 1 :: Int32]) else xZeros + variable <- TF.zeroInitializedVariable $ TF.Shape [2, 1, 2 :: Int64] + let wv = if adjW then TF.conjugateTranspose (TF.readValue variable) (TF.vector [0, 2, 1 :: Int32]) else TF.readValue variable + let f = TF.batchMatMul' (adjAttrs adjX adjW) x wv :: TF.Tensor TF.Build Float + w <- TF.render wv + ds <- TF.gradients f [x, w] + return (x, w, ds) + + TF.runSession $ do + (x, w, [dx, dw]) <- TF.build dfBuild + xShape <- TF.run $ TF.shape x + dxShape <- TF.run $ TF.shape dx + liftIO $ assertEqual "xShape must match dxShape" xShape (dxShape :: V.Vector Int32) + + wShape <- TF.run $ TF.shape w + dwShape <- TF.run $ TF.shape dw + liftIO $ assertEqual "wShape must match dwShape" wShape (dwShape :: V.Vector Int32) + +adjAttrs :: (TF.Attribute x, + TF.Attribute y) => + x -> y -> TF.OpDef -> TF.OpDef +adjAttrs x y = + (TF.opAttr "adj_x" .~ x) . (TF.opAttr "adj_y" .~ y) + + +-- TODO check gradient with regard to filter also testConv2DBackpropInputGrad :: Test testConv2DBackpropInputGrad = testCase "testConv2DBackpropInputGrad" $ do (dx, shapeDX, shapeX) <- TF.runSession $ do @@ -410,15 +703,60 @@ shapeX @=? (shapeDX :: V.Vector Int32) V.fromList [4::Float] @=? (dx :: V.Vector Float) +testDepthwiseConv2dGrad :: Test +testDepthwiseConv2dGrad = testCase "testDepthwiseConv2dGrad" $ do + (dx, shapeDX, shapeX) <- TF.runSession $ do + let conv_input_shape = TF.vector [1, 2, 2, 1 :: Int32] + x <- TF.render $ TF.fill conv_input_shape (TF.scalar (2 :: Float)) + let filterShape = TF.vector [2, 2, 1, 1 :: Int32] + filter' <- TF.render $ TF.fill filterShape (TF.scalar (1 :: Float)) + let y = TF.depthwiseConv2dNative' + ( (TF.opAttr "strides" .~ [1 :: Int64, 1, 1, 1]) + . (TF.opAttr "padding" .~ (BS.pack "VALID")) + . (TF.opAttr "data_format" .~ (BS.pack "NHWC")) + ) + x filter' + + [dx] <- TF.gradients y [x] + TF.run (dx, TF.shape dx, TF.shape x) + shapeX @=? (shapeDX :: V.Vector Int32) + V.fromList [1, 1, 1, 1 :: Float] @=? (dx :: V.Vector Float) + +-- TODO also test filter gradient +testDepthwiseConv2dBackpropInputGrad :: Test +testDepthwiseConv2dBackpropInputGrad = testCase "testDepthwiseConv2dBackpropInputGrad" $ do + (dx, shapeDX, shapeX) <- TF.runSession $ do + let conv_input_shape = TF.vector [1, 2, 2, 1 :: Int32] + let conv_out_shape = TF.vector [1, 1, 1, 1 :: Int32] -- [batch, h, w, out_channels] + x <- TF.render $ TF.fill conv_out_shape (TF.scalar (1::Float)) + + let filterShape = TF.vector [2, 2, 1, 1 :: Int32] + filter' <- TF.render $ TF.fill filterShape (TF.scalar (1 :: Float)) + let y = TF.depthwiseConv2dNativeBackpropInput' + ( (TF.opAttr "strides" .~ [1 :: Int64, 1, 1, 1]) + . (TF.opAttr "padding" .~ (BS.pack "VALID")) + . (TF.opAttr "data_format" .~ (BS.pack "NHWC")) + ) + conv_input_shape filter' x + + [dx] <- TF.gradients y [x] + TF.run (dx, TF.shape dx, TF.shape x) + shapeX @=? (shapeDX :: V.Vector Int32) + V.fromList [4::Float] @=? (dx :: V.Vector Float) + main :: IO () main = defaultMain [ testGradientSimple , testGradientDisconnected + , testGradientIncidental + , testGradientPruning , testCreateGraphStateful , testCreateGraphNameScopes , testDiamond , testAddNGradient + , testMeanGradient + , testMeanGradGrad , testMaxGradient , testConcatGradient , testConcatGradientSimple @@ -427,14 +765,33 @@ , testMaximumGradGrad , testReluGrad , testReluGradGrad + , testTanhGrad + , testSigmoidGrad + , testExpandDims + , testReshape + , testPad + , testSqrt + , testSlice + , testBatchToSpaceND + , testSpaceToBatchND + , testSqueeze , testFillGrad , testTileGrad , testTile2DGrad + , testResizeBilinearGrad , matMulGradient , matMulGradGrad , matMulTransposeGradient (False, False) , matMulTransposeGradient (False, True) , matMulTransposeGradient (True, False) , matMulTransposeGradient (True, True) + , batchMatMulGradient + , batchMatMulGradGrad + , batchMatMulAdjointGradient (False, False) + , batchMatMulAdjointGradient (False, True) + , batchMatMulAdjointGradient (True, False) + , batchMatMulAdjointGradient (True, True) , testConv2DBackpropInputGrad + , testDepthwiseConv2dGrad + , testDepthwiseConv2dBackpropInputGrad ]