diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2019, Jiasen Wu +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,6 @@ +# fei-examples + ++ MNIST ++ CIFAR10 + Resnet / ResNext ++ mxnet custom operator ++ Faster RCNN diff --git a/fei-examples.cabal b/fei-examples.cabal new file mode 100644 --- /dev/null +++ b/fei-examples.cabal @@ -0,0 +1,95 @@ +name: fei-examples +version: 0.3.0 +synopsis: fei examples +description: Various fei examples +homepage: https://github.com/pierric/fei-examples#readme +bug-reports: https://github.com/pierric/fei-examples/issues +author: Jiasen Wu +maintainer: jiasenwu@hotmail.com +copyright: 2019 Jiasen Wu +license: BSD3 +license-file: LICENSE +category: Machine Learning, AI +build-type: Simple +cabal-version: >= 1.10 + +extra-source-files: + README.md + +source-repository head + type: git + location: https://github.com/pierric/fei-examples + +Executable lenet + main-is: lenet.hs + other-modules: Model.Lenet + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , unordered-containers >= 0.2.8 + , vector >= 0.12 + , fei-base + , fei-nn + , fei-dataiter + default-extensions: OverloadedLabels + , TypeFamilies + + +Executable cifar10 + main-is: cifar10.hs + other-modules: Model.Resnet, + Model.Resnext + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , unordered-containers >= 0.2.8 + , vector >= 0.12 + , optparse-applicative + , lens >= 4.12 + , fei-base + , fei-nn + , fei-dataiter + default-extensions: OverloadedLabels + , TypeFamilies + +Executable custom-op + main-is: custom-op.hs + hs-source-dirs: src + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , fei-base + , fei-nn + , fei-dataiter + , unordered-containers >= 0.2.8 + , vector >= 0.12 + default-extensions: OverloadedLabels + , TypeFamilies + , FlexibleInstances +Executable rcnn + main-is: rcnn.hs + other-modules: Model.VGG, + Model.FasterRCNN + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0 + , unordered-containers >= 0.2.10 + , vector >= 0.12 + , optparse-applicative + , attoparsec + , text + , lens >= 4.12 + , repa + , random-fu + , directory + , mtl + , conduit + , fei-base + , fei-nn + , fei-dataiter + , fei-cocoapi + default-extensions: OverloadedLabels + , TypeFamilies + diff --git a/src/Model/FasterRCNN.hs b/src/Model/FasterRCNN.hs new file mode 100644 --- /dev/null +++ b/src/Model/FasterRCNN.hs @@ -0,0 +1,606 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeApplications #-} +module Model.FasterRCNN where + +import qualified Data.Vector as V +import qualified Data.Vector.Storable as SV +import qualified Data.Vector.Unboxed as UV +import qualified Data.Vector.Unboxed.Mutable as UVM +import qualified Data.HashMap.Strict as M +import Data.IORef +import Data.Array.Repa.Index +import Data.Array.Repa.Shape +import Data.Array.Repa.Slice +import qualified Data.Array.Repa as Repa +import Data.Random (shuffle, runRVar, StdRandom(..)) +import Data.Random.Vector (randomElement) +import Control.Exception.Base(assert) +import Control.Lens ((^.), makeLenses) +import Control.Monad (replicateM, forM_, join) +import Control.Monad.IO.Class (liftIO) +import Text.Printf (printf) + +import MXNet.Base +import MXNet.Base.Operators.NDArray (_set_value_upd, argmax, argmax_channel) +import MXNet.Base.Operators.Symbol ( + elemwise_mul, elemwise_sub, smooth_l1, softmax, _SoftmaxOutput, _ROIPooling, + _MakeLoss, _contrib_MultiProposal, _BlockGrad, _Custom) +import qualified MXNet.Base.NDArray as A +import qualified MXNet.NN.NDArray as A +import MXNet.NN.Layer +import MXNet.NN.EvalMetric +import qualified Model.VGG as VGG + +import Debug.Trace + +data RcnnConfiguration = RcnnConfiguration { + rpn_anchor_scales :: [Int], + rpn_anchor_ratios :: [Float], + rpn_feature_stride :: Int, + rpn_batch_rois :: Int, + rpn_pre_topk :: Int, + rpn_post_topk :: Int, + rpn_nms_thresh :: Float, + rpn_min_size :: Int, + rpn_fg_fraction :: Float, + rpn_fg_overlap :: Float, + rpn_bg_overlap :: Float, + rpn_allowd_border :: Int, + rcnn_num_classes :: Int, + rcnn_feature_stride :: Int, + rcnn_pooled_size :: [Int], + rcnn_batch_rois :: Int, + rcnn_batch_size :: Int, + rcnn_fg_fraction :: Float, + rcnn_fg_overlap :: Float, + rcnn_bbox_stds :: [Float], + pretrained_weights :: String +} deriving Show + +symbolTrain :: RcnnConfiguration -> IO (Symbol Float) +symbolTrain RcnnConfiguration{..} = do + let numAnchors = length rpn_anchor_scales * length rpn_anchor_ratios + -- dat: + dat <- variable "data" + -- imInfo: + imInfo <- variable "im_info" + -- gtBoxes: + gtBoxes <- variable "gt_boxes" + -- rpnLabel: (batch_size, 1, numAnchors * feat_height, feat_width)) + rpnLabel <- variable "label" + -- rpnBoxTarget: (batch_size, 4 * numAnchors, feat_height, feat_width) + rpnBoxTarget <- variable "bbox_target" + -- rpnBoxWeight: (batch_size, 4 * numAnchors, feat_height, feat_width) + rpnBoxWeight <- variable "bbox_weight" + + -- VGG-15 without the last pooling layer + convFeat <- VGG.getFeature dat [2, 2, 3, 3, 3] [64, 128, 256, 512, 512] False False + + rpnConv <- convolution "rpn_conv_3x3" (#data := convFeat .& #kernel := [3,3] .& #pad := [1,1] .& #num_filter := 512 .& Nil) + rpnRelu <- activation "rpn_relu" (#data := rpnConv .& #act_type := #relu .& Nil) + + --------------------------- + -- rpn_clas_prob part + -- + -- per pixel: fore/back-ground classification + rpnClsScore <- convolution "rpn_cls_score" (#data := rpnRelu .& #kernel := [1,1] .& #pad := [0,0] .& #num_filter := 2 * numAnchors .& Nil) + rpnClsScoreReshape <- reshape "rpn_cls_score_reshape" (#data := rpnClsScore .& #shape := [0, 2, -1, 0] .& Nil) + -- rpnClsProb output shape: (batch_size, [Pr(foreground), Pr(background)], numAnchors * feat_height, feat_width) + rpnClsProb <- _SoftmaxOutput "rpn_cls_prob" (#data := rpnClsScoreReshape .& #label := rpnLabel .& #multi_output := True + .& #normalization := #valid .& #use_ignore := True .& #ignore_label := -1 .& Nil) + + --------------------------- + -- rpn_bbox part + rpnBBoxPred <- convolution "rpn_bbox_pred" (#data := rpnRelu .& #kernel := [1,1] .& #pad := [0,0] .& #num_filter := 4 * numAnchors .& Nil) + rpnBBoxPredReg <- elemwise_sub "rpn_bbox_pred_reg" (#lhs := rpnBBoxPred .& #rhs := rpnBoxTarget .& Nil) + rpnBBoxPredRegSmooth <- smooth_l1 "rpn_bbox_pred_reg_smooth" (#data := rpnBBoxPredReg .& #scalar := 3.0 .& Nil) + rpnBBoxLoss_ <- elemwise_mul "rpn_bbox_loss_" (#lhs := rpnBoxWeight .& #rhs := rpnBBoxPredRegSmooth .& Nil) + rpnBBoxLoss <- _MakeLoss "rpn_bbox_loss" (#data := rpnBBoxLoss_ .& #grad_scale := 1.0 / fromIntegral rpn_batch_rois .& Nil) + + --------------------------- + rpnClsAct <- softmax "rpn_cls_act" (#data := rpnClsScoreReshape .& #axis := 1 .& Nil) + rpnClsActReshape <- reshape "rpn_cls_act_reshape" (#data := rpnClsAct .& #shape := [0, 2 * numAnchors, -1, 0] .& Nil) + rois <- _contrib_MultiProposal "rois" (#cls_prob := rpnClsActReshape .& #bbox_pred := rpnBBoxPred .& #im_info := imInfo + .& #feature_stride := rpn_feature_stride .& #scales := map fromIntegral rpn_anchor_scales .& #ratios := rpn_anchor_ratios + .& #rpn_pre_nms_top_n := rpn_pre_topk .& #rpn_post_nms_top_n := rpn_post_topk + .& #threshold := rpn_nms_thresh .& #rpn_min_size := rpn_min_size .& Nil) + + proposal <- _Custom "proposal" (#data := [rois, gtBoxes] + .& #op_type := "proposal_target" + .& #num_classes :≅ rcnn_num_classes + .& #batch_images:≅ rcnn_batch_size + .& #batch_rois :≅ rcnn_batch_rois + .& #fg_fraction :≅ rcnn_fg_fraction + .& #fg_overlap :≅ rcnn_fg_overlap + .& #box_stds :≅ rcnn_bbox_stds + .& Nil) + [rois, label, bboxTarget, bboxWeight] <- mapM (at proposal) [0..3] + + --------------------------- + -- cls_prob part + -- + roiPool <- _ROIPooling "roi_pool" (#data := convFeat .& #rois := rois + .& #pooled_size := rcnn_pooled_size + .& #spatial_scale := 1.0 / fromIntegral rcnn_feature_stride .& Nil) + topFeat <- VGG.getTopFeature (Just "rcnn_") roiPool + clsScore <- fullyConnected "cls_score" (#data := topFeat .& #num_hidden := rcnn_num_classes .& Nil) + clsProb <- _SoftmaxOutput "cls_prob" (#data := clsScore .& #label := label .& #normalization := #batch .& Nil) + + --------------------------- + -- bbox_loss part + -- + bboxPred <- fullyConnected "bbox_pred" (#data := topFeat .& #num_hidden := 4 * rcnn_num_classes .& Nil) + bboxPredReg <- elemwise_sub "bbox_pred_reg" (#lhs := bboxPred .& #rhs := bboxTarget .& Nil) + bboxPredRegSmooth <- smooth_l1 "bbox_pred_reg_smooth" (#data := bboxPredReg .& #scalar := 1.0 .& Nil) + bboxLoss_ <- elemwise_mul "bbox_loss_" (#lhs := bboxPredRegSmooth .& #rhs := bboxWeight .& Nil) + bboxLoss <- _MakeLoss "bbox_loss" (#data := bboxLoss_ .& #grad_scale := 1.0 / fromIntegral rcnn_batch_rois .& Nil) + + labelReshape <- reshape "label_reshape" (#data := label .& #shape := [rcnn_batch_size, -1] .& Nil) + clsProbReshape <- reshape "cls_prob_reshape" (#data := clsProb .& #shape := [rcnn_batch_size, -1, rcnn_num_classes] .& Nil) + bboxLossReshape <- reshape "bbox_loss_reshape" (#data := bboxLoss .& #shape := [rcnn_batch_size, -1, 4 * rcnn_num_classes] .& Nil) + labelSG <- _BlockGrad "label_sg" (#data := labelReshape .& Nil) + + Symbol <$> group [rpnClsProb, rpnBBoxLoss, clsProbReshape, bboxLossReshape, labelSG] + +-------------------------------- + +data ProposalTargetProp = ProposalTargetProp { + _num_classes :: Int, + _batch_images :: Int, + _batch_rois :: Int, + _fg_fraction :: Float, + _fg_overlap :: Float, + _box_stds :: [Float] +} +makeLenses ''ProposalTargetProp + +instance CustomOperationProp ProposalTargetProp where + prop_list_arguments _ = ["rois", "gt_boxes"] + prop_list_outputs _ = ["rois_output", "label", "bbox_target", "bbox_weight"] + prop_list_auxiliary_states _ = [] + prop_infer_shape prop [rpn_rois_shape, gt_boxes_shape] = + let prop_batch_size = prop ^. batch_rois + prop_num_classes = prop ^. num_classes + output_rois_shape = [prop_batch_size, 5] + label_shape = [prop_batch_size] + bbox_target_shape = [prop_batch_size, prop_num_classes * 4] + bbox_weight_shape = [prop_batch_size, prop_num_classes * 4] + in ([rpn_rois_shape, gt_boxes_shape], + [output_rois_shape, label_shape, bbox_target_shape, bbox_weight_shape], + []) + prop_declare_backward_dependency prop grad_out data_in data_out = [] + + data Operation ProposalTargetProp = ProposalTarget ProposalTargetProp + prop_create_operator prop _ _ = return (ProposalTarget prop) + +instance CustomOperation (Operation ProposalTargetProp) where + forward (ProposalTarget prop) [ReqWrite, ReqWrite, ReqWrite, ReqWrite] inputs outputs aux is_train = do + -- :param: rois, shape of (N*nms_top_n, 5), [image_index_in_batch, bbox0, bbox1, bbox2, bbox3] + -- :param: gt_boxes, shape of (N, M, 5), M varies per image. [bbox0, bbox1, bbox2, bbox3, class] + let [rois, gt_boxes] = inputs + [rois_output, label_output, bbox_target_output, bbox_weight_output] = outputs + batch_size = prop ^. batch_images + + -- convert NDArray to Vector of Repa array. + r_rois <- toRepa @DIM2 (NDArray rois) >>= return . toRows2 + r_gt <- toRepa @DIM3 (NDArray gt_boxes) >>= return . toRows3 + + assert (batch_size == length r_gt) (return ()) + + (rois, labels, bbox_targets, bbox_weights) <- V.unzip4 <$> V.mapM (sample_batch r_rois r_gt) (V.enumFromN (0 :: Int) batch_size) + let rois' = vstack $ V.map (Repa.reshape (Z :. 1 :. 5)) $ join rois + labels' = join labels + bbox_targets' = vstack bbox_targets + bbox_weights' = vstack bbox_weights + + rois_output_nd = NDArray rois_output :: NDArray Float + bbox_target_output_nd = NDArray bbox_target_output :: NDArray Float + bbox_weight_output_nd = NDArray bbox_weight_output :: NDArray Float + label_output_nd = NDArray label_output :: NDArray Float + + ndsize rois_output_nd >>= \s -> assert (s == Repa.size (Repa.extent rois')) (return ()) + ndsize bbox_target_output_nd >>= \s -> assert (s == Repa.size (Repa.extent bbox_targets')) (return ()) + ndsize bbox_weight_output_nd >>= \s -> assert (s == Repa.size (Repa.extent bbox_weights')) (return ()) + + copyFromRepa rois_output_nd rois' + copyFromRepa bbox_target_output_nd bbox_targets' + copyFromRepa bbox_weight_output_nd bbox_weights' + copyFromVector label_output_nd $ V.convert labels' + + where + toRows2 arr = let Z :. rows :._ = Repa.extent arr + range = V.enumFromN (0 :: Int) rows + in V.map (\i -> Repa.computeUnboxedS $ Repa.slice arr (Z :. i :. All)) range + + toRows3 arr = let Z :. rows :. _ :. _ = Repa.extent arr + range = V.enumFromN (0 :: Int) rows + in V.map (\i -> Repa.computeUnboxedS $ Repa.slice arr (Z :. i :. All :. All)) range + + sample_batch :: V.Vector (Repa.Array Repa.U DIM1 Float) -> V.Vector (Repa.Array _ DIM2 Float) -> Int -> IO (_, _, _, _) + sample_batch r_rois r_gt index = do + let rois_this_image = V.filter (\roi -> floor (roi #! 0) == index) r_rois + all_gt_this_image = toRows2 $ r_gt %! index + gt_this_image = V.filter (\gt -> gt #! 4 > 0) all_gt_this_image + + let num_rois_per_image = (prop ^. batch_rois) `div` (prop ^. batch_images) + fg_rois_per_image = round (prop ^. fg_fraction * fromIntegral num_rois_per_image) + + -- WHY? + -- append gt boxes to rois + let prepend_index = Repa.computeUnboxedS . (Repa.fromListUnboxed (Z :. 1) [fromIntegral index] Repa.++) + gt_boxes_as_rois = V.map (\gt -> prepend_index $ Repa.extract (Z :. 0) (Z :. 4) gt) gt_this_image + rois_this_image' = rois_this_image V.++ gt_boxes_as_rois + + sample_rois rois_this_image' gt_this_image + (prop ^. num_classes) num_rois_per_image fg_rois_per_image (prop ^. fg_overlap) (prop ^. box_stds) + + backward _ [ReqWrite, ReqWrite] _ _ [in_grad_0, in_grad_1] _ _ = do + _set_value_upd [in_grad_0] (#src := 0 .& Nil) + _set_value_upd [in_grad_1] (#src := 0 .& Nil) + + +sample_rois :: V.Vector (Repa.Array Repa.U DIM1 Float) -> V.Vector (Repa.Array Repa.U DIM1 Float) -> Int -> Int -> Int -> Float -> [Float] + -> IO (V.Vector (Repa.Array Repa.U Repa.DIM1 Float), + V.Vector Float, + Repa.Array _ Repa.DIM2 Float, + Repa.Array _ Repa.DIM2 Float) +sample_rois rois gt num_classes rois_per_image fg_rois_per_image fg_overlap box_stds = do + -- :param rois: [num_rois, 5] (batch_index, x1, y1, x2, y2) + -- :param gt: [num_rois, 5] (x1, y1, x2, y2, cls) + -- + -- :returns: sampled (rois, labels, regression, weight) + let num_rois = V.length rois + -- print(num_rois, V.length gt_boxes) + -- assert (num_rois == V.length gt_boxes) (return ()) + let aoi_boxes = V.map (Repa.computeUnboxedS . Repa.extract (Z:.1) (Z:.4)) rois + gt_boxes = V.map (Repa.computeUnboxedS . Repa.extract (Z:.0) (Z:.4)) gt + overlaps = Repa.computeUnboxedS $ overlapMatrix aoi_boxes gt_boxes + + let maxIndices = argMax overlaps + gt_chosen = V.map (gt %!) maxIndices + + -- a uniform sampling w/o replacement from the fg boxes if there are too many + fg_indexes <- let fg_indexes = V.filter (\(i, j) -> Repa.index overlaps (Z :. i :. j) >= fg_overlap) (V.indexed maxIndices) + in if length fg_indexes > fg_rois_per_image then + V.fromList . take fg_rois_per_image <$> runRVar' (shuffle $ V.toList fg_indexes) + else + return fg_indexes + + -- slightly different from the orignal implemetation: + -- a uniform sampling w/ replacement if not enough bg boxes + let bg_rois_this_image = rois_per_image - length fg_indexes + bg_indexes <- let bg_indexes = V.filter (\(i, j) -> Repa.index overlaps (Z :. i :. j) < fg_overlap) (V.indexed maxIndices) + num_bg_indexes = length bg_indexes + in case compare num_bg_indexes bg_rois_this_image of + GT -> V.fromList . take bg_rois_this_image <$> runRVar' (shuffle $ V.toList bg_indexes) + LT -> V.fromList <$> runRVar' (replicateM bg_rois_this_image (randomElement bg_indexes)) + EQ -> return bg_indexes + + let keep_indexes = V.map fst $ fg_indexes V.++ bg_indexes + + rois_keep = V.map (rois %!) keep_indexes + roi_box_keep = V.map (asTuple . Repa.computeUnboxedS . Repa.extract (Z:.1) (Z:.4)) rois_keep + + gt_keep = V.map (gt_chosen %!) keep_indexes + gt_box_keep = V.map (asTuple . Repa.computeUnboxedS . Repa.extract (Z:.0) (Z:.4)) gt_keep + labels_keep = V.take (length fg_indexes) (V.map (#! 4) gt_keep) V.++ V.replicate bg_rois_this_image 0 + + targets = V.zipWith (bboxTransform box_stds) roi_box_keep gt_box_keep + + -- regression is indexed by class. + bbox_target <- UVM.replicate (rois_per_image * 4 * num_classes) (0 :: Float) + bbox_weight <- UVM.replicate (rois_per_image * 4 * num_classes) (0 :: Float) + + -- only assign regression and weights for the foreground boxes. + forM_ [0..length fg_indexes-1] $ \i -> do + let lbl = floor (labels_keep %! i) + (tgt0, tgt1, tgt2, tgt3) = targets %! i :: Box + assert (lbl >= 0 && lbl < num_classes) (return ()) + let tgt_dst = UVM.slice (i * 4 * num_classes + 4 * lbl) 4 bbox_target + UVM.write tgt_dst 0 tgt0 + UVM.write tgt_dst 1 tgt1 + UVM.write tgt_dst 2 tgt2 + UVM.write tgt_dst 3 tgt3 + let wgh_dst = UVM.slice (i * 4 * num_classes + 4 * lbl) 4 bbox_weight + UVM.set wgh_dst 1 + + let shape = Z :. rois_per_image :. 4 * num_classes + bbox_target <- Repa.fromUnboxed shape <$> UV.freeze bbox_target + bbox_weight <- Repa.fromUnboxed shape <$> UV.freeze bbox_weight + return (rois_keep, labels_keep, bbox_target, bbox_weight) + + where + runRVar' = flip runRVar StdRandom + +overlapMatrix :: V.Vector (Repa.Array Repa.U Repa.DIM1 Float) -> V.Vector (Repa.Array Repa.U Repa.DIM1 Float) -> Repa.Array Repa.D Repa.DIM2 Float +overlapMatrix rois gt = Repa.fromFunction (Z :. width :. height) calcOvp + where + width = length rois + height = length gt + + calcArea box = (box #! 2 - box #! 0 + 1) * (box #! 3 - box #! 1 + 1) + area1 = V.map calcArea rois + area2 = V.map calcArea gt + + calcOvp (Z :. ind_rois :. ind_gt) = + let b1 = rois %! ind_rois + b2 = gt %! ind_gt + iw = min (b1 #! 2) (b2 #! 2) - max (b1 #! 0) (b2 #! 0) + 1 + ih = min (b1 #! 3) (b2 #! 3) - max (b1 #! 1) (b2 #! 1) + 1 + areaI = iw * ih + areaU = area1 %! ind_rois + area2 %! ind_gt - areaI + in if iw > 0 && ih > 0 then areaI / areaU else 0 + +argMax overlaps = + let Z :. m :. n = Repa.extent overlaps + findMax row = UV.maxIndex $ Repa.toUnboxed $ Repa.computeS $ Repa.slice overlaps (Z :. row :. All) + in V.map findMax $ V.enumFromN (0 :: Int) m + +type Box = (Float, Float, Float, Float) +whctr :: Box -> Box +whctr (x0, y0, x1, y1) = (w, h, x, y) + where + w = x1 - x0 + 1 + h = y1 - y0 + 1 + x = x0 + 0.5 * (w - 1) + y = y0 + 0.5 * (h - 1) + +asTuple :: Repa.Array Repa.U Repa.DIM1 Float -> (Float, Float, Float, Float) +asTuple box = (box #! 0, box #! 1, box #! 2, box #! 3) + +bboxTransform :: [Float] -> Box -> Box -> Box +bboxTransform [std0, std1, std2, std3] box1 box2 = + let (w1, h1, cx1, cy1) = whctr box1 + (w2, h2, cx2, cy2) = whctr box2 + dx = (cx2 - cx1) / (w1 + 1e-14) / std0 + dy = (cy2 - cy1) / (h1 + 1e-14) / std1 + dw = log (w2 / w1) / std2 + dh = log (h2 / h1) / std3 + in (dx, dy, dw, dh) + +(#!) :: (Shape sh, UV.Unbox e) => Repa.Array Repa.U sh e -> Int -> e +(#!) = Repa.linearIndex +(%!) = (V.!) + +vstack :: Repa.Source r Float => V.Vector (Repa.Array r Repa.DIM2 Float) -> Repa.Array Repa.D Repa.DIM2 Float +vstack = Repa.transpose . V.foldl1 (Repa.++) . V.map Repa.transpose + + +test_sample_rois = let + v1 = Repa.fromListUnboxed (Z:.5::DIM1) [0, 0.8, 0.8, 2.2, 2.2] + v2 = Repa.fromListUnboxed (Z:.5::DIM1) [0, 2.2, 2.2, 4.5, 4.5] + v3 = Repa.fromListUnboxed (Z:.5::DIM1) [0, 4.2, 1, 6.5, 2.8] + v4 = Repa.fromListUnboxed (Z:.5::DIM1) [0, 6, 3, 7, 4] + rois = V.fromList [v1, v2, v3, v4] + g1 = Repa.fromListUnboxed (Z:.5::DIM1) [1,1,2,2,1] + g2 = Repa.fromListUnboxed (Z:.5::DIM1) [2,3,3,4,1] + g3 = Repa.fromListUnboxed (Z:.5::DIM1) [4,1,6,3,2] + gt_boxes = V.fromList [g1, g2, g3] + in sample_rois rois gt_boxes 3 6 2 0.5 [0.1, 0.1, 0.1, 0.1] + + +data RPNAccMetric a = RPNAccMetric Int String + +instance EvalMetricMethod RPNAccMetric where + data MetricData RPNAccMetric a = RPNAccMetricData String Int String (IORef Int) (IORef Int) + newMetric phase (RPNAccMetric oindex label) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RPNAccMetricData phase oindex label a b + + format (RPNAccMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (100 * fromIntegral s / fromIntegral n :: Float) + + evaluate (RPNAccMetricData phase oindex lname cntRef sumRef) bindings outputs = liftIO $ do + let label = bindings M.! lname + pred = outputs !! oindex + + pred <- A.makeNDArrayLike pred contextCPU >>= A.copy pred + [pred_label] <- argmax_channel (#data := unNDArray pred .& Nil) + pred_label <- V.convert <$> toVector (NDArray pred_label) + label <- V.convert <$> toVector label + + let pairs = V.filter ((/= -1) . fst) $ V.zip label pred_label + equal = V.filter (uncurry (==)) pairs + + modifyIORef' sumRef (+ length equal) + modifyIORef' cntRef (+ length pairs) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = fromIntegral s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + + +data RCNNAccMetric a = RCNNAccMetric Int Int + +instance EvalMetricMethod RCNNAccMetric where + data MetricData RCNNAccMetric a = RCNNAccMetricData String Int Int (IORef Int) (IORef Int) + newMetric phase (RCNNAccMetric cindex lindex) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RCNNAccMetricData phase cindex lindex a b + + format (RCNNAccMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (100 * fromIntegral s / fromIntegral n :: Float) + + evaluate (RCNNAccMetricData phase cindex lindex cntRef sumRef) bindings outputs = liftIO $ do + -- cls_prob: (batch_size, #num_anchors*feat_w*feat_h, #num_classes) + -- label: (batch_size, #num_anchors*feat_w*feat_h) + let cls_prob = outputs !! cindex + label = outputs !! lindex + + cls_prob <- A.makeNDArrayLike cls_prob contextCPU >>= A.copy cls_prob + [pred_class] <- argmax (#data := unNDArray cls_prob .& #axis := Just 2 .& Nil) + + pred_class <- toRepa @DIM2 (NDArray pred_class) + label <- toRepa @DIM2 label + + let pairs = UV.zip (Repa.toUnboxed label) (Repa.toUnboxed pred_class) + equal = UV.filter (uncurry (==)) pairs + + modifyIORef' sumRef (+ UV.length equal) + modifyIORef' cntRef (+ UV.length pairs) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = fromIntegral s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + +data RPNLogLossMetric a = RPNLogLossMetric Int String + +instance EvalMetricMethod RPNLogLossMetric where + data MetricData RPNLogLossMetric a = RPNLogLossMetricData String Int String (IORef Int) (IORef Double) + newMetric phase (RPNLogLossMetric cindex lname) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RPNLogLossMetricData phase cindex lname a b + + format (RPNLogLossMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (realToFrac s / fromIntegral n :: Float) + + evaluate (RPNLogLossMetricData phase cindex lname cntRef sumRef) bindings outputs = liftIO $ do + let cls_prob = outputs !! cindex + label = bindings M.! lname + + -- (batch_size, #num_anchors*feat_w*feat_h) to (batch_size*#num_anchors*feat_w*feat_h,) + label <- A.reshape label [-1] + label <- toRepa @DIM1 label + let Z :. size = Repa.extent label + + -- (batch_size, #channel, #num_anchors*feat_w, feat_h) to (batch_size, #channel, #num_anchors*feat_w*feat_h) + -- to (batch_size, #num_anchors*feat_w*feat_h, #channel) to (batch_size*#num_anchors*feat_w*feat_h, #channel) + cls_prob <- A.makeNDArrayLike cls_prob contextCPU >>= A.copy cls_prob + pred <- A.reshape cls_prob [0, 0, -1] >>= flip A.transpose [0, 2, 1] >>= flip A.reshape [size, -1] + pred <- toRepa @DIM2 pred + + -- mark out labels where value -1 + let mask = Repa.computeUnboxedS $ Repa.map (/= -1) label + + pred <- Repa.selectP (mask #!) (\i -> pred Repa.! (Z :. i :. (floor $ label #! i))) size + -- traceShowM pred + label <- Repa.selectP (mask #!) (label #!) size + + let pred_with_ep = Repa.map ((0 -) . log) (pred Repa.+^ constant (Z :. size) 1e-14) + cls_loss <- Repa.foldP (+) 0 pred_with_ep + + let cls_loss_val = realToFrac (cls_loss #! 0) + modifyIORef' sumRef (+ cls_loss_val) + modifyIORef' cntRef (+ size) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + +data RCNNLogLossMetric a = RCNNLogLossMetric Int Int + +instance EvalMetricMethod RCNNLogLossMetric where + data MetricData RCNNLogLossMetric a = RCNNLogLossMetricData String Int Int (IORef Int) (IORef Double) + newMetric phase (RCNNLogLossMetric cindex lindex) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RCNNLogLossMetricData phase cindex lindex a b + + format (RCNNLogLossMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (realToFrac s / fromIntegral n :: Float) + + evaluate (RCNNLogLossMetricData phase cindex lindex cntRef sumRef) bindings outputs = liftIO $ do + let cls_prob = outputs !! cindex + label = outputs !! lindex + + cls_prob <- toRepa @DIM3 cls_prob + label <- toRepa @DIM2 label + + let lbl_shp@(Z :. _ :. size) = Repa.extent label + cls = Repa.fromFunction lbl_shp (\ pos@(Z :. bi :. ai) -> cls_prob Repa.! (Z :. bi :. ai :. (floor $ label Repa.! pos))) + + cls_loss_val <- Repa.sumAllP $ Repa.map (\v -> - log(1e-14 + v)) cls + -- traceShowM cls_loss_val + modifyIORef' sumRef (+ realToFrac cls_loss_val) + modifyIORef' cntRef (+ size) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + +data RPNL1LossMetric a = RPNL1LossMetric Int String + +instance EvalMetricMethod RPNL1LossMetric where + data MetricData RPNL1LossMetric a = RPNL1LossMetricData String Int String (IORef Int) (IORef Double) + newMetric phase (RPNL1LossMetric bindex blabel) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RPNL1LossMetricData phase bindex blabel a b + + format (RPNL1LossMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (realToFrac s / fromIntegral n :: Float) + + evaluate (RPNL1LossMetricData phase bindex blabel cntRef sumRef) bindings outputs = liftIO $ do + let bbox_loss = outputs !! bindex + bbox_weight = bindings M.! blabel + + bbox_loss <- toRepa @DIM4 bbox_loss + all_loss <- Repa.sumAllP bbox_loss + + bbox_weight <- toRepa @DIM4 bbox_weight + all_pos_weight <- Repa.sumAllP $ Repa.map (\w -> if w > 0 then 1 else 0) bbox_weight + + modifyIORef' sumRef (+ realToFrac all_loss) + modifyIORef' cntRef (+ (all_pos_weight `div` 4)) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + +data RCNNL1LossMetric a = RCNNL1LossMetric Int Int + +instance EvalMetricMethod RCNNL1LossMetric where + data MetricData RCNNL1LossMetric a = RCNNL1LossMetricData String Int Int (IORef Int) (IORef Double) + newMetric phase (RCNNL1LossMetric bindex lindex) = do + a <- liftIO $ newIORef 0 + b <- liftIO $ newIORef 0 + return $ RCNNL1LossMetricData phase bindex lindex a b + + format (RCNNL1LossMetricData _ _ _ cntRef sumRef) = liftIO $ do + s <- liftIO $ readIORef sumRef + n <- liftIO $ readIORef cntRef + return $ printf "" (realToFrac s / fromIntegral n :: Float) + + evaluate (RCNNL1LossMetricData phase bindex lindex cntRef sumRef) bindings outputs = liftIO $ do + let bbox_loss = outputs !! bindex + label = outputs !! lindex + + bbox_loss <- toRepa @DIM3 bbox_loss + all_loss <- Repa.sumAllP bbox_loss + + label <- toRepa @DIM2 label + all_pos <- Repa.sumAllP $ Repa.map (\w -> if w > 0 then 1 else 0) label + + modifyIORef' sumRef (+ realToFrac all_loss) + modifyIORef' cntRef (+ all_pos) + + s <- readIORef sumRef + n <- readIORef cntRef + let acc = s / fromIntegral n + return $ M.singleton (phase ++ "_acc") acc + +constant :: (Shape sh, UV.Unbox a) => sh -> a -> Repa.Array Repa.U sh a +constant shp val = Repa.fromListUnboxed shp (replicate (size shp) val) diff --git a/src/Model/Lenet.hs b/src/Model/Lenet.hs new file mode 100644 --- /dev/null +++ b/src/Model/Lenet.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module Model.Lenet (symbol) where + +import MXNet.Base +import MXNet.NN.Layer + +-- # first conv +-- conv1 = mx.symbol.Convolution(data=data, kernel=(5,5), num_filter=20) +-- tanh1 = mx.symbol.Activation(data=conv1, act_type="tanh") +-- pool1 = mx.symbol.Pooling(data=tanh1, pool_type="max", kernel=(2,2), stride=(2,2)) +-- # second conv +-- conv2 = mx.symbol.Convolution(data=pool1, kernel=(5,5), num_filter=50) +-- tanh2 = mx.symbol.Activation(data=conv2, act_type="tanh") +-- pool2 = mx.symbol.Pooling(data=tanh2, pool_type="max", kernel=(2,2), stride=(2,2)) +-- # first fullc +-- flatten = mx.symbol.Flatten(data=pool2) +-- fc1 = mx.symbol.FullyConnected(data=flatten, num_hidden=500) +-- tanh3 = mx.symbol.Activation(data=fc1, act_type="tanh") +-- # second fullc +-- fc2 = mx.symbol.FullyConnected(data=tanh3, num_hidden=num_classes) +-- # loss +-- lenet = mx.symbol.SoftmaxOutput(data=fc2, name='softmax') + +symbol :: DType a => IO (Symbol a) +symbol = do + x <- variable "x" + y <- variable "y" + + v1 <- convolution "conv1" (#data := x .& #kernel := [5,5] .& #num_filter := 20 .& Nil) + a1 <- activation "conv1-a" (#data := v1 .& #act_type := #tanh .& Nil) + p1 <- pooling "conv1-p" (#data := a1 .& #kernel := [2,2] .& #pool_type := #max .& Nil) + + v2 <- convolution "conv2" (#data := p1 .& #kernel := [5,5] .& #num_filter := 50 .& Nil) + a2 <- activation "conv2-a" (#data := v2 .& #act_type := #tanh .& Nil) + p2 <- pooling "conv2-p" (#data := a2 .& #kernel := [2,2] .& #pool_type := #max .& Nil) + + fl <- flatten "flatten" (#data := p2 .& Nil) + + v3 <- fullyConnected "fc1" (#data := fl .& #num_hidden := 500 .& Nil) + a3 <- activation "fc1-a" (#data := v3 .& #act_type := #tanh .& Nil) + + v4 <- fullyConnected "fc2" (#data := a3 .& #num_hidden := 10 .& Nil) + a4 <- softmaxoutput "softmax" (#data := v4 .& #label := y .& Nil) + return $ Symbol a4 diff --git a/src/Model/Resnet.hs b/src/Model/Resnet.hs new file mode 100644 --- /dev/null +++ b/src/Model/Resnet.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} + +module Model.Resnet (symbol) where + +import Control.Monad (foldM, when, void) +import Control.Exception.Base (Exception, throw, throwIO) +import Data.Maybe (fromMaybe) +import Data.Typeable (Typeable) + +import MXNet.Base +import MXNet.NN.Layer + +data NoKnownExperiment = NoKnownExperiment Int + deriving (Typeable, Show) +instance Exception NoKnownExperiment + +------------------------------------------------------------------------------- +-- ResNet + +symbol :: DType a => Int -> Int -> [Int] -> IO (Symbol a) +symbol num_classes num_layers image_shape@[_, height, _] = + if height <= 28 then do + handle <- if (num_layers - 2) `mod` 9 == 0 && num_layers >= 164 then + resnet $ + #image_shape := image_shape .& + #num_classes := num_classes .& + #num_stages := 3 .& + #filter_list := [64, 64, 128, 256] .& + #units := replicate 3 ((num_layers - 2) `div` 9) .& + #bottle_neck := True .& + #workspace := 256 .& Nil + else if (num_layers - 2) `mod` 6 == 0 && num_layers < 164 then + resnet $ + #image_shape := image_shape .& + #num_classes := num_classes .& + #num_stages := 3 .& + #filter_list := [64, 64, 32, 64] .& + #units := replicate 3 ((num_layers - 2) `div` 6) .& + #bottle_neck := False .& + #workspace := 256 .& Nil + else + throwIO $ NoKnownExperiment num_layers + return $ Symbol handle + else do + handle <- resnet $ #image_shape := image_shape .& #num_classes := num_classes .& #num_stages := 4 .& case num_layers of + 18 -> #filter_list := [64, 64, 128, 256, 512] .& #units := [2,2,2,2] .& #bottle_neck := False .& #workspace := 256 .& Nil + 34 -> #filter_list := [64, 64, 128, 256, 512] .& #units := [3,4,6,3] .& #bottle_neck := False .& #workspace := 256 .& Nil + 50 -> #filter_list := [64, 256, 512, 1024, 2048] .& #units := [3,4,6,3] .& #bottle_neck := True .& #workspace := 256 .& Nil + 101 -> #filter_list := [64, 256, 512, 1024, 2048] .& #units := [3,4,23,3] .& #bottle_neck := True .& #workspace := 256 .& Nil + 152 -> #filter_list := [64, 256, 512, 1024, 2048] .& #units := [3,8,36,3] .& #bottle_neck := True .& #workspace := 256 .& Nil + 200 -> #filter_list := [64, 256, 512, 1024, 2048] .& #units := [3,24,36,3] .& #bottle_neck := True .& #workspace := 256 .& Nil + 269 -> #filter_list := [64, 256, 512, 1024, 2048] .& #units := [3,30,48,8] .& #bottle_neck := True .& #workspace := 256 .& Nil + _ -> throw $ NoKnownExperiment num_layers + return $ Symbol handle + +type instance ParameterList "resnet" = + '[ '("num_classes", 'AttrReq Int) + , '("num_stages" , 'AttrReq Int) + , '("filter_list", 'AttrReq [Int]) + , '("units" , 'AttrReq [Int]) + , '("bottle_neck", 'AttrReq Bool) + , '("workspace" , 'AttrReq Int) + , '("image_shape", 'AttrReq [Int])] +resnet :: (Fullfilled "resnet" args) => ArgsHMap "resnet" args -> IO SymbolHandle +resnet args = do + x <- variable "x" + y <- variable "y" + + xcp <- identity "id" ( + #data := x .& Nil) + + bnx <- batchnorm "bn-x" ( + #data := xcp .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := True .& Nil) + + let [_, height, _] = args ! #image_shape + filter0 : filter_list = args ! #filter_list + bdy <- if height <= 32 + then + convolution "conv-bn-x" ( + #data := bnx .& + #kernel := [3,3] .& + #num_filter:= filter0 .& + #stride := [1,1] .& + #pad := [1,1] .& + #workspace := conv_workspace .& + #no_bias := True .& Nil) + else do + bdy <- convolution "conv-bn-x" ( + #data := bnx .& + #kernel := [7,7] .& + #num_filter:= filter0 .& + #stride := [2,2] .& + #pad := [3,3] .& + #workspace := conv_workspace .& + #no_bias := True .& Nil) + bdy <- batchnorm "bn-0" ( + #data := bdy .& + #fix_gamma := False .& + #eps := eps .& + #momentum := bn_mom .& Nil) + bdy <- activation "relu0" ( + #data := bdy .& + #act_type := #relu .& Nil) + pooling "max" ( + #data := bdy .& + #kernel := [3,3] .& + #stride := [2,2] .& + #pad := [1,1] .& + #pool_type := #max .& Nil) + + bdy <- foldM build_layer bdy (zip3 [0::Int ..] filter_list (args ! #units)) + + bn1 <- batchnorm "bn-1" ( + #data := bdy .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + ac1 <- activation "relu-1" ( + #data := bn1 .& + #act_type := #relu .& Nil) + pl1 <- pooling "pool-1" ( + #data := ac1 .& + #kernel := [7,7] .& + #pool_type := #avg .& + #global_pool := True .& Nil) + + flt <- flatten "flt-1" ( + #data := pl1 .& Nil) + fc1 <- fullyConnected "fc-1" ( + #data := flt .& + #num_hidden := args ! #num_classes .& Nil) + + softmaxoutput "softmax" ( + #data := fc1 .& + #label := y .& Nil) + where + bn_mom = 0.9 :: Float + conv_workspace = 256 :: Int + eps = 2e-5 :: Double + + build_layer bdy (stage_id, filter_size, unit) = do + let stride0 = if stage_id == 0 then [1,1] else [2,2] + name unit_id = "stage" ++ show stage_id ++ "_unit" ++ show unit_id + resargs = #bottle_neck := False .& #workspace := conv_workspace .& #memonger := False .& Nil + bdy <- residual (name 0) (#data := bdy .& #num_filter := filter_size .& #stride := stride0 .& #dim_match := False .& resargs) + foldM (\bdy unit_id -> + residual (name unit_id) (#data := bdy .& #num_filter := filter_size .& #stride := [1,1] .& #dim_match := True .& resargs)) + bdy [1..unit] + +type instance ParameterList "_residual_layer(resnet)" = + '[ '("data" , 'AttrReq SymbolHandle) + , '("num_filter" , 'AttrReq Int) + , '("stride" , 'AttrReq [Int]) + , '("dim_match" , 'AttrReq Bool) + , '("bottle_neck", 'AttrOpt Bool) + , '("bn_mom" , 'AttrOpt Float) + , '("workspace" , 'AttrOpt Int) + , '("memonger" , 'AttrOpt Bool) ] +residual :: (Fullfilled "_residual_layer(resnet)" args) + => String -> ArgsHMap "_residual_layer(resnet)" args -> IO SymbolHandle +residual name args = do + let dat = args ! #data + num_filter = args ! #num_filter + stride = args ! #stride + dim_match = args ! #dim_match + bottle_neck= fromMaybe True $ args !? #bottle_neck + bn_mom = fromMaybe 0.9 $ args !? #bn_mom + workspace = fromMaybe 256 $ args !? #workspace + memonger = fromMaybe False$ args !? #memonger + eps = 2e-5 :: Double + if bottle_neck + then do + bn1 <- batchnorm (name ++ "-bn1") ( + #data := dat .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act1 <- activation (name ++ "-relu1") ( + #data := bn1 .& + #act_type := #relu .& Nil) + conv1 <- convolution (name ++ "-conv1") ( + #data := act1 .& + #kernel := [1,1] .& + #num_filter := num_filter `div` 4 .& + #stride := [1,1] .& + #pad := [0,0] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn2 <- batchnorm (name ++ "-bn2") ( + #data := conv1 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act2 <- activation (name ++ "-relu2") ( + #data := bn2 .& + #act_type := #relu .& Nil) + conv2 <- convolution (name ++ "-conv2") ( + #data := act2 .& + #kernel := [3,3] .& + #num_filter := (num_filter `div` 4) .& + #stride := stride .& + #pad := [1,1] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn3 <- batchnorm (name ++ "-bn3") ( + #data := conv2 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act3 <- activation (name ++ "-relu3") ( + #data := bn3 .& + #act_type := #relu .& Nil) + conv3 <- convolution (name ++ "-conv3") ( + #data := act3 .& + #kernel := [1,1] .& + #num_filter := num_filter .& + #stride := [1,1] .& + #pad := [0,0] .& + #workspace := workspace .& + #no_bias := True .& Nil) + shortcut <- if dim_match + then return dat + else convolution (name ++ "-sc") ( + #data := act1 .& + #kernel := [1,1] .& + #num_filter := num_filter .& + #stride := stride .& + #workspace := workspace .& + #no_bias := True .& Nil) + when memonger $ + void $ mxSymbolSetAttr shortcut "mirror_stage" "true" + plus name (#lhs := conv3 .& #rhs := shortcut .& Nil) + else do + bn1 <- batchnorm (name ++ "-bn1") ( + #data := dat .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act1 <- activation (name ++ "-relu1") ( + #data := bn1 .& + #act_type := #relu .& Nil) + conv1 <- convolution (name ++ "-conv1") ( + #data := act1 .& + #kernel := [3,3] .& + #num_filter:= num_filter .& + #stride := stride .& + #pad := [1,1] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn2 <- batchnorm (name ++ "-bn2") ( + #data := conv1 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act2 <- activation (name ++ "-relu2") ( + #data := bn2 .& + #act_type := #relu .& Nil) + conv2 <- convolution (name ++ "-conv2") ( + #data := act2 .& + #kernel := [3,3] .& + #num_filter:= num_filter .& + #stride := [1,1] .& + #pad := [1,1] .& + #workspace := workspace .& + #no_bias := True .& Nil) + shortcut <- if dim_match + then return dat + else convolution (name ++ "-sc") ( + #data := act1 .& + #kernel := [1,1] .& + #num_filter:= num_filter .& + #stride := stride .& + #workspace := workspace.& + #no_bias := True .& Nil) + when memonger $ + void $ mxSymbolSetAttr shortcut "mirror_stage" "true" + plus name (#lhs := conv2 .& #rhs := shortcut .& Nil) diff --git a/src/Model/Resnext.hs b/src/Model/Resnext.hs new file mode 100644 --- /dev/null +++ b/src/Model/Resnext.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +module Model.Resnext where + +import Control.Monad (foldM, when, void) +import Data.Maybe (fromMaybe) + +import MXNet.Base +import MXNet.NN.Layer + +-- ResNet +-- #layer: 164 +-- #stage: 3 +-- #layer per stage: 18 +-- #filter of stage 1: 64 +-- #filter of stage 2: 128 +-- #filter of stage 3: 256 + +symbol :: DType a => IO (Symbol a) +symbol = do + x <- variable "x" + y <- variable "y" + + xcp <- identity "id" ( + #data := x .& Nil) + + bnx <- batchnorm "bn-x" ( + #data := xcp .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := True .& Nil) + + cvx <- convolution "conv-bn-x" ( + #data := bnx .& + #kernel := [3,3] .& + #num_filter := 16 .& + #stride := [1,1] .& + #pad := [1,1] .& + #workspace := conv_workspace .& + #no_bias := True .& Nil) + + bdy <- foldM (\layer (num_filter, stride, dim_match, name) -> + residual name (#data := layer .& + #num_filter := num_filter .& + #stride := stride .& + #dim_match := dim_match .& resargs)) + cvx + residual'parms + + pool1 <- pooling "pool1" ( + #data := bdy .& + #kernel := [7,7] .& + #pool_type := #avg .& + #global_pool := True .& Nil) + flat <- flatten "flat-1" ( + #data := pool1 .& Nil) + fc1 <- fullyConnected "fc-1" ( + #data := flat .& + #num_hidden := 10 .& Nil) + Symbol <$> softmaxoutput "softmax" ( + #data := fc1 .& + #label := y .& Nil) + where + bn_mom = 0.9 :: Float + conv_workspace = 256 :: Int + eps = 2e-5 :: Double + residual'parms = [ (64, [1,1], False, "stage1-unit1") ] ++ map (\i -> (64, [1,1], True, "stage1-unit" ++ show i)) [2..18 :: Int] + ++ [ (128, [2,2], False, "stage2-unit1") ] ++ map (\i -> (128, [1,1], True, "stage2-unit" ++ show i)) [2..18 :: Int] + ++ [ (256, [2,2], False, "stage3-unit1") ] ++ map (\i -> (256, [1,1], True, "stage3-unit" ++ show i)) [2..18 :: Int] + resargs = #bottle_neck := True .& #workspace := conv_workspace .& #memonger := False .& Nil + +type instance ParameterList "_residual_layer(resnext)" = + '[ '("data" , 'AttrReq SymbolHandle) + , '("num_filter" , 'AttrReq Int) + , '("stride" , 'AttrReq [Int]) + , '("dim_match" , 'AttrReq Bool) + , '("bottle_neck", 'AttrOpt Bool) + , '("num_group" , 'AttrOpt Int) + , '("bn_mom" , 'AttrOpt Float) + , '("workspace" , 'AttrOpt Int) + , '("memonger" , 'AttrOpt Bool) ] +residual :: (Fullfilled "_residual_layer(resnext)" args) + => String -> ArgsHMap "_residual_layer(resnext)" args -> IO SymbolHandle +residual name args = do + let dat = args ! #data + num_filter = args ! #num_filter + stride = args ! #stride + dim_match = args ! #dim_match + bottle_neck= fromMaybe True $ args !? #bottle_neck + num_group = fromMaybe 32 $ args !? #num_group + bn_mom = fromMaybe 0.9 $ args !? #bn_mom + workspace = fromMaybe 256 $ args !? #workspace + memonger = fromMaybe False$ args !? #memonger + eps = 2e-5 :: Double + if bottle_neck + then do + conv1 <- convolution (name ++ "-conv1") ( + #data := dat .& + #kernel := [1,1] .& + #num_filter:= num_filter `div` 2 .& + #stride := [1,1] .& + #pad := [0,0] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn1 <- batchnorm (name ++ "-bn1") ( + #data := conv1 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act1 <- activation (name ++ "-relu1") ( + #data := bn1 .& + #act_type := #relu .& Nil) + conv2 <- convolution (name ++ "-conv2") ( + #data := act1 .& + #kernel := [3,3] .& + #num_filter:= num_filter `div` 2 .& + #stride := stride .& + #pad := [1,1] .& + #num_group := num_group .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn2 <- batchnorm (name ++ "-bn2") ( + #data := conv2 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act2 <- activation (name ++ "-relu2") ( + #data := bn2 .& + #act_type := #relu .& Nil) + conv3 <- convolution (name ++ "-conv3") ( + #data := act2 .& + #kernel := [1,1] .& + #num_filter:= num_filter .& + #stride := [1,1] .& + #pad := [0,0] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn3 <- batchnorm (name ++ "-bn3") ( + #data := conv3 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + shortcut <- if dim_match + then return dat + else do + shortcut_conv <- convolution (name ++ "-sc") ( + #data := dat .& + #kernel := [1,1] .& + #num_filter := num_filter .& + #stride := stride .& + #workspace := workspace .& + #no_bias := True .& Nil) + batchnorm (name ++ "-sc-bn") ( + #data := shortcut_conv .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + when memonger $ + void $ mxSymbolSetAttr shortcut "mirror_stage" "true" + eltwise <- plus name ( + #lhs := bn3 .& + #rhs := shortcut .& Nil) + activation (name ++ "-relu") ( + #data := eltwise .& + #act_type := #relu .& Nil) + else do + conv1 <- convolution (name ++ "-conv1") ( + #data := dat .& + #kernel := [3,3] .& + #num_filter := num_filter .& + #stride := stride .& + #pad := [1,1] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn1 <- batchnorm (name ++ "-bn1") ( + #data := conv1 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + act1 <- activation (name ++ "-relu1") ( + #data := bn1 .& + #act_type := #relu .& Nil) + conv2 <- convolution (name ++ "-conv2") ( + #data := act1 .& + #kernel := [3,3] .& + #num_filter := num_filter .& + #stride := [1,1] .& + #pad := [1,1] .& + #workspace := workspace .& + #no_bias := True .& Nil) + bn2 <- batchnorm (name ++ "-bn2") ( + #data := conv2 .& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + shortcut <- if dim_match + then return dat + else do + shortcut_conv <- convolution (name ++ "-sc") ( + #data := act1 .& + #kernel := [1,1] .& + #num_filter := num_filter .& + #stride := stride .& + #workspace := workspace .& + #no_bias := True .& Nil) + batchnorm (name ++ "-sc-bn") ( + #data := shortcut_conv.& + #eps := eps .& + #momentum := bn_mom .& + #fix_gamma := False .& Nil) + when memonger $ + void $ mxSymbolSetAttr shortcut "mirror_stage" "true" + eltwise <- plus name ( + #lhs := bn2 .& + #rhs := shortcut .& Nil) + activation (name ++ "-relu") ( + #data := eltwise .& + #act_type := #relu .& Nil) diff --git a/src/Model/VGG.hs b/src/Model/VGG.hs new file mode 100644 --- /dev/null +++ b/src/Model/VGG.hs @@ -0,0 +1,66 @@ +module Model.VGG where + +import Text.Printf (printf) +import Control.Monad (foldM) +import Data.Maybe (fromMaybe) + +import MXNet.Base +import MXNet.NN.Layer + +getFeature :: SymbolHandle -> [Int] -> [Int] -> Bool -> Bool -> IO SymbolHandle +getFeature internalLayer layers filters with_batch_norm with_last_pooling= do + sym <- foldM build1 internalLayer specs + -- inlining the build1 below, and omit pooling depending on the with_last_pooling + case last_group of + (idx, num, filter) -> do + sym <- foldM (build2 idx) sym $ zip [1::Int ..] (replicate num filter) + if not with_last_pooling + then return sym + else pooling (printf "pool%d" idx) (#data := sym .& #pool_type := #max .& #kernel := [2,2] .& #stride := [2,2] .& Nil) + + where + last_group:groups = reverse $ zip3 [1::Int ..] layers filters + specs = reverse groups + + build1 sym (idx, num, filter) = do + sym <- foldM (build2 idx) sym $ zip [1::Int ..] (replicate num filter) + pooling (printf "pool%d" idx) (#data := sym .& #pool_type := #max .& #kernel := [2,2] .& #stride := [2,2] .& Nil) + + build2 idx1 sym (idx2, filter) = do + let ident = printf "%d_%d" idx1 idx2 + sym <- convolution ("conv" ++ ident) (#data := sym .& #kernel := [3,3] .& #pad := [1,1] .& #num_filter := filter .& #workspace := 2048 .& Nil) + sym <- if with_batch_norm then batchnorm ("bn" ++ ident) (#data := sym .& Nil) else return sym + activation ("relu" ++ ident) (#data := sym .& #act_type := #relu .& Nil) + +getTopFeature :: Maybe String -> SymbolHandle -> IO SymbolHandle +getTopFeature prefix input_data = do + let addPrefix = (fromMaybe "" prefix ++) + sym <- flatten (addPrefix "flatten") (#data := input_data .& Nil) + sym <- fullyConnected (addPrefix "fc6") (#data := sym .& #num_hidden := 4096 .& Nil) + sym <- activation (addPrefix "relu6") (#data := sym .& #act_type := #relu .& Nil) + sym <- dropout (addPrefix "drop6") (#data := sym .& #p := 0.5 .& Nil) + sym <- fullyConnected (addPrefix "fc7") (#data := sym .& #num_hidden := 4096 .& Nil) + sym <- activation (addPrefix "relu7") (#data := sym .& #act_type := #relu .& Nil) + dropout (addPrefix "drop7") (#data := sym .& #p := 0.5 .& Nil) + +getClassifier :: Maybe String -> SymbolHandle -> Int -> IO SymbolHandle +getClassifier prefix input_data num_classes = do + let addPrefix = (fromMaybe "" prefix ++) + sym <- getTopFeature prefix input_data + fullyConnected (addPrefix "fc8") (#data := sym .& #num_hidden := num_classes .& Nil) + +symbol :: Int -> Int -> Bool -> IO (Symbol Float) +symbol num_classes num_layers with_batch_norm = do + sym <- variable "data" + sym <- getFeature sym layers filters with_batch_norm True + sym <- getClassifier Nothing sym num_classes + sym <- softmaxoutput "softmax" (#data := sym .& Nil) + return (Symbol sym) + + where + (layers, filters) = case num_layers of + 11 -> ([1, 1, 2, 2, 2], [64, 128, 256, 512, 512]) + 13 -> ([2, 2, 2, 2, 2], [64, 128, 256, 512, 512]) + 16 -> ([2, 2, 3, 3, 3], [64, 128, 256, 512, 512]) + 19 -> ([2, 2, 4, 4, 4], [64, 128, 256, 512, 512]) + diff --git a/src/cifar10.hs b/src/cifar10.hs new file mode 100644 --- /dev/null +++ b/src/cifar10.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +module Main where + +import qualified Data.HashMap.Strict as M +import Control.Monad (forM_, void) +import qualified Data.Vector.Storable as SV +import Control.Monad.IO.Class +import Control.Lens ((%~)) +import System.IO (hFlush, stdout) +import Options.Applicative (Parser, execParser, header, info, fullDesc, helper, value, option, auto, metavar, short, showDefault, (<**>)) +import Data.Semigroup ((<>)) + +import MXNet.Base (NDArray(..), contextCPU, contextGPU0, mxListAllOpNames, toVector, (.&), HMap(..), ArgOf(..)) +import qualified MXNet.Base.Operators.NDArray as A +import MXNet.NN +import MXNet.NN.Utils +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Streaming +import qualified Model.Resnet as Resnet +import qualified Model.Resnext as Resnext + +type ArrayF = NDArray Float +type DS = StreamData (TrainM Float IO) (ArrayF, ArrayF) + +data Model = Resnet | Resnext deriving (Show, Read) +data ProgArg = ProgArg Model +cmdArgParser :: Parser ProgArg +cmdArgParser = ProgArg <$> (option auto $ short 'm' <> metavar "MODEL" <> showDefault <> value Resnet) + +range :: Int -> [Int] +range = enumFromTo 1 + +default_initializer :: Initializer Float +default_initializer name shp + | endsWith "-bias" name = zeros name shp + | endsWith "-beta" name = zeros name shp + | endsWith "-gamma" name = ones name shp + | endsWith "-moving-mean" name = zeros name shp + | endsWith "-moving-var" name = ones name shp + | otherwise = case shp of + [_,_] -> xavier 2.0 XavierGaussian XavierIn name shp + _ -> normal 0.1 name shp + +main :: IO () +main = do + ProgArg model <- execParser $ info (cmdArgParser <**> helper) (fullDesc <> header "CIFAR-10 solver") + -- call mxListAllOpNames can ensure the MXNet itself is properly initialized + -- i.e. MXNet operators are registered in the NNVM + _ <- mxListAllOpNames + net <- case model of + Resnet -> Resnet.symbol 10 34 [3,32,32] + Resnext -> Resnext.symbol + sess <- initialize net $ Config { + _cfg_data = M.singleton "x" [3,32,32], + _cfg_label = ["y"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_context = contextGPU0 + } + + cbTP <- dumpThroughputEpoch + sess <- return $ (sess_callbacks %~ ([Callback DumpLearningRate, cbTP, Callback (Checkpoint "tmp")] ++)) sess + + optimizer <- makeOptimizer ADAM (lrOfPoly $ #maxnup := 10000 .& #base := 0.05 .& #power := 1 .& Nil) Nil + + train sess $ do + + let trainingData = imageRecordIter (#path_imgrec := "data/cifar10_train.rec" .& + #data_shape := [3,32,32] .& + #batch_size := 128 .& Nil) + let testingData = imageRecordIter (#path_imgrec := "data/cifar10_val.rec" .& + #data_shape := [3,32,32] .& + #batch_size := 32 .& Nil) + fitDataset trainingData testingData bind optimizer (CrossEntropy "y" :* Accuracy "y" :* MNil) 18 + + where + bind ["x", "y"] (dat, lbl) = M.fromList [("x", dat), ("y", lbl)] diff --git a/src/custom-op.hs b/src/custom-op.hs new file mode 100644 --- /dev/null +++ b/src/custom-op.hs @@ -0,0 +1,152 @@ +module Main where + +import MXNet.Base +import qualified MXNet.Base.Operators.NDArray as A +import qualified MXNet.Base.Operators.Symbol as S +import qualified MXNet.NN as NN +import qualified MXNet.NN.Utils as NN +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Streaming +import qualified Data.HashMap.Strict as M +import qualified Data.Vector.Storable as SV +import Control.Monad.IO.Class +import Control.Monad (forM_, void) +import System.IO (hFlush, stdout) + +type ArrayF = NDArray Float + +data SoftmaxProp = SoftmaxProp + +instance CustomOperationProp SoftmaxProp where + prop_list_arguments _ = ["data", "label"] + prop_list_outputs _ = ["output"] + prop_list_auxiliary_states _ = [] + prop_infer_shape _ [data_shape, _] = + let output_shape = data_shape + in ([data_shape, [head data_shape]], [output_shape], []) + prop_declare_backward_dependency _ grad_out data_in data_out = data_in ++ data_out + + data Operation SoftmaxProp = Softmax + prop_create_operator _ _ _ = return Softmax + +instance CustomOperation (Operation SoftmaxProp) where + forward _ [ReqWrite] [in_data,_] [out_data] aux is_train = do + -- let in_data_ = (NDArray in_data :: ArrayF) + -- [_, num_classes] <- ndshape in_data_ + -- vec <- toVector in_data_ + -- let batch_exp = L.toRows $ exp $ L.reshape num_classes vec :: [L.Vector Float] + -- norm1 = map (realToFrac . L.sumElements) $ batch_exp :: [L.Vector Float] + -- output = L.fromRows $ zipWith (/) batch_exp norm1 + -- copyFromVector (NDArray out_data :: ArrayF) vec + + [result] <- A.softmax (#data := in_data .& #axis := 1 .& Nil) + A._copyto_upd [out_data] (#data := result .& Nil) + + backward _ [ReqWrite] [_, label] [out_data] [in_grad, _] _ aux = do + -- let out_data_ = NDArray out_data :: ArrayF + -- label_ = NDArray label :: ArrayF + -- out_shp@[_, num_classes] <- ndshape out_data_ + -- vec_lbl <- toVector label_ + -- vec_out <- toVector out_data_ + -- let rows = L.toRows $ L.reshape num_classes vec_out :: [L.Vector Float] + -- upd :: L.Vector Float -> Float -> L.Vector Float + -- upd row n = let n_ = floor n + -- in row SV.// [(n_, row SV.! n_ - 1)] + -- result = L.fromRows $ zipWith upd rows (L.toList vec_lbl) :: L.Matrix Float + -- copyFromVector (NDArray in_grad :: ArrayF) (L.flatten result) + + out_shp@[_, num_classes] <- ndshape (NDArray out_data :: ArrayF) + [label_onehot] <- A.one_hot (#indices := label .& #depth := num_classes .& Nil) + [result] <- A.elemwise_sub (#lhs := out_data .& #rhs := label_onehot .& Nil) + A._copyto_upd [in_grad] (#data := result .& Nil) + + +symbol :: DType a => IO (Symbol a) +symbol = do + x <- NN.variable "x" + y <- NN.variable "y" + + v1 <- NN.convolution "conv1" (#data := x .& #kernel := [5,5] .& #num_filter := 20 .& Nil) + a1 <- NN.activation "conv1-a" (#data := v1 .& #act_type := #tanh .& Nil) + p1 <- NN.pooling "conv1-p" (#data := a1 .& #kernel := [2,2] .& #pool_type := #max .& Nil) + + v2 <- NN.convolution "conv2" (#data := p1 .& #kernel := [5,5] .& #num_filter := 50 .& Nil) + a2 <- NN.activation "conv2-a" (#data := v2 .& #act_type := #tanh .& Nil) + p2 <- NN.pooling "conv2-p" (#data := a2 .& #kernel := [2,2] .& #pool_type := #max .& Nil) + + fl <- NN.flatten "flatten" (#data := p2 .& Nil) + + v3 <- NN.fullyConnected "fc1" (#data := fl .& #num_hidden := 500 .& Nil) + a3 <- NN.activation "fc1-a" (#data := v3 .& #act_type := #tanh .& Nil) + + v4 <- NN.fullyConnected "fc2" (#data := a3 .& #num_hidden := 10 .& Nil) + a4 <- S._Custom "softmax" (#data := [v4, y] .& #op_type := "softmax_custom" .& Nil) + return $ Symbol a4 + +default_initializer :: NN.Initializer Float +default_initializer name shp + | NN.endsWith "-bias" name = NN.zeros name shp + | otherwise = NN.normal 0.1 name shp + +main :: IO () +main = do + -- call mxListAllOpNames can ensure the MXNet itself is properly initialized + -- i.e. MXNet operators are registered in the NNVM + _ <- mxListAllOpNames + registerCustomOperator ("softmax_custom", \_ -> return SoftmaxProp) + net <- symbol + + sess <- NN.initialize net $ NN.Config { + NN._cfg_data = M.singleton "x" [1,28,28], + NN._cfg_label = ["y"], + NN._cfg_initializers = M.empty, + NN._cfg_default_initializer = default_initializer, + NN._cfg_context = contextCPU + } + optimizer <- NN.makeOptimizer NN.SGD'Mom (NN.Const 0.0002) Nil + + NN.train sess $ do + + let trainingData = mnistIter (#image := "data/train-images-idx3-ubyte" .& + #label := "data/train-labels-idx1-ubyte" .& + #batch_size := 128 .& Nil) + let testingData = mnistIter (#image := "data/t10k-images-idx3-ubyte" .& + #label := "data/t10k-labels-idx1-ubyte" .& + #batch_size := 16 .& Nil) + + total1 <- sizeD trainingData + liftIO $ putStrLn $ "[Train] " + forM_ (enumFromTo 1 20 :: [Int]) $ \ind -> do + liftIO $ putStrLn $ "iteration " ++ show ind + metric <- NN.newMetric "train" (NN.CrossEntropy "y") + void $ forEachD_i trainingData $ \(i, (x, y)) -> do + NN.fitAndEval optimizer (M.fromList [("x", x), ("y", y)]) metric + eval <- NN.format metric + liftIO $ do + putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total1 ++ " " ++ eval + hFlush stdout + liftIO $ putStrLn "" + + liftIO $ putStrLn $ "[Test] " + + total2 <- sizeD testingData + result <- forEachD_i testingData $ \(i, (x, y)) -> do + liftIO $ do + putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total2 + hFlush stdout + [y'] <- NN.forwardOnly (M.fromList [("x", Just x), ("y", Nothing)]) + ind1 <- liftIO $ toVector y + ind2 <- liftIO $ argmax y' >>= toVector + return (ind1, ind2) + liftIO $ putStr "\r\ESC[K" + + let (ls,ps) = unzip result + ls_unbatched = mconcat ls + ps_unbatched = mconcat ps + total_test_items = SV.length ls_unbatched + correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched + liftIO $ putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total_test_items + + where + argmax :: ArrayF -> IO ArrayF + argmax (NDArray ys) = NDArray . head <$> A.argmax (#data := ys .& #axis := Just 1 .& Nil) diff --git a/src/lenet.hs b/src/lenet.hs new file mode 100644 --- /dev/null +++ b/src/lenet.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +module Main where + +import qualified Data.HashMap.Strict as M +import Control.Monad (forM_, void) +import qualified Data.Vector.Storable as SV +import Control.Monad.IO.Class +import System.IO (hFlush, stdout) + +import MXNet.Base (NDArray(..), contextCPU, contextGPU0, mxListAllOpNames, toVector, (.&), HMap(..), ArgOf(..), waitAll) +import qualified MXNet.Base.Operators.NDArray as A +import MXNet.NN +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Conduit +import qualified Model.Lenet as Model + +type ArrayF = NDArray Float +type DS = ConduitData (TrainM Float IO) (ArrayF, ArrayF) + +range :: Int -> [Int] +range = enumFromTo 1 + +default_initializer :: Initializer Float +default_initializer name shp@[_] = zeros name shp +default_initializer name shp@[_,_] = xavier 2.0 XavierGaussian XavierIn name shp +default_initializer name shp = normal 0.1 name shp + +main :: IO () +main = do + -- call mxListAllOpNames can ensure the MXNet itself is properly initialized + -- i.e. MXNet operators are registered in the NNVM + _ <- mxListAllOpNames + net <- Model.symbol + sess <- initialize net $ Config { + _cfg_data = M.singleton "x" [1,28,28], + _cfg_label = ["y"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_context = contextCPU + } + optimizer <- makeOptimizer SGD'Mom (Const 0.0002) Nil + + train sess $ do + + let trainingData = mnistIter (#image := "data/train-images-idx3-ubyte" .& + #label := "data/train-labels-idx1-ubyte" .& + #batch_size := 128 .& Nil) + let testingData = mnistIter (#image := "data/t10k-images-idx3-ubyte" .& + #label := "data/t10k-labels-idx1-ubyte" .& + #batch_size := 16 .& Nil) + + total1 <- sizeD trainingData + total2 <- sizeD testingData + + liftIO $ putStrLn $ "[Train] " + forM_ (range 1) $ \ind -> do + liftIO $ putStrLn $ "iteration " ++ show ind + -- metric <- newMetric "train" (CrossEntropy "y") + metric <- newMetric "train" MNil + void $ forEachD_i trainingData $ \(i, (x, y)) -> do + -- liftIO $ putStrLn "A" + fitAndEval optimizer (M.fromList [("x", x), ("y", y)]) metric + -- liftIO $ putStrLn "B" + eval <- format metric + liftIO $ do + putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total1 ++ " " ++ eval + hFlush stdout + -- putStrLn "C" + waitAll + liftIO $ putStrLn "D" + + metric <- newMetric "val" (Accuracy "y") + result <- forEachD_i testingData $ \(i, (x, y)) -> do + pred <- forwardOnly (M.fromList [("x", Just x), ("y", Nothing)]) + evaluate metric (M.singleton "y" y) pred + eval <- format metric + liftIO $ do + putStr $ "\r\ESC[K" ++ show i ++ "/" ++ show total2 ++ " " ++ eval + hFlush stdout + liftIO $ putStrLn "" + + -- let (ls,ps) = unzip result + -- ls_unbatched = mconcat ls + -- ps_unbatched = mconcat ps + -- total_test_items = SV.length ls_unbatched + -- correct = SV.length $ SV.filter id $ SV.zipWith (==) ls_unbatched ps_unbatched + -- liftIO $ putStrLn $ "Accuracy: " ++ show correct ++ "/" ++ show total_test_items + +-- where +-- argmax :: ArrayF -> IO ArrayF +-- argmax (NDArray ys) = NDArray . head <$> A.argmax (#data := ys .& #axis := Just 1 .& Nil) diff --git a/src/rcnn.hs b/src/rcnn.hs new file mode 100644 --- /dev/null +++ b/src/rcnn.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RecordWildCards #-} +module Main where + +import qualified Data.HashMap.Strict as M +import Control.Monad (forM_, void, unless) +import Control.Applicative (liftA2) +import qualified Data.Vector.Storable as SV +import Control.Monad.IO.Class +import Control.Lens ((.=)) +import System.IO (hFlush, stdout) +import Options.Applicative ( + Parser, execParser, + long, value, option, auto, strOption, metavar, showDefault, eitherReader, help, + info, helper, fullDesc, header, (<**>)) +import Data.Attoparsec.Text (sepBy, char, rational, decimal, endOfInput, parseOnly) +import qualified Data.Text as T +import System.Directory (doesFileExist, canonicalizePath) + +import MXNet.Base ( + NDArray(..), toVector, + contextCPU, contextGPU0, + mxListAllOpNames, mxNotifyShutdown, mxNDArraySave, + registerCustomOperator, + ndshape, + listOutputs, internals, inferShape, at', at, + HMap(..), (.&), ArgOf(..)) +import MXNet.NN +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Conduit +import MXNet.NN.DataIter.Coco as Coco +import MXNet.NN.Utils (loadSession, saveSession) +import Model.FasterRCNN + + +data CocoConfig = CocoConfig { + coco_base_path :: String, + coco_img_short_side :: Int, + coco_img_long_side :: Int, + coco_img_pixel_means :: [Float], + coco_img_pixel_stds :: [Float] +} deriving Show + +cmdArgParser :: Parser (RcnnConfiguration, CocoConfig) +cmdArgParser = liftA2 (,) + (RcnnConfiguration + <$> option intList (long "rpn-anchor-scales" <> metavar "SCALES" <> showDefault <> value [8,16,32] <> help "rpn anchor scales") + <*> option floatList (long "rpn-anchor-ratios" <> metavar "RATIOS" <> showDefault <> value [0.5,1,2] <> help "rpn anchor ratios") + <*> option auto (long "rpn-feat-stride" <> metavar "STRIDE" <> showDefault <> value 16 <> help "rpn feature stride") + <*> option auto (long "rpn-batch-rois" <> metavar "BATCH-ROIS" <> showDefault <> value 256 <> help "rpn number of rois per batch") + <*> option auto (long "rpn-pre-nms-topk" <> metavar "PRE-NMS-TOPK" <> showDefault <> value 12000 <> help "rpn nms pre-top-k") + <*> option auto (long "rpn-post-nms-topk" <> metavar "POST-NMS-TOPK" <> showDefault <> value 2000 <> help "rpn nms post-top-k") + <*> option auto (long "rpn-nms-thresh" <> metavar "NMS-THRESH" <> showDefault <> value 0.7 <> help "rpn nms threshold") + <*> option auto (long "rpn-min-size" <> metavar "MIN-SIZE" <> showDefault <> value 16 <> help "rpn min size") + <*> option auto (long "rpn-fg-fraction" <> metavar "FG-FRACTION" <> showDefault <> value 0.5 <> help "rpn foreground fraction") + <*> option auto (long "rpn-fg-overlap" <> metavar "FG-OVERLAP" <> showDefault <> value 0.7 <> help "rpn foreground iou threshold") + <*> option auto (long "rpn-bg-overlap" <> metavar "BG-OVERLAP" <> showDefault <> value 0.3 <> help "rpn background iou threshold") + <*> option auto (long "rpn-allowed-border"<> metavar "ALLOWED-BORDER" <> showDefault <> value 0 <> help "rpn allowed border") + <*> option auto (long "rcnn-num-classes" <> metavar "NUM-CLASSES" <> showDefault <> value 81 <> help "rcnn number of classes") + <*> option auto (long "rcnn-feat-stride" <> metavar "FEATURE-STRIDE" <> showDefault <> value 16 <> help "rcnn feature stride") + <*> option intList (long "rcnn-pooled-size" <> metavar "POOLED-SIZE" <> showDefault <> value [7,7] <> help "rcnn pooled size") + <*> option auto (long "rcnn-batch-rois" <> metavar "BATCH_ROIS" <> showDefault <> value 128 <> help "rcnn batch rois") + <*> option auto (long "rcnn-batch-size" <> metavar "BATCH-SIZE" <> showDefault <> value 1 <> help "rcnn batch size") + <*> option auto (long "rcnn-fg-fraction" <> metavar "FG-FRACTION" <> showDefault <> value 0.25 <> help "rcnn foreground fraction") + <*> option auto (long "rcnn-fg-overlap" <> metavar "FG-OVERLAP" <> showDefault <> value 0.5 <> help "rcnn foreground iou threshold") + <*> option floatList (long "rcnn-bbox-stds" <> metavar "BBOX-STDDEV" <> showDefault <> value [0.1, 0.1, 0.2, 0.2] <> help "standard deviation of bbox") + <*> strOption (long "pretrained" <> metavar "PATH" <> value "" <> help "path to pretrained model")) + (CocoConfig + <$> strOption (long "coco" <> metavar "PATH" <> help "path to the coco dataset") + <*> option auto (long "img-short-side" <> metavar "SIZE" <> showDefault <> value 600 <> help "short side of image") + <*> option auto (long "img-long-side" <> metavar "SIZE" <> showDefault <> value 1000 <> help "long side of image") + <*> option floatList (long "img-pixel-means" <> metavar "RGB-MEAN" <> showDefault <> value [0,0,0] <> help "RGB mean of images") + <*> option floatList (long "img-pixel-stds" <> metavar "RGB-STDS" <> showDefault <> value [1,1,1] <> help "RGB std-dev of images")) + where + list obj = parseOnly (sepBy obj (char ',') <* endOfInput) . T.pack + floatList = eitherReader $ list rational + intList = eitherReader $ list decimal + +buildProposalTargetProp params = do + let params' = M.fromList params + return $ ProposalTargetProp { + _num_classes = read $ params' M.! "num_classes", + _batch_images= read $ params' M.! "batch_images", + _batch_rois = read $ params' M.! "batch_rois", + _fg_fraction = read $ params' M.! "fg_fraction", + _fg_overlap = read $ params' M.! "fg_overlap", + _box_stds = read $ params' M.! "box_stds" + } + +toTriple [a, b, c] = (a, b, c) +toTriple x = error (show x) + + +default_initializer :: Initializer Float +default_initializer name = case name of + "rpn_conv_3x3_weight" -> normal 0.01 name + "rpn_conv_3x3_bias" -> zeros name + "rpn_cls_score_weight" -> normal 0.01 name + "rpn_cls_score_bias" -> zeros name + "rpn_bbox_pred_weight" -> normal 0.01 name + "rpn_bbox_pred_bias" -> zeros name + "cls_score_weight" -> normal 0.01 name + "cls_score_bias" -> zeros name + "bbox_pred_weight" -> normal 0.001 name + "bbox_pred_bias" -> zeros name + _ -> empty name + +loadWeights weights_path = do + weights_path <- liftIO $ canonicalizePath weights_path + e <- liftIO $ doesFileExist (weights_path ++ ".params") + if not e + then liftIO $ putStrLn $ "'" ++ weights_path ++ ".params' doesn't exist." + else loadSession weights_path ["rpn_conv_3x3_weight", + "rpn_conv_3x3_bias", + "rpn_cls_score_weight", + "rpn_cls_score_bias", + "rpn_bbox_pred_weight", + "rpn_bbox_pred_bias", + "cls_score_weight", + "cls_score_bias", + "bbox_pred_weight", + "bbox_pred_bias"] + +main :: IO () +main = do + _ <- mxListAllOpNames + registerCustomOperator ("proposal_target", buildProposalTargetProp) + (rcnn_conf@RcnnConfiguration{..}, CocoConfig{..}) <- execParser $ info (cmdArgParser <**> helper) (fullDesc <> header "Faster-RCNN") + sym <- symbolTrain rcnn_conf + + rpn_cls_score_output <- internals sym >>= flip at' "rpn_cls_score_output" + let extr_feature_shape (w, h) = do + -- get the feature (width, height) at the top of feature extraction. + (_, [(_, [_, _, feat_width, feat_height])], _, _) <- inferShape rpn_cls_score_output [("data", [1, 3,w, h])] + return (feat_width, feat_height) + + -- Coco x y inst <- coco coco_base_path "train2017" + -- let coco_inst = Coco x y (inst & images %~ V.filter (\img_desc -> (img_desc ^. img_id) `elem` [97733])) -- , 123980, 111549])) + -- let data_iter = cocoImagesWithAnchors' (cocoImages coco_inst False) extr_feature_shape + -- (#anchor_scales := rpn_anchor_scales + -- .& #anchor_ratios := rpn_anchor_ratios + -- .& #batch_rois := rpn_batch_rois + -- .& #feature_stride:= rpn_feature_stride + -- .& #allowed_border:= rpn_allowd_border + -- .& #fg_fraction := rpn_fg_fraction + -- .& #fg_overlap := rpn_fg_overlap + -- .& #bg_overlap := rpn_bg_overlap + -- .& #short_size := coco_img_short_side + -- .& #long_size := coco_img_long_side + -- .& #mean := toTriple coco_img_pixel_means + -- .& #std := toTriple coco_img_pixel_stds + -- .& #batch_size := rcnn_batch_size + -- .& Nil) + + coco_inst <- coco coco_base_path "train2017" + let data_iter = cocoImagesWithAnchors coco_inst extr_feature_shape + (#anchor_scales := rpn_anchor_scales + .& #anchor_ratios := rpn_anchor_ratios + .& #batch_rois := rpn_batch_rois + .& #feature_stride:= rpn_feature_stride + .& #allowed_border:= rpn_allowd_border + .& #fg_fraction := rpn_fg_fraction + .& #fg_overlap := rpn_fg_overlap + .& #bg_overlap := rpn_bg_overlap + .& #short_size := coco_img_short_side + .& #long_size := coco_img_long_side + .& #mean := toTriple coco_img_pixel_means + .& #std := toTriple coco_img_pixel_stds + .& #batch_size := rcnn_batch_size + .& #shuffle := True + .& Nil) + + sess <- initialize sym $ Config { + _cfg_data = M.fromList [("data", [3, coco_img_short_side, coco_img_long_side]), + ("im_info", [3]), + ("gt_boxes", [0, 5])], + _cfg_label = ["label", "bbox_target", "bbox_weight"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_context = contextGPU0 + } + optimizer <- makeOptimizer SGD'Mom (Const 0.001) (#momentum := 0.9 + .& #wd := 0.0005 + .& #rescale_grad := 1 / (fromIntegral rcnn_batch_size) + .& #clip_gradient := 5 + .& Nil) + + train sess $ do + sess_callbacks .= [Callback DumpLearningRate, Callback (Checkpoint "checkpoints")] + + unless (null pretrained_weights) (loadWeights pretrained_weights) + + metric <- newMetric "train" (RPNAccMetric 0 "label" :* RCNNAccMetric 2 4 :* RPNLogLossMetric 0 "label" :* RCNNLogLossMetric 2 4 :* RPNL1LossMetric 1 "bbox_weight" :* RCNNL1LossMetric 3 4 :* MNil) + forM_ [1..40] $ \ ei -> do + liftIO $ putStrLn $ "Epoch " ++ show ei + liftIO $ hFlush stdout + let dd = takeD 200 data_iter + void $ forEachD_i dd $ \(i, ((x0, x1, x2), (y0, y1, y2))) -> do + let binding = M.fromList [ ("data", x0) + , ("im_info", x1) + , ("gt_boxes", x2) + , ("label", y0) + , ("bbox_target", y1) + , ("bbox_weight", y2) ] + fitAndEval optimizer binding metric + eval <- format metric + liftIO $ do + putStrLn $ show i ++ " " ++ eval + hFlush stdout + liftIO $ putStrLn "" + saveSession "sav" + + -- CUDA.stop + mxNotifyShutdown