179 lines
7.2 KiB
Haskell
179 lines
7.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 FlexibleContexts #-}
|
|
{-# LANGUAGE RankNTypes #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
|
|
-- | Tests for EmbeddingOps.
|
|
module Main where
|
|
|
|
import Control.Monad.IO.Class (liftIO)
|
|
import Data.Int (Int32, Int64)
|
|
import Data.List (genericLength)
|
|
import TensorFlow.EmbeddingOps (embeddingLookup)
|
|
import Test.Framework (defaultMain, Test)
|
|
import Test.Framework.Providers.QuickCheck2 (testProperty)
|
|
import Test.HUnit ((@=?))
|
|
import Test.Framework.Providers.HUnit (testCase)
|
|
import Test.QuickCheck (Arbitrary(..), Property, choose, vectorOf)
|
|
import Test.QuickCheck.Monadic (monadicIO, run)
|
|
import TensorFlow.Test (assertAllClose)
|
|
|
|
import qualified Data.Vector as V
|
|
import qualified TensorFlow.GenOps.Core as CoreOps
|
|
import qualified TensorFlow.Ops as TF
|
|
import qualified TensorFlow.Session as TF
|
|
import qualified TensorFlow.Tensor as TF
|
|
import qualified TensorFlow.Types as TF
|
|
import qualified TensorFlow.Gradient as TF
|
|
import qualified TensorFlow.Build as TF
|
|
|
|
|
|
-- | Tries to perform a simple embedding lookup, with two partitions.
|
|
testEmbeddingLookupHasRightShapeWithPartition :: Test
|
|
testEmbeddingLookupHasRightShapeWithPartition =
|
|
testCase "testEmbeddingLookupHasRightShapeWithPartition" $ do
|
|
let embShape = TF.Shape [1, 3] -- Consider a 3-dim embedding of two items.
|
|
let embedding1 = [1, 1, 1 :: Int32]
|
|
let embedding2 = [0, 0, 0 :: Int32]
|
|
|
|
let idValues = [0, 1 :: Int32]
|
|
|
|
(values, shape) <- TF.runSession $ do
|
|
embedding <- mapM TF.render [ TF.constant embShape embedding1
|
|
, TF.constant embShape embedding2
|
|
]
|
|
let ids = TF.constant (TF.Shape [1, 2]) idValues
|
|
vs <- embeddingLookup embedding ids
|
|
TF.run (vs, TF.shape vs)
|
|
|
|
-- This is the shape that is returned in the equiv. Python.
|
|
shape @=? V.fromList [1, 2, 3]
|
|
|
|
-- "[0, 1]" should pull out the resulting vector.
|
|
values @=? V.fromList [1, 1, 1, 0, 0, 0]
|
|
|
|
|
|
-- | Tries to perform a simple embedding lookup, with only a single partition.
|
|
testEmbeddingLookupHasRightShape :: Test
|
|
testEmbeddingLookupHasRightShape =
|
|
testCase "testEmbeddingLookupHasRightShape" $ do
|
|
-- Consider a 3-dim embedding of two items
|
|
let embShape = TF.Shape [2, 3]
|
|
let embeddingInit = [ 1, 1, 1
|
|
, 0, 0, 0 :: Int32
|
|
]
|
|
|
|
let idValues = [0, 1 :: Int32]
|
|
|
|
(values, shape) <- TF.runSession $ do
|
|
embedding <- TF.render $ TF.constant embShape embeddingInit
|
|
let ids = TF.constant (TF.Shape [1, 2]) idValues
|
|
vs <- embeddingLookup [embedding] ids
|
|
TF.run (vs, TF.shape vs)
|
|
|
|
-- This is the shape that is returned in the equiv. Python.
|
|
shape @=? V.fromList [1, 2, 3]
|
|
|
|
-- "[0, 1]" should pull out the resulting vector.
|
|
values @=? V.fromList [1, 1, 1, 0, 0, 0]
|
|
|
|
-- | Check that we can calculate gradients w.r.t embeddings.
|
|
testEmbeddingLookupGradients :: Test
|
|
testEmbeddingLookupGradients = testCase "testEmbeddingLookupGradients" $ do
|
|
-- Agrees with "embedding", so gradient should be zero.
|
|
let xVals = V.fromList ([20, 20 :: Float])
|
|
let shape = TF.Shape [2]
|
|
|
|
gs <- TF.runSession $ do
|
|
let embShape = TF.Shape [2, 1]
|
|
let embeddingInit = [1, 20 ::Float]
|
|
let idValues = [1, 1 :: Int32]
|
|
let ids = TF.constant (TF.Shape [1, 2]) idValues
|
|
|
|
x <- TF.placeholder (TF.Shape [2])
|
|
embedding <- TF.initializedVariable
|
|
(TF.constant embShape embeddingInit)
|
|
|
|
op <- embeddingLookup [embedding] ids
|
|
let twoNorm = CoreOps.square $ TF.abs (op `TF.sub` x)
|
|
loss = TF.mean twoNorm (TF.scalar (0 :: Int32))
|
|
|
|
grad <- fmap head (TF.gradients loss [embedding])
|
|
TF.runWithFeeds
|
|
[TF.feed x $ TF.encodeTensorData shape xVals]
|
|
grad
|
|
-- Gradients should be zero (or close)
|
|
assertAllClose gs (V.fromList ([0, 0 :: Float]))
|
|
|
|
|
|
-- Verifies that direct gather is the same as dynamic split into
|
|
-- partitions, followed by embedding lookup.
|
|
testEmbeddingLookupUndoesSplit ::
|
|
forall a. (TF.TensorDataType V.Vector a, Show a, Eq a)
|
|
=> LookupExample a -> Property
|
|
testEmbeddingLookupUndoesSplit
|
|
(LookupExample numParts
|
|
shape@(TF.Shape (firstDim : restDims))
|
|
values
|
|
indices) = monadicIO $ run $ TF.runSession $ do
|
|
let shapedValues = TF.constant shape values
|
|
indicesVector <- TF.render $ TF.vector indices
|
|
let directs = CoreOps.gather shapedValues indicesVector
|
|
let cyclicCounter :: TF.Tensor TF.Build Int32 =
|
|
TF.vector [0..fromIntegral firstDim-1]
|
|
`CoreOps.mod` fromIntegral numParts
|
|
modShardedValues :: [TF.Tensor TF.Value a] <-
|
|
mapM TF.render $ CoreOps.dynamicPartition numParts shapedValues cyclicCounter
|
|
embeddings <- embeddingLookup modShardedValues indicesVector
|
|
(shapeOut, got, want :: V.Vector a) <-
|
|
TF.run (TF.cast (TF.shape embeddings), embeddings, directs)
|
|
-- Checks the explicitly documented invariant of embeddingLookup.
|
|
liftIO $ shapeOut @=? V.fromList (genericLength indices : restDims)
|
|
liftIO $ got @=? want
|
|
testEmbeddingLookupUndoesSplit _ = error "Bug in Arbitrary (LookupExample)"
|
|
|
|
-- | Consistent set of parameters for EmbeddingLookupUndoesSplit.
|
|
data LookupExample a = LookupExample
|
|
Int64 -- ^ number of ways to split.
|
|
TF.Shape -- ^ shape of the generated tensor
|
|
[a] -- ^ data for the tensor
|
|
[Int32] -- ^ indices to split the tensor by
|
|
deriving Show
|
|
|
|
instance Arbitrary a => Arbitrary (LookupExample a) where
|
|
arbitrary = do
|
|
rank <- choose (1, 4)
|
|
-- Takes rank-th root of 100 to cap the tensor size.
|
|
let maxDim = fromIntegral (ceiling doubleMaxDim :: Int64)
|
|
doubleMaxDim :: Double
|
|
doubleMaxDim = 100 ** (1 / fromIntegral rank)
|
|
shape <- vectorOf rank (choose (1, maxDim))
|
|
let firstDim = head shape
|
|
values <- vectorOf (fromIntegral $ product shape) arbitrary
|
|
numParts <- choose (2, 15)
|
|
indSize <- choose (0, fromIntegral $ firstDim - 1)
|
|
indices <- vectorOf indSize (choose (0, fromIntegral firstDim - 1))
|
|
return $ LookupExample numParts (TF.Shape shape) values indices
|
|
|
|
main :: IO ()
|
|
main = defaultMain
|
|
[ testProperty "EmbeddingLookupUndoesSplit"
|
|
(testEmbeddingLookupUndoesSplit :: LookupExample Double -> Property)
|
|
, testEmbeddingLookupHasRightShape
|
|
, testEmbeddingLookupHasRightShapeWithPartition
|
|
, testEmbeddingLookupGradients
|
|
]
|