diff --git a/NN/Backend/Torch/Codegen.hs b/NN/Backend/Torch/Codegen.hs new file mode 100644 --- /dev/null +++ b/NN/Backend/Torch/Codegen.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} +module NN.Backend.Torch.Codegen where + +import Control.Applicative +import Control.Lens hiding (assign) +import Control.Monad.State.Strict +import Gen.Caffe.LayerParameter as LP +import Language.Lua.PrettyPrinter +import Language.Lua.Syntax +import Text.Printf + +import NN.Backend.Torch.Lua +import NN.Backend.Torch.Torch + +data TorchState = TorchState { + _statements :: [Stat], + _sequential :: Maybe String, + _criteria :: [String], + _count :: Int + } +makeLenses ''TorchState + +newtype Torch a = Torch { _unTorch :: State TorchState a } + deriving (Functor, Applicative, Monad, MonadState TorchState) + +initialize :: Torch () +initialize = do + seq' <- fresh "seq" + sequential ?= seq' + + statements <>= [require "nn"] + statements <>= [assign seq' $ torchExp (TorchModule "nn" "Sequential" [])] + where + require module' = funCall "require" [toLua $ L module'] + +fresh :: String -> Torch String +fresh prefix = do + c <- use count + count += 1 + return $ printf "%s%d" prefix c + +insertModule :: Module Exp -> Torch () +insertModule (Criterion exp') = do + name' <- fresh "criterion" + criteria <>= [name'] + statements <>= [assign name' exp'] + +insertModule (Inner exp') = do + Just seq' <- use sequential + statements <>= [methCall seq' "add" [exp']] + +finalize :: Torch Block +finalize = do + Just seq' <- use sequential + criteria' <- use criteria + statements' <- use statements + return $ Block statements' (Just $ return' <$> seq':criteria') + +runTorch :: [LayerParameter] -> Torch Block +runTorch layers = do + initialize + forM_ exps insertModule + finalize + where + exps = concatMap torchExps layers + torchExps lp = (torchExp <$>) <$> torchModules lp + +lower :: [LayerParameter] -> Block +lower layers = (evalState . _unTorch) (runTorch layers) emptyTorch + where + emptyTorch = TorchState [] Nothing [] 0 + +codegen :: Block -> String +codegen block = pprint block & renderPretty 0.4 150 & displayS & \f -> f "" diff --git a/NN/Backend/Torch/Lua.hs b/NN/Backend/Torch/Lua.hs new file mode 100644 --- /dev/null +++ b/NN/Backend/Torch/Lua.hs @@ -0,0 +1,39 @@ +module NN.Backend.Torch.Lua where + +import Data.Word +import Language.Lua.Syntax + +newtype LS = L String + +-- |Handy typeclass for converting arguments +class ToLua a where + toLua :: a -> Exp + +instance ToLua Word32 where + toLua = Number . show + +instance ToLua LS where + toLua (L s') = String s' + +instance ToLua Float where + toLua = Number . show + +instance (ToLua a) => ToLua (Maybe a) where + toLua Nothing = Nil + toLua (Just a) = toLua a + +-- Helpers for Lua code generation +assign :: Name -> Exp -> Stat +assign lval exp' = LocalAssign [lval] (Just [exp']) + +funCall :: Name -> [Exp] -> Stat +funCall name' args = FunCall (NormalFunCall (var name') (Args args)) + +methCall :: Name -> Name -> [Exp] -> Stat +methCall table field args = FunCall (MethodCall (var table) field (Args args)) + +return' :: Name -> Exp +return' name' = PrefixExp (var name') + +var :: Name -> PrefixExp +var name' = PEVar (VarName name') diff --git a/NN/Backend/Torch/Torch.hs b/NN/Backend/Torch/Torch.hs new file mode 100644 --- /dev/null +++ b/NN/Backend/Torch/Torch.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE DeriveFunctor #-} +module NN.Backend.Torch.Torch where + +import Gen.Caffe.ConvolutionParameter as CP +import Gen.Caffe.DropoutParameter as DP +import Gen.Caffe.InnerProductParameter as IP +import Gen.Caffe.LayerParameter as LP +import Gen.Caffe.PoolingParameter as PP +import Gen.Caffe.PoolingParameter.PoolMethod as PP + +import Control.Applicative +import Control.Lens +import Data.Graph.Inductive.Graph hiding ((&)) +import Data.Graph.Inductive.Query +import Language.Lua.Syntax + +import NN.Backend.Torch.Lua +import NN.DSL + +-- Modules are either sequential or criterion - which are treated +-- differently by Torch +data Module a = Criterion a | Inner a deriving (Functor, Show) +data TorchModule = TorchModule Name Name [Exp] deriving (Show) + +torchExp :: TorchModule -> Exp +torchExp module' = PrefixExp (PEFunCall (construct module')) + where + construct (TorchModule luaModule torchModule args) = + NormalFunCall (PEVar (SelectName (var luaModule) torchModule)) (Args args) + +torchModules :: LayerParameter -> [Module TorchModule] +torchModules lp = go (layerTy lp) + where + nn name' args = Inner $ TorchModule "nn" name' (toLua <$> args) + criterion name' = Criterion $ TorchModule "nn" name' [] + nn' name' = nn name' ([] :: [Float]) + + -- Ugly case anaysis, sorry. + go Pool = [nn ty' [kW, kH, dW, dH]] + where + kW = poolP PP._kernel_size + kH = kW + dW = poolP PP._stride + dH = dW + ty' = case poolP PP._pool of + Just MAX -> "SpatialMaxPooling" + Just AVE -> "SpatialAveragePooling" + _ -> error "Unsupported Pooling Type" + poolP f = lp ^. LP._pooling_param ^? _Just . f . _Just + go Conv = [nn "SpatialConvolutionMM" [nInputPlane, nOutputPlane, kW, kH, dW, dH, padding]] + where + kW = convP CP._kernel_size + kH = kW + dW = convP CP._stride + dH = dW + padding = convP CP._pad + -- TODO - propagation pass to size the layers + nInputPlane = Nothing + nOutputPlane = convP CP._num_output + convP f = lp ^. LP._convolution_param ^? _Just . f . _Just + go ReLU = [nn' "Threshold"] + go IP = [nn "Linear" [nInput, nOutput]] + where + -- TODO - propagation pass to size the layers + nInput = Nothing + nOutput = lp ^. LP._inner_product_param ^? _Just . IP._num_output . _Just + go Dropout = [nn "Dropout" [ratio]] where Just ratio = lp ^. LP._dropout_param ^? _Just . DP._dropout_ratio . _Just + go SoftmaxWithLoss = [nn' "LogSoftMax", criterion "ClassNLLCriterion"] + go ty' = error $ "Unhandled layer type: " ++ show ty' + +torchLayers :: [LayerTy] +torchLayers = [Pool, Conv, ReLU, IP, Dropout, SoftmaxWithLoss] + +-- Graph validation +-- A graph is `sequential` if and only if +-- - It has n-1 edges +-- - It is connected +-- - Every node has an out degree of zero or one. +isSequential :: Net -> Bool +isSequential gr = e == (n-1) && length (dff' gr) == 1 && and [l `elem` [0, 1] | i <- nodes gr, let l = (length . suc gr) i] + where + e = length (edges gr) + n = length (nodes gr) + +clean :: Net -> Net +clean gr = foldl (flip delNode) gr toDelete + where + toDelete = filter (\n -> layerTy (label n) `notElem` torchLayers) (nodes gr) + label n = lab' (context gr n) + +linearize :: Net -> Maybe [LayerParameter] +linearize gr = if isSequential gr then Just (topsort' gr) else Nothing diff --git a/NN/Examples/MLPSweep.hs b/NN/Examples/MLPSweep.hs new file mode 100644 --- /dev/null +++ b/NN/Examples/MLPSweep.hs @@ -0,0 +1,54 @@ +module NN.Examples.MLPSweep where + +import Control.Applicative +import Control.Concurrent +import Control.Lens +import Control.Monad +import Data.Function +import Data.List +import Data.Word +import GHC.IO.Handle +import System.Exit +import System.IO.Temp +import System.Process +import Text.Read + +import NN.Backend.Torch as Torch +import NN.DSL +import NN.Graph +import NN.Passes + +-- A simple example of performing a parameter sweep over the number of +-- hidden units in an MLP. +parameterSweepMLP :: Int -> IO ([Word32], Maybe Float) +parameterSweepMLP numWorkers = maximumBy (compare `on` snd) <$> parMapIO numWorkers candidates assess + where + mlp hiddenUnits = do + _ <- sequential (concatMap (\n -> [ip n, relu]) hiddenUnits ++ [softmax]) + return () + + candidates = [[i, j, k] | let xs = [10..15], i <- xs, j <- xs, k <- xs] + + assess experiment = do + let Just torchCode = mlp experiment & parse & Torch.backend + (file, handle) <- openTempFile "/tmp" "mlp.lua" + hPutStr handle torchCode + hClose handle + (rc, stdout, _) <- readProcessWithExitCode "NN/Examples/scripts/run_mlp.lua" [file] "" + return $ case rc of + ExitSuccess -> readMaybe stdout + _ -> Nothing + +parMapIO :: Int -> [a] -> (a -> IO b) -> IO [(a, b)] +parMapIO n xs f = do + jobs <- newChan + results <- newChan + forM_ [1..n] $ \_ -> forkIO $ worker jobs results + forM_ xs (writeChan jobs) + forM xs $ \_ -> readChan results + where + worker jobs results = + forever $ do + job <- readChan jobs + result <- f job + writeChan results (job, result) diff --git a/caffegraph.cabal b/caffegraph.cabal --- a/caffegraph.cabal +++ b/caffegraph.cabal @@ -1,5 +1,5 @@ name: caffegraph -version: 0.1.0.1 +version: 0.1.0.2 description: A compiler for building, optimizing, visualizing, and generating (Caffe/Torch) DNNs license: BSD3 license-file: LICENSE @@ -18,17 +18,21 @@ GHC-Options: -Wall Hs-Source-Dirs: . exposed-modules: NN, - NN.CLI, - NN.DSL, - NN.Graph, - NN.Passes, - NN.Visualize, - NN.Backend.Caffe, + NN.Backend.Caffe + NN.Backend.Torch.Codegen + NN.Backend.Torch.Lua + NN.Backend.Torch.Torch NN.Backend.Torch - NN.Examples.ImageNet, - NN.Examples.AlexNet, - NN.Examples.GoogLeNet, - NN.Examples.Demo, + NN.CLI + NN.DSL + NN.Examples.AlexNet + NN.Examples.Demo + NN.Examples.GoogLeNet + NN.Examples.ImageNet + NN.Examples.MLPSweep + NN.Graph + NN.Passes + NN.Visualize Gen.Caffe.AccuracyParameter, Gen.Caffe.ArgMaxParameter, Gen.Caffe.BlobProto,