diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -4,3 +4,4 @@ + CIFAR10 + Resnet / ResNext + mxnet custom operator + Faster RCNN + + `LD_LIBRARY_PATH= stack run faster-rcnn -- --backbone RESNET50FPN --strides [4,8,16,32] --pretrained params/resnet50_v2 --base --img-size 512 --img-pixel-means [0.5,0.5,0.5] --train-epochs 20 --train-iter-per-epoch 300 --batch-size=1 --rcnn-batch-rois=256 +RTS -N6` diff --git a/fei-examples.cabal b/fei-examples.cabal --- a/fei-examples.cabal +++ b/fei-examples.cabal @@ -1,17 +1,17 @@ +cabal-version: 2.2 name: fei-examples -version: 0.3.0 +version: 1.0.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 +copyright: 2020 - Jiasen Wu +license: BSD-3-Clause license-file: LICENSE category: Machine Learning, AI build-type: Simple -cabal-version: >= 1.10 extra-source-files: README.md @@ -20,76 +20,70 @@ 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 +common common-options + ghc-options: -Wall -threaded default-language: Haskell2010 build-depends: base >= 4.7 && < 5.0 - , unordered-containers >= 0.2.8 - , vector >= 0.12 + , rio + , lens + , formatting , fei-base , fei-nn - , fei-dataiter + , fei-modelzoo + , resourcet default-extensions: OverloadedLabels + , OverloadedStrings + , OverloadedLists , TypeFamilies + , DataKinds + , TypeApplications + , NoImplicitPrelude + , FlexibleInstances + , FlexibleContexts +Executable lenet + import: common-options + main-is: lenet.hs + hs-source-dirs: src Executable cifar10 + import: common-options 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 + build-depends: optparse-applicative 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 + import: common-options + main-is: custom-op.hs + hs-source-dirs: src + +Executable faster-rcnn + import: common-options + hs-source-dirs: src/RCNN + main-is: faster-rcnn.hs + other-modules: RCNN + build-depends: optparse-applicative , attoparsec - , text - , lens >= 4.12 , repa - , random-fu - , directory - , mtl , conduit - , fei-base - , fei-nn - , fei-dataiter + , resourcet + , store + , JuicyPixels + , random-source , fei-cocoapi - default-extensions: OverloadedLabels - , TypeFamilies + , fei-datasets +Executable mask-rcnn + import: common-options + hs-source-dirs: src/RCNN + main-is: mask-rcnn.hs + other-modules: RCNN + build-depends: optparse-applicative + , attoparsec + , repa + , conduit + , resourcet + , store + , random-source + , fei-cocoapi + , fei-datasets diff --git a/src/Model/FasterRCNN.hs b/src/Model/FasterRCNN.hs deleted file mode 100644 --- a/src/Model/FasterRCNN.hs +++ /dev/null @@ -1,606 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Model/Lenet.hs +++ /dev/null @@ -1,47 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Model/Resnet.hs +++ /dev/null @@ -1,284 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Model/Resnext.hs +++ /dev/null @@ -1,221 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Model/VGG.hs +++ /dev/null @@ -1,66 +0,0 @@ -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/RCNN/RCNN.hs b/src/RCNN/RCNN.hs new file mode 100644 --- /dev/null +++ b/src/RCNN/RCNN.hs @@ -0,0 +1,437 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedLists #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TemplateHaskell #-} +module RCNN where + +import Control.Applicative (ZipList (..)) +import Control.Lens (_1, _2, makePrisms) +import Control.Monad.Trans.Resource +import qualified Data.Array.Repa as Repa +import Formatting (sformat, string, (%)) +import Options.Applicative (Parser, ReadM, auto, + eitherReader, help, long, + metavar, option, + showDefault, strOption, + value) +import RIO +import RIO.Directory (canonicalizePath, + doesFileExist) +import qualified RIO.HashSet as S +import RIO.List (lastMaybe, unzip, unzip3) +import RIO.List.Partial (maximum) +import qualified RIO.NonEmpty as RNE +import qualified RIO.NonEmpty.Partial as RNE +import qualified RIO.Text as T +import qualified RIO.Vector.Boxed as V + +import MXNet.Base +import MXNet.Base.ParserUtils (decimal, endOfInput, list, + parseOnly, rational) +import MXNet.NN +import qualified MXNet.NN.DataIter.Anchor as Anchor +import qualified MXNet.NN.DataIter.Coco as Coco +import qualified MXNet.NN.Initializer as I +import MXNet.NN.ModelZoo.RCNN.FasterRCNN + +instance Coco.HasDatasetConfig (FeiApp t n Coco.CocoConfig) where + type DatasetTag (FeiApp t n Coco.CocoConfig) = "coco" + datasetConfig = fa_extra + +data CommonArgs = CommonArgs + { ds_base_path :: String + , ds_img_size :: Int + , ds_img_pixel_means :: [Float] + , ds_img_pixel_stds :: [Float] + } + deriving Show + +data ExtraArgs = TrainArgs + { pg_train_epochs :: Int + , pg_train_iter_per_epoch :: Int + } + | NoExtraArgs + deriving Show + +apRcnn :: (Parser RcnnConfiguration, Parser RcnnConfiguration) +apRcnn = (train, infr) + where + train = RcnnConfigurationTrain + <$> backbone + <*> batch_size + <*> feature_strides + <*> strOption (long "pretrained" <> metavar "PATH" + <> value "" + <> help "path to pretrained model") + <*> option floatx4 (long "bbox-reg-stds" <> metavar "BBOX_STDS" + <> value (0.1, 0.1, 0.2, 0.2)) + <*> rpn_anchor_scales + <*> rpn_anchor_ratios + <*> rpn_base_size + <*> rpn_pre_topk + <*> rpn_post_topk + <*> rpn_nms_threshold + <*> rpn_min_size + <*> option auto (long "rpn-batch-rois" <> metavar "BATCH-ROIS" + <> showDefault + <> value 256 + <> help "rpn number of rois per batch") + <*> 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") + <*> rcnn_num_classes + <*> rcnn_pool_sized + <*> rcnn_batch_rois + <*> 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 auto (long "rcnn-max-num-gt" <> metavar "NUM-GT" + <> showDefault + <> value 100 + <> help "rcnn max number of gt") + + infr = RcnnConfigurationInference + <$> backbone + <*> batch_size + <*> feature_strides + <*> strOption (long "checkpoint" <> metavar "PATH" + <> value "" + <> help "path to a saved model") + <*> option floatx4 (long "bbox-reg-stds" <> metavar "BBOX_STDS" <> value (0.1, 0.1, 0.2, 0.2)) + <*> rpn_anchor_scales + <*> rpn_anchor_ratios + <*> rpn_base_size + <*> rpn_pre_topk + <*> rpn_post_topk + <*> rpn_nms_threshold + <*> rpn_min_size + <*> rcnn_num_classes + <*> rcnn_pool_sized + <*> rcnn_batch_rois + <*> option auto (long "rcnn-force-nms" <> metavar "FORCE_NMS" <> value False) + <*> option auto (long "rcnn-nms-threshold" <> metavar "NMS_THRESH" <> value 0.5) + <*> option auto (long "rcnn-nms-topk" <> metavar "NMS_TOPK" <> value (-1)) + + + rpn_anchor_scales = option intList (long "rpn-anchor-scales" <> metavar "SCALES" + <> showDefault + <> value [8,16,32] + <> help "rpn anchor scales") + rpn_anchor_ratios = option floatList (long "rpn-anchor-ratios" <> metavar "RATIOS" + <> showDefault + <> value [0.5,1,2] + <> help "rpn anchor ratios") + rpn_base_size = option auto (long "rpn-anchor-bsize" <> metavar "BSIZE" + <> showDefault + <> value 16 + <> help "rpn anchor base size") + rpn_pre_topk = option auto (long "rpn-pre-nms-topk" <> metavar "PRE-NMS-TOPK" + <> showDefault + <> value 12000 + <> help "rpn nms pre-top-k") + rpn_post_topk = option auto (long "rpn-post-nms-topk" <> metavar "POST-NMS-TOPK" + <> showDefault + <> value 2000 + <> help "rpn nms post-top-k") + rpn_nms_threshold = option auto (long "rpn-nms-thresh" <> metavar "NMS-THRESH" + <> showDefault + <> value 0.7 + <> help "rpn nms threshold") + rpn_min_size = option auto (long "rpn-min-size" <> metavar "MIN-SIZE" + <> showDefault + <> value 16 + <> help "rpn min size") + rcnn_num_classes = option auto (long "rcnn-num-classes" <> metavar "NUM-CLASSES" + <> showDefault + <> value 81 + <> help "rcnn number of classes") + rcnn_batch_rois = option auto (long "rcnn-batch-rois" <> metavar "BATCH_ROIS" + <> showDefault + <> value 128 + <> help "rcnn batch rois") + rcnn_pool_sized = option auto (long "rcnn-pooled-size" <> metavar "POOLED-SIZE" + <> showDefault + <> value 14 + <> help "rcnn pooled size") + feature_strides = option intList (long "strides" <> metavar "STRIDE" + <> showDefault + <> value [16] + <> help "feature stride") + batch_size = option auto (long "batch-size" <> metavar "BATCH-SIZE" + <> showDefault + <> value 1 + <> help "batch size") + backbone = option auto (long "backbone" <> metavar "BACKBONE" + <> value VGG16 + <> help "vgg-16 or resnet-50") + +apCommon :: Parser CommonArgs +apCommon = CommonArgs + <$> strOption (long "base" <> metavar "PATH" + <> help "path to the dataset") + <*> option auto (long "img-size" <> metavar "SIZE" + <> showDefault + <> value 1024 + <> 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") + +apTrain :: Parser ExtraArgs +apTrain = TrainArgs + <$> option auto (long "train-epochs" <> metavar "EPOCHS" + <> value 500 + <> help "number of epochs to train") + <*> option auto (long "train-iter-per-epoch" <> metavar "ITER-PER-EPOCH" + <> value 100 + <> help "number of iter per epoch") + +floatList :: ReadM [Float] +floatList = eitherReader $ parseOnly (list rational<* endOfInput) . T.pack + +floatx4 :: ReadM (Float, Float, Float, Float) +floatx4 = let p = do fs <- list rational + case fs of + [a, b, c, d] -> return (a, b, c, d) + _ -> fail "should be exactly 4 floats" + in eitherReader $ parseOnly (p <* endOfInput) . T.pack + +intList :: ReadM [Int] +intList = eitherReader $ parseOnly (list decimal <* endOfInput) . T.pack + +toTriple [a, b, c] = (a, b, c) +toTriple x = error (show x) + + +default_initializer :: Initializer Float +default_initializer name = case name of + "features.rpn.rpn_conv_3x3.weight" -> I.normal 0.01 name + "features.rpn.rpn_conv_3x3.bias" -> I.zeros name + "features.rpn.rpn_cls_score.weight" -> I.normal 0.01 name + "features.rpn.rpn_cls_score.bias" -> I.zeros name + "features.rpn.rpn_bbox_pred.weight" -> I.normal 0.01 name + "features.rpn.rpn_bbox_pred.bias" -> I.zeros name + "features.rcnn.rcnn_cls_score.weight" -> I.normal 0.01 name + "features.rcnn.rcnn_cls_score.bias" -> I.zeros name + "features.rcnn.rcnn_bbox_feature.weight" -> I.normal 0.01 name + "features.rcnn.rcnn_bbox_feature.bias" -> I.zeros name + "features.rcnn.rcnn_bbox_pred.weight" -> I.normal 0.001 name + "features.rcnn.rcnn_bbox_pred.bias" -> I.zeros name + "features.fpn.0.conv1.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.0.conv2.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.1.conv1.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.1.conv2.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.2.conv1.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.2.conv2.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.3.conv1.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + "features.fpn.3.conv2.weight" -> I.xavier 1 I.XavierUniform I.XavierIn name + _ | T.isSuffixOf ".running_mean" name -> I.zeros name + | T.isSuffixOf ".running_var" name -> I.ones name + | T.isSuffixOf ".beta" name -> I.zeros name + | T.isSuffixOf ".gamma" name -> I.ones name + | otherwise -> I.zeros name + +loadWeights :: (DType a, MonadIO m, MonadReader env m, HasLogFunc env, HasCallStack) + => String -> Module t a m () +loadWeights weights_path = do + weights_path <- liftIO $ canonicalizePath weights_path + e <- liftIO $ doesFileExist (weights_path ++ ".params") + if not e + then lift . logInfo . display $ sformat ("'" % string % ".params' doesn't exist.") weights_path + else loadState weights_path ["features.9.gamma", + "features.9.beta", + "features.9.running_var", + "features.9.running_mean", + "output.weight", + "output.bias" + ] + +data Stage = TRAIN + | INFERENCE + +fixedParams :: Backbone -> Stage -> SymbolHandle -> IO (HashSet Text) +fixedParams backbone stage symbol = do + argnames <- listArguments symbol + return $ case (stage, backbone) of + (INFERENCE, _) + -> S.fromList argnames + (TRAIN, VGG16) + -> S.fromList [n | n <- argnames + -- fix conv_1_1, conv_1_2, conv_2_1, conv_2_2 + , elemM [0, 2, 5, 7] (layer n)] + (TRAIN, r) | r `elem` ([RESNET50, RESNET50FPN, RESNET101] :: [Backbone]) + -> S.fromList [n | n <- argnames + -- fix conv_0, stage_1_*, *_gamma, *_beta + , let layer_idx = layer n + , elemM [0, 1, 5] layer_idx || + (leqM 9 layer_idx && elemM ["gamma", "beta"] (lastName n))] + + where + toMaybe = either (const Nothing) Just + layer param = case T.split (=='.') param of + "features":n:_ -> toMaybe $ parseOnly decimal n + _ -> Nothing + lastName = lastMaybe . T.split (=='.') + elemM :: Eq a => [a] -> Maybe a -> Bool + elemM b = isJust . (>>= guard) . liftM (`elem` b) + leqM n = isJust . (>>= guard) . liftM (<= n) + +data App c = App LogFunc c +makePrisms ''App + +instance HasLogFunc (App c) where + logFuncL = _App . _1 + +instance Coco.HasDatasetConfig (App Coco.CocoConfig) where + type DatasetTag (App Coco.CocoConfig) = "coco" + datasetConfig = _App . _2 + +runApp :: c -> ReaderT (App c) (ResourceT IO) a -> IO a +runApp conf body = do + logopt <- logOptionsHandle stdout False + runResourceT $ withLogFunc logopt $ \logfunc -> + flip runReaderT (App logfunc conf) body + +generateTargets :: (SymbolHandle -> Layer (NonEmpty SymbolHandle)) + -> Coco.ImageInfo + -> NonEmpty Int + -> Anchor.Configuration + -> [Anchor.GTBox Repa.U] + -> IO (NDArray Float, NDArray Float, NDArray Float) +generateTargets feature_net im_info strides anchor_conf gt_boxes = do + feats <- runLayerBuilder $ variable "data" >>= feature_net + + -- there should equally number of features and strides, and pair them. + let feat_stride = RNE.zip feats strides + + layers <- mapM (uncurry make) feat_stride + let (cls_targets, box_targets, box_masks) = unzip3 $ RNE.toList layers + cls_targets <- mapM fromRepa cls_targets + box_targets <- mapM fromRepa box_targets + box_masks <- mapM fromRepa box_masks + cls_targets <- concat_ 0 cls_targets + box_targets <- concat_ 0 box_targets + box_masks <- concat_ 0 box_masks + return (cls_targets, box_targets, box_masks) + + where + [img_height, img_width, _] = Repa.toList im_info + -- we have padded the image to a square + img_size = floor (max img_height img_width) + base_size = anchor_conf ^. Anchor.conf_anchor_base_size + scales = anchor_conf ^. Anchor.conf_anchor_scales + ratios = anchor_conf ^. Anchor.conf_anchor_ratios + make :: SymbolHandle -> Int -> IO (Anchor.Labels, Anchor.Targets, Anchor.Weights) + make feat stride = do + (_, outputs, _, _) <- inferShape feat [("data", STensor [1,3,img_size,img_size])] + let [(_, STensor [_, _, h, w])] = outputs + anchors = Anchor.anchors (h, w) stride base_size scales ratios + runReaderT (Anchor.assign (V.fromList gt_boxes) img_size img_size anchors) anchor_conf + +padLength :: DType a => [NDArray a] -> a -> IO [NDArray a] +padLength arrays value = do + shps <- mapM ndshape arrays + let max_num = maximum $ map RNE.head shps + forM (zip arrays shps) $ \(a, n :| shp) -> + if n == max_num + then return a + else do + padding <- full value ((max_num - n) :| shp) + concat_ 0 [a, padding] + +withRpnTargets :: MonadIO m + => RcnnConfiguration + -> (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes) + -> m (String, [NDArray Float]) +withRpnTargets RcnnConfigurationTrain{..} dat = liftIO $ do + (cls_targets, box_targets, box_weights) <- + generateTargets extract info (RNE.fromList feature_strides) conf (V.toList gt) + imgA <- fromRepa img + infoA <- fromRepa info + gtA <- stack 0 . V.toList =<< mapM fromRepa gt + return (filename, [gtA, imgA, infoA, cls_targets, box_targets, box_weights]) + where + (filename, img, info, gt) = dat + conf = Anchor.Configuration + { Anchor._conf_anchor_scales = rpn_anchor_scales + , Anchor._conf_anchor_ratios = rpn_anchor_ratios + , Anchor._conf_anchor_base_size = rpn_anchor_base_size + , Anchor._conf_allowed_border = rpn_allowd_border + , Anchor._conf_fg_num = floor $ (rpn_fg_fraction * fromIntegral rpn_batch_rois) + , Anchor._conf_batch_num = rpn_batch_rois + , Anchor._conf_bg_overlap = rpn_bg_overlap + , Anchor._conf_fg_overlap = rpn_fg_overlap + } + extract = sequential "features" . features1 backbone + + +withRpnTargets'Mask :: MonadIO m + => RcnnConfiguration + -> (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes, Coco.Masks) + -> m (String, [NDArray Float]) +withRpnTargets'Mask conf dat = do + let (filename, img, info, gt, msks) = dat + (_, ret) <- withRpnTargets conf (filename, img, info, gt) + liftIO $ do + msksA <- stack 0 . V.toList =<< mapM fromRepa msks + msksA <- divScalar 255 =<< cast #float32 msksA :: IO (NDArray Float) + return (filename, msksA : ret) + +toListNDArray :: MonadIO m + => (String, Coco.ImageTensor, Coco.ImageInfo, Coco.GTBoxes) + -> m (String, [NDArray Float]) +toListNDArray (filename, img, info, gt) = liftIO $ do + imgA <- fromRepa img + infoA <- fromRepa info + gtA <- stack 0 . V.toList =<< mapM fromRepa gt + return (filename, [gtA, imgA, infoA]) + +concatBatch :: MonadIO m => [(String, [NDArray Float])] -> m ([String], [NDArray Float]) +concatBatch batch = liftIO $ do + let (filenames, tensors) = unzip batch + gt : others = unzipList tensors + -- gt in the batch may not have the same number + -- must be padded with -1 before stacking + gt <- stack 0 =<< padLength gt (-1) + -- other tensors can be simply stacked + others <- mapM (stack 0) others + return (filenames, gt : others) + + +concatBatch'Mask :: MonadIO m => [(String, [NDArray Float])] -> m ([String], [NDArray Float]) +concatBatch'Mask batch = liftIO $ do + let (filenames, tensors) = unzip batch + mask_gt : box_gt : others = unzipList tensors + mask_gt <- stack 0 =<< padLength mask_gt 0 + box_gt <- stack 0 =<< padLength box_gt (-1) + others <- mapM (stack 0) others + return (filenames, mask_gt : box_gt : others) + +unzipList :: [[a]] -> [[a]] +unzipList = getZipList . traverse ZipList + + diff --git a/src/RCNN/faster-rcnn.hs b/src/RCNN/faster-rcnn.hs new file mode 100644 --- /dev/null +++ b/src/RCNN/faster-rcnn.hs @@ -0,0 +1,293 @@ +{-# LANGUAGE OverloadedLists #-} +{-# LANGUAGE RecordWildCards #-} +module Main where + +import Codec.Picture (PixelRGBA8 (..), writePng) +import Control.Lens (ix, use, (.=), (^?!)) +import Data.Array.Repa (Array, DIM1, DIM2, DIM3, + DIM4, U) +import qualified Data.Array.Repa as Repa +import Data.Conduit ((.|)) +import qualified Data.Conduit.List as C +import Data.Random.Source.StdGen (mkStdGen) +import Formatting (fixed, formatToString, int, + left, sformat, stext, + string, (%)) +import Options.Applicative (command, execParser, + fullDesc, header, helper, + hsubparser, info, progDesc, + (<**>)) +import RIO hiding (Const) +import RIO.Char (isDigit) +import RIO.FilePath +import qualified RIO.HashMap as M +import qualified RIO.HashSet as S +import RIO.List (sort) +import qualified RIO.Text as T +import qualified RIO.Vector.Boxed as VB +import qualified RIO.Vector.Storable as VS + +import MXNet.Base +import qualified MXNet.Base.Operators.Tensor as Ops +import qualified MXNet.Base.ParserUtils as P +import MXNet.NN +import qualified MXNet.NN.DataIter.Anchor as Anchor +import qualified MXNet.NN.DataIter.Coco as Coco +import MXNet.NN.DataIter.Conduit +import MXNet.NN.ModelZoo.RCNN.FasterRCNN +import MXNet.NN.Utils.Render + +import RCNN + +main :: IO () +main = do + mxRandomSeed 8 + registerCustomOperator ("anchor_generator", Anchor.buildAnchorGenerator) + + let (apRcnnT, apRcnnI) = apRcnn + apT = liftA3 (,,) apRcnnT apCommon apTrain + apI = liftA3 (,,) apRcnnI apCommon (pure NoExtraArgs) + whole = hsubparser + ( command "train" (info apT (progDesc "Train")) + <> command "inference" (info apI (progDesc "Run inference")) + ) + args <- liftIO $ execParser $ info (whole <**> helper) (fullDesc <> header "Faster-RCNN") + case args of + (RcnnConfigurationTrain{}, _, _) -> mainTrain args + (RcnnConfigurationInference{}, _, _) -> mainInfer args + + +-- data Dbg e a = Dbg (e a) +-- +-- instance EvalMetricMethod e => EvalMetricMethod (Dbg e) where +-- data MetricData (Dbg e) a = DbgPriv (MetricData e a) +-- newMetric phase (Dbg conf) = do +-- p <- newMetric phase conf +-- return $ DbgPriv p +-- evalMetric (DbgPriv p) bindings outputs = do +-- liftIO $ do +-- a <- toCPU $ bindings ^?! ix "rpn_cls_targets" +-- a <- toVector =<< prim Ops._norm (#ord := 1 .& #data := a .& Nil) +-- b <- toCPU $ outputs ^?! ix 0 +-- b <- toVector =<< prim Ops._norm (#ord := 1 .& #data := b .& Nil) +-- traceShowM (a, b) +-- evalMetric p bindings outputs +-- formatMetric (DbgPriv p) = formatMetric p +-- +-- data AccDbg a = AccDbg (Accuracy a) +-- +-- instance EvalMetricMethod AccDbg where +-- data MetricData AccDbg a = AccDbgPriv (MetricData Accuracy a) +-- newMetric phase (AccDbg conf) = do +-- priv <- newMetric phase conf +-- return $ AccDbgPriv priv +-- evalMetric (AccDbgPriv accpriv) bindings outputs = do +-- liftIO $ do +-- let AccuracyPriv acc _ _ _ = accpriv +-- lbl <- toCPU $ _mtr_acc_get_gt acc bindings outputs +-- -- x <- toCPU $ outputs ^?! ix 6 +-- traceShowM =<< toVector lbl +-- -- traceShowM . VS.take 20 =<< toVector x +-- ret <- evalMetric accpriv bindings outputs +-- return ret +-- formatMetric (AccDbgPriv acc) = formatMetric acc +-- +-- +-- data Dbg a = Dbg (M.HashMap Text (NDArray a) -> [NDArray a] -> NDArray a) +-- +-- instance EvalMetricMethod Dbg where +-- data MetricData Dbg a = DbgPriv (Dbg a) +-- newMetric phase a = return (DbgPriv a) +-- evalMetric (DbgPriv (Dbg __get)) bindings outputs = liftIO $ do +-- array <- toCPU $ __get bindings outputs +-- shp <- ndshape array +-- val <- toVector array +-- traceShowM (shp, val) +-- return M.empty +-- formatMetric _ = return "" + + +mainTrain (rcnn_conf@RcnnConfigurationTrain{..}, CommonArgs{..}, TrainArgs{..}) = do + rand_gen <- liftIO $ newIORef $ mkStdGen 19 + coco_inst <- Coco.coco ds_base_path "train2017" + let coco_conf = Coco.CocoConfig coco_inst ds_img_size + (toTriple ds_img_pixel_means) + (toTriple ds_img_pixel_stds) + -- There is a serious problem with asyncConduit. It made the training loop running + -- in different threads, which is very bad because the execution of ExecutorForward + -- has a thread-local state (saving the temporary workspace for cudnn) + -- + -- data_iter = asyncConduit (Just batch_size) $ + -- + data_iter = ConduitData (Just batch_size) $ + Coco.cocoImagesBBoxes rand_gen .| + C.mapM (Coco.augmentWithBBoxes rand_gen) .| + C.mapM (withRpnTargets rcnn_conf) .| + C.chunksOf batch_size .| + C.mapM concatBatch + + runFeiM coco_conf $ do + (_, sym) <- runLayerBuilder $ graphT rcnn_conf + fixed_params <- liftIO $ fixedParams backbone TRAIN sym + + initSession @"faster_rcnn" sym (Config { + _cfg_data = M.fromList [("data", (STensor [batch_size, 3, ds_img_size, ds_img_size])) + ,("im_info", (STensor [batch_size, 3])) + ,("gt_boxes", (STensor [batch_size, 1, 5])) + ], + _cfg_label = ["rpn_cls_targets" + ,"rpn_box_targets" + ,"rpn_box_masks" + ], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = fixed_params, + _cfg_context = contextGPU0 }) + + let lr_sched = lrOfFactor (#base := 0.01 .& #factor := 0.5 .& #step := 5000 .& Nil) + optm <- makeOptimizer SGD'Mom lr_sched (#momentum := 0.9 + .& #wd := 0.0001 + .& #rescale_grad := 1 / (fromIntegral batch_size) + .& #clip_gradient := 10 + .& Nil) + -- optm <- makeOptimizer ADAMW lr_sched (#rescale_grad := 1 / (fromIntegral batch_size) + -- .& #eta := 0.001 + -- .& #wd := 0.0001 + -- .& #clip_gradient := 10 .& Nil) + + checkpoint <- lastSavedState "checkpoints" "faster_rcnn" + start_epoch <- case checkpoint of + Nothing -> do + logInfo . display $ sformat string pretrained_weights + unless (null pretrained_weights) + (askSession $ loadWeights pretrained_weights) + return (1 :: Int) + Just filename -> do + askSession $ loadState filename [] + let (base, _) = splitExtension filename + fn_rev = T.reverse $ T.pack base + epoch = P.parseR (P.takeWhile isDigit <* P.takeText) fn_rev + epoch_next = (P.parseR P.decimal $ T.reverse epoch) + 1 + return epoch_next + logInfo . display $ sformat ("fixed parameters: " % stext) (tshow (sort $ S.toList fixed_params)) + + metric <- newMetric "train" (Accuracy (Just "RPN-acc") (PredByThreshold 0.5) 0 + (\_ preds -> preds ^?! ix 0) + (\bindings _ -> bindings ^?! ix "rpn_cls_targets") + :* Accuracy (Just "RCNN-acc") PredByArgmax 1 + (\_ preds -> preds ^?! ix 3) + (\_ preds -> preds ^?! ix 5) + :* CrossEntropy (Just "RPN-ce") False + (\_ preds -> preds ^?! ix 0) + (\bindings _ -> bindings ^?! ix "rpn_cls_targets") + :* CrossEntropy (Just "RCNN-ce") True + (\_ preds -> preds ^?! ix 3) + (\_ preds -> preds ^?! ix 5) + :* Norm (Just "RPN-L1") 1 + (\_ preds -> preds ^?! ix 2) + :* Norm (Just "RCNN-L1") 1 + (\_ preds -> preds ^?! ix 4) + :* MNil) + + -- update the internal counting of the iterations + -- the lr is updated as per to it + askSession $ do + untag . mod_statistics . stat_num_upd .= (start_epoch - 1) * pg_train_iter_per_epoch + + forM_ ([start_epoch..pg_train_epochs] :: [Int]) $ \ ei -> do + logInfo . display $ sformat ("Epoch " % int) ei + let slice = takeD pg_train_iter_per_epoch data_iter + void $ forEachD_i slice $ \(i, (fn, [x0, x1, x2, y0, y1, y2])) -> askSession $ do + let binding = M.fromList [ ("gt_boxes", x0) + , ("data", x1) + , ("im_info", x2) + , ("rpn_cls_targets", y0) + , ("rpn_box_targets", y1) + , ("rpn_box_masks", y2) + ] + fitAndEval optm binding metric + eval <- metricFormat metric + lr <- use (untag . mod_statistics . stat_last_lr) + + logInfo . display $ sformat (int % " " % stext % " LR: " % fixed 5) i eval lr + + askSession $ saveState (ei == 1) + (formatToString ("checkpoints/faster_rcnn_epoch_" % left 3 '0') ei) + + +mainInfer (rcnn_conf@RcnnConfigurationInference{..}, CommonArgs{..}, NoExtraArgs) = do + coco_inst@(Coco.Coco _ _ coco_inst_ _) <- Coco.coco ds_base_path "val2017" + rand_gen <- newIORef $ mkStdGen 24 + + let coco_conf = Coco.CocoConfig coco_inst ds_img_size + (toTriple ds_img_pixel_means) + (toTriple ds_img_pixel_stds) + data_iter = ConduitData (Just batch_size) ( + Coco.cocoImagesBBoxes rand_gen .| + C.mapM toListNDArray .| + C.chunksOf batch_size .| + C.mapM concatBatch + ) & takeD 5 + + runFeiM coco_conf $ do + (_, sym) <- runLayerBuilder $ graphI rcnn_conf + fixed_params <- liftIO $ fixedParams backbone INFERENCE sym + fixed_params <- return $ S.difference fixed_params (S.fromList ["data", "im_info"]) + + initSession @"faster_rcnn" sym (Config { + _cfg_data = M.fromList [("data", (STensor [batch_size, 3, ds_img_size, ds_img_size])), + ("im_info", (STensor [batch_size, 3]))], + _cfg_label = [], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = fixed_params, + _cfg_context = contextGPU0 + }) + + askSession $ do + loadState checkpoint [] + void $ forEachD_i data_iter $ \(i, (fn, [x0, x1, x2])) -> do + let bindings = M.fromList [ ("data", x1) + , ("im_info", x2) + ] + [cls_ids, scores, boxes] <- forwardOnly bindings + + -- cls_ids: (B, num_fg_classes * rcnn_nms_topk, 1) + -- scores : (B, num_fg_classes * rcnn_nms_topk, 1) + -- boxes : (B, num_fg_classes * rcnn_nms_topk, 4) + + liftIO $ do + fn <- pure $ VB.fromList fn + infos <- vunstack <$> toRepa @DIM2 x2 + -- gt_boxes<- vunstack <$> toRepa @DIM3 x0 + cls_ids <- vunstack <$> toRepa @DIM3 cls_ids + scores <- vunstack <$> toRepa @DIM3 scores + boxes <- vunstack <$> toRepa @DIM3 boxes + mean <- fromVector [3] (VS.fromList ds_img_pixel_means) + std <- fromVector [3] (VS.fromList ds_img_pixel_stds) + images <- transpose x1 [0, 2, 3, 1] >>= + mulBroadcast std >>= + addBroadcast mean >>= + mulScalar 255 + images <- vunstack <$> toRepa @DIM4 images + forM_ (VB.zip6 fn images infos cls_ids scores boxes) renderImageBBoxes + + +renderImageBBoxes :: (String, Array U DIM3 Float, Array U DIM1 Float, Array U DIM2 Float, Array U DIM2 Float, Array U DIM2 Float) -> IO () +renderImageBBoxes (filename, image, info, cls_ids, scores, boxes) = do + let [height, width, scale] = Repa.toUnboxed info + jp_image = imageFromRepa image + -- TODO scale the image and bboxes back to orignal size + width <- pure $ floor width + height <- pure $ floor height + writePng filename $ render width height $ do + drawImage jp_image + let all_boxes = vunstack boxes + all_scores = VB.convert $ Repa.toUnboxed scores + all_cls = VB.convert $ Repa.toUnboxed cls_ids + forM_ (VB.zip3 all_cls all_scores all_boxes) $ \(cls, score, box) -> do + let [x, y, x', y'] = Repa.toUnboxed box + when (score >= 0.5) $ do + traceShowM (score, box) + drawBox (PixelRGBA8 255 0 0 255) 1.0 x y x' y' Nothing + diff --git a/src/RCNN/mask-rcnn.hs b/src/RCNN/mask-rcnn.hs new file mode 100644 --- /dev/null +++ b/src/RCNN/mask-rcnn.hs @@ -0,0 +1,156 @@ +{-# LANGUAGE OverloadedLists #-} +{-# LANGUAGE RecordWildCards #-} +module Main where + +import Control.Lens (ix, use, (.=), (^?!)) +import Data.Conduit ((.|)) +import qualified Data.Conduit.List as C +import Data.Random.Source.StdGen (mkStdGen) +import Formatting (fixed, formatToString, int, + left, sformat, stext, + string, (%)) +import Options.Applicative (command, execParser, + fullDesc, header, helper, + hsubparser, info, progDesc, + (<**>)) +import RIO hiding (Const) +import RIO.Char (isDigit) +import RIO.FilePath +import qualified RIO.HashMap as M +import qualified RIO.HashSet as S +import RIO.List (sort) +import qualified RIO.Text as T + +import MXNet.Base +import qualified MXNet.Base.ParserUtils as P +import MXNet.NN +import qualified MXNet.NN.DataIter.Anchor as Anchor +import qualified MXNet.NN.DataIter.Coco as Coco +import MXNet.NN.DataIter.Conduit +import MXNet.NN.ModelZoo.RCNN.FasterRCNN (RcnnConfiguration (..)) +import MXNet.NN.ModelZoo.RCNN.MaskRCNN + +import RCNN + + +main :: IO () +main = do + mxRandomSeed 8 + registerCustomOperator ("anchor_generator", Anchor.buildAnchorGenerator) + + let (apRcnnT, apRcnnI) = apRcnn + apT = liftA3 (,,) apRcnnT apCommon apTrain + apI = liftA3 (,,) apRcnnI apCommon (pure NoExtraArgs) + whole = hsubparser + ( command "train" (info apT (progDesc "Train")) + <> command "inference" (info apI (progDesc "Run inference")) + ) + args <- liftIO $ execParser $ info (whole <**> helper) (fullDesc <> header "Faster-RCNN") + case args of + (RcnnConfigurationTrain{}, _, _) -> mainTrain args + +mainTrain (rcnn_conf@RcnnConfigurationTrain{..}, CommonArgs{..}, TrainArgs{..}) = do + rand_gen <- liftIO $ newIORef $ mkStdGen 42 + coco_inst <- Coco.coco ds_base_path "train2017" + + let coco_conf = Coco.CocoConfig coco_inst ds_img_size (toTriple ds_img_pixel_means) (toTriple ds_img_pixel_stds) + -- There is a serious problem with asyncConduit. It made the training loop running + -- in different threads, which is very bad because the execution of ExecutorForward + -- has a thread-local state (saving the temporary workspace for cudnn) + -- + -- data_iter = asyncConduit (Just batch_size) $ + -- + data_iter = ConduitData (Just batch_size) $ + Coco.cocoImagesBBoxesMasks rand_gen .| + C.mapM (withRpnTargets'Mask rcnn_conf) .| + C.chunksOf batch_size .| + C.mapM concatBatch'Mask + + runFeiM coco_conf $ do + (_, sym) <- runLayerBuilder $ graphT rcnn_conf + fixed_params <- liftIO $ fixedParams backbone TRAIN sym + + initSession @"mask_rcnn" sym (Config { + _cfg_data = M.fromList [ ("data", STensor [batch_size, 3, ds_img_size, ds_img_size]) + , ("im_info", STensor [batch_size, 3]) + , ("gt_boxes", STensor [batch_size, 1, 5]) + , ("gt_masks", STensor [batch_size, 1, ds_img_size, ds_img_size]) + ], + _cfg_label = ["rpn_cls_targets" + ,"rpn_box_targets" + ,"rpn_box_masks" + ], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = fixed_params, + _cfg_context = contextGPU0 }) + + let lr_sched = lrOfFactor (#base := 0.004 .& #factor := 0.5 .& #step := 2000 .& Nil) + -- optm <- makeOptimizer SGD'Mom lr_sched (#momentum := 0.9 + -- .& #wd := 0.0001 + -- .& #rescale_grad := 1 / (fromIntegral batch_size) + -- .& #clip_gradient := 10 + -- .& Nil) + optm <- makeOptimizer ADAM lr_sched (#rescale_grad := 1 / (fromIntegral batch_size) + .& #wd := 0.0001 + .& #clip_gradient := 10 .& Nil) + + checkpoint <- lastSavedState "checkpoints" "mask_rcnn" + start_epoch <- case checkpoint of + Nothing -> do + logInfo . display $ sformat string pretrained_weights + unless (null pretrained_weights) + (askSession $ loadWeights pretrained_weights) + return (1 :: Int) + Just filename -> do + askSession $ loadState filename [] + let (base, _) = splitExtension filename + fn_rev = T.reverse $ T.pack base + epoch = P.parseR (P.takeWhile isDigit <* P.takeText) fn_rev + epoch_next = P.parseR P.decimal $ T.reverse epoch + return epoch_next + logInfo . display $ sformat ("fixed parameters: " % stext) (tshow (sort $ S.toList fixed_params)) + + metric <- newMetric "train" (Accuracy (Just "RPN-acc") (PredByThreshold 0.5) 0 + (\_ preds -> preds ^?! ix 0) + (\bindings _ -> bindings ^?! ix "rpn_cls_targets") + :* Accuracy (Just "RCNN-acc") PredByArgmax 1 + (\_ preds -> preds ^?! ix 3) + (\_ preds -> preds ^?! ix 5) + :* CrossEntropy (Just "RPN-ce") False + (\_ preds -> preds ^?! ix 0) + (\bindings _ -> bindings ^?! ix "rpn_cls_targets") + :* CrossEntropy (Just "RCNN-ce") True + (\_ preds -> preds ^?! ix 3) + (\_ preds -> preds ^?! ix 5) + :* Norm (Just "RPN-L1") 1 + (\_ preds -> preds ^?! ix 2) + :* Norm (Just "RCNN-L1") 1 + (\_ preds -> preds ^?! ix 4) + :* MNil) + + -- update the internal counting of the iterations + -- the lr is updated as per to it + askSession $ do + untag . mod_statistics . stat_num_upd .= (start_epoch - 1) * pg_train_iter_per_epoch + + forM_ ([start_epoch..pg_train_epochs] :: [Int]) $ \ ei -> do + logInfo . display $ sformat ("Epoch " % int) ei + let slice = takeD pg_train_iter_per_epoch data_iter + void $ forEachD_i slice $ \(i, (fn, [x0, x1, x2, x3, y0, y1, y2])) -> askSession $ do + let binding = M.fromList [ ("gt_masks", x0) + , ("gt_boxes", x1) + , ("data", x2) + , ("im_info", x3) + , ("rpn_cls_targets", y0) + , ("rpn_box_targets", y1) + , ("rpn_box_masks", y2) + ] + fitAndEval optm binding metric + eval <- metricFormat metric + lr <- use (untag . mod_statistics . stat_last_lr) + logInfo . display $ sformat (int % " " % stext % " LR: " % fixed 5) i eval lr + + askSession $ saveState (ei == 1) + (formatToString ("checkpoints/mask_rcnn_epoch_" % left 3 '0') ei) + diff --git a/src/cifar10.hs b/src/cifar10.hs --- a/src/cifar10.hs +++ b/src/cifar10.hs @@ -1,80 +1,122 @@ -{-# 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 Control.Lens (ix, use, (^?!)) +import Formatting (float, int, sformat, stext, (%)) +import Options.Applicative +import RIO +import qualified RIO.HashMap as M +import qualified RIO.HashSet as S +import RIO.List.Partial (last) +import qualified RIO.Text as T -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 +import MXNet.Base (ArgOf (..), FShape (..), + HMap (..), contextCPU, + contextGPU0, listArguments, (.&)) +import MXNet.NN +import MXNet.NN.DataIter.Streaming +import qualified MXNet.NN.Initializer as I +import qualified MXNet.NN.ModelZoo.Resnet as Resnet +import qualified MXNet.NN.ModelZoo.Resnext as Resnext -type ArrayF = NDArray Float -type DS = StreamData (TrainM Float IO) (ArrayF, ArrayF) +batch_size = 128 -data Model = Resnet | Resnext deriving (Show, Read) -data ProgArg = ProgArg Model +data Model = Resnet + | Resnext + deriving (Show, Read) +data ProgArg = ProgArg Model (Maybe String) cmdArgParser :: Parser ProgArg -cmdArgParser = ProgArg <$> (option auto $ short 'm' <> metavar "MODEL" <> showDefault <> value Resnet) - -range :: Int -> [Int] -range = enumFromTo 1 +cmdArgParser = ProgArg + <$> (option auto $ short 'm' <> metavar "MODEL" <> showDefault <> value Resnet) + <*> (option maybe $ short 'p' <> metavar "PRETRAINED" <> showDefault <> value Nothing) + where + maybe = maybeReader (Just . Just) 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 + | T.isSuffixOf ".bias" name = I.zeros name shp + | T.isSuffixOf ".beta" name = I.zeros name shp + | T.isSuffixOf ".gamma" name = I.ones name shp + | T.isSuffixOf ".running_mean" name = I.zeros name shp + | T.isSuffixOf ".running_var" name = I.ones name shp + | otherwise = case shp of + [_,_] -> I.xavier 2.0 I.XavierGaussian I.XavierIn name shp + _ -> I.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 - } +main = runFeiM () $ do + ProgArg model pretrained <- liftIO $ execParser $ info + (cmdArgParser <**> helper) (fullDesc <> header "CIFAR-10 solver") - 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 + net <- runLayerBuilder $ do + dat <- variable "x" + lbl <- variable "y" + logits <- case model of + Resnet -> Resnet.resnet50 10 dat + Resnext -> Resnext.symbol dat + named "softmax" $ softmaxoutput (#data := logits .& #label := lbl .& Nil) - train sess $ do + fixed <- case pretrained of + Nothing -> return S.empty + Just _ -> fixedParams net model - 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 + initSession @"cifar10" net (Config { + _cfg_data = M.singleton "x" (STensor [batch_size, 3,32,32]), + _cfg_label = ["y"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = fixed, + _cfg_context = contextGPU0 }) + let lr_scheduler = lrOfMultifactor $ #steps := [100, 200, 300] + .& #base := 0.0001 + .& #factor:= 0.75 .& Nil + ce = CrossEntropy Nothing True + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") + acc = Accuracy Nothing PredByArgmax 0 + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") + + optimizer <- makeOptimizer SGD'Mom lr_scheduler Nil + + let trainingData = imageRecordIter (#path_imgrec := "data/cifar10_train.rec" + .& #data_shape := [3,32,32] + .& #batch_size := batch_size .& Nil) + valData = imageRecordIter (#path_imgrec := "data/cifar10_val.rec" + .& #data_shape := [3,32,32] + .& #batch_size := 16 .& Nil) + askSession $ case pretrained of + Just path -> loadState path ["output.weight", "output.bias"] + Nothing -> return () + + forM_ ([1..20] :: [Int]) $ \ ei -> do + logInfo . display $ sformat ("Epoch " % int) ei + metric <- newMetric "train" (ce :* acc :* MNil) + void $ forEachD_i trainingData $ \(i, (x, y)) -> askSession $ do + let binding = M.fromList [("x", x), ("y", y)] + fitAndEval optimizer binding metric + eval <- metricFormat metric + lr <- use (untag . mod_statistics . stat_last_lr) + when (i `mod` 20 == 0) $ do + logInfo . display $ sformat (int % " " % stext % " LR: " % float) i eval lr + + metric <- newMetric "val" (acc :* MNil) + void $ forEachD_i valData $ \(_, (x, y)) -> askSession $ do + pred <- forwardOnly (M.singleton "x" x) + void $ metricUpdate metric (M.singleton "y" y) pred + eval <- metricFormat metric + logInfo . display $ sformat ("Validation: " % stext) eval + +fixedParams symbol _ = do + argnames <- listArguments symbol + return $ S.fromList [n | n <- argnames + -- fix conv_0, stage_1_*, *_gamma, *_beta + , layer n `elemL` ["1", "5"] || name n `elemL` ["gamma", "beta"]] + where - bind ["x", "y"] (dat, lbl) = M.fromList [("x", dat), ("y", lbl)] + layer param = case T.split (=='.') param of + "features":n:_ -> n + _ -> "" + name param = last $ T.split (=='.') param + elemL :: Eq a => a -> [a] -> Bool + elemL = elem diff --git a/src/custom-op.hs b/src/custom-op.hs --- a/src/custom-op.hs +++ b/src/custom-op.hs @@ -1,20 +1,26 @@ 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) +import Control.Lens (ix, (^?!)) +import Formatting +import RIO hiding (Const) +import qualified RIO.HashMap as M +import qualified RIO.HashSet as S +import RIO.List (unzip) +import qualified RIO.NonEmpty as RNE +import qualified RIO.Text as T +import qualified RIO.Vector.Boxed as V +import qualified RIO.Vector.Storable as SV -type ArrayF = NDArray Float +import MXNet.Base +import MXNet.Base.Operators.Tensor +import MXNet.NN +import MXNet.NN.DataIter.Class +import MXNet.NN.DataIter.Streaming +import qualified MXNet.NN.Initializer as I +import MXNet.NN.Layer +batch_size = 128 + data SoftmaxProp = SoftmaxProp instance CustomOperationProp SoftmaxProp where @@ -22,131 +28,98 @@ 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], []) + -- data: [batch_size, N] + -- label: [batch_size] + -- output: [batch_size, N] + -- loss: [batch_size] + let STensor (batch_size :| _) = data_shape + out_shape = STensor (batch_size :| []) + in ([data_shape, out_shape], [data_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) + forward _ [ReqWrite] [in_data, label] [out] aux is_train = do + label <- prim _one_hot (#indices := label .& #depth := 10 .& Nil) + r <- prim _softmax (#data := in_data .& Nil) + void $ copy r out - 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) + backward _ [ReqWrite] [_, label] [out] [in_grad, _] _ aux = do + label <- prim _one_hot (#indices := label .& #depth := 10 .& Nil) + result <- prim _elemwise_sub (#lhs := out .& #rhs := label .& Nil) + void $ copy result in_grad -symbol :: DType a => IO (Symbol a) +symbol :: Layer SymbolHandle symbol = do - x <- NN.variable "x" - y <- NN.variable "y" + x <- variable "x" + y <- 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) + sequential "custom-op" $ do + v1 <- convolution (#data := x .& #kernel := [5,5] .& #num_filter := 20 .& Nil) + a1 <- activation (#data := v1 .& #act_type := #tanh .& Nil) + p1 <- pooling (#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) + v2 <- convolution (#data := p1 .& #kernel := [5,5] .& #num_filter := 50 .& Nil) + a2 <- activation (#data := v2 .& #act_type := #tanh .& Nil) + p2 <- pooling (#data := a2 .& #kernel := [2,2] .& #pool_type := #max .& Nil) - fl <- NN.flatten "flatten" (#data := p2 .& Nil) + fl <- flatten p2 - v3 <- NN.fullyConnected "fc1" (#data := fl .& #num_hidden := 500 .& Nil) - a3 <- NN.activation "fc1-a" (#data := v3 .& #act_type := #tanh .& Nil) + v3 <- fullyConnected (#data := fl .& #num_hidden := 500 .& Nil) + a3 <- activation (#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 + v4 <- fullyConnected (#data := a3 .& #num_hidden := 10 .& Nil) + named "softmax" $ prim _Custom (#data := [v4, y] .& #op_type := "softmax_custom" .& Nil) -default_initializer :: NN.Initializer Float +default_initializer :: Initializer Float default_initializer name shp - | NN.endsWith "-bias" name = NN.zeros name shp - | otherwise = NN.normal 0.1 name shp + | T.isSuffixOf "-bias" name = I.zeros name shp + | otherwise = I.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 "" +main = runFeiM () $ do + liftIO $ registerCustomOperator ("softmax_custom", \_ -> return SoftmaxProp) + net <- runLayerBuilder symbol + initSession @"lenet" net (Config { + _cfg_data = M.singleton "x" (STensor [batch_size, 1,28,28]), + _cfg_label = ["y"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = S.fromList [], + _cfg_context = contextGPU0 }) + optimizer <- makeOptimizer SGD'Mom (Const 0.0002) Nil - liftIO $ putStrLn $ "[Test] " + let ce = CrossEntropy Nothing True + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") + acc = Accuracy Nothing PredByArgmax 0 + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") - 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" + trainingData = mnistIter (#image := "data/train-images-idx3-ubyte" + .& #label := "data/train-labels-idx1-ubyte" + .& #batch_size := batch_size .& Nil) + testingData = mnistIter (#image := "data/t10k-images-idx3-ubyte" + .& #label := "data/t10k-labels-idx1-ubyte" + .& #batch_size := 16 .& Nil) - 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 + total <- sizeD trainingData + logInfo . display $ sformat "[Train] " + forM_ (V.enumFromTo 1 10) $ \ind -> do + logInfo . display $ sformat ("iteration " % int) ind + metric <- newMetric "train" (ce :* acc :* MNil) + void $ forEachD_i trainingData $ \(i, (x, y)) -> askSession $ do + fitAndEval optimizer (M.fromList [("x", x), ("y", y)]) metric + eval <- metricFormat metric + when (i `mod` 100 == 1) $ + logInfo . display $ sformat (int % "/" % int % ":" % stext) i total eval - where - argmax :: ArrayF -> IO ArrayF - argmax (NDArray ys) = NDArray . head <$> A.argmax (#data := ys .& #axis := Just 1 .& Nil) + metric <- newMetric "val" (acc :* MNil) + forEachD_i testingData $ \(i, (x, y)) -> askSession $ do + pred <- forwardOnly (M.singleton "x" x) + void $ metricUpdate metric (M.singleton "y" y) pred + eval <- metricFormat metric + logInfo . display $ sformat ("Validation: " % stext) eval diff --git a/src/lenet.hs b/src/lenet.hs --- a/src/lenet.hs +++ b/src/lenet.hs @@ -1,95 +1,80 @@ -{-# 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 Control.Lens (ix, (^?!)) +import Formatting (int, sformat, stext, (%)) +import RIO hiding (Const) +import qualified RIO.HashMap as M +import qualified RIO.HashSet as S +import qualified RIO.Vector.Boxed as V -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 +import MXNet.Base +import MXNet.NN +import MXNet.NN.DataIter.Conduit +import qualified MXNet.NN.Initializer as I +import qualified MXNet.NN.ModelZoo.Lenet as Model -type ArrayF = NDArray Float -type DS = ConduitData (TrainM Float IO) (ArrayF, ArrayF) +batch_size = 128 -range :: Int -> [Int] -range = enumFromTo 1 +range :: Int -> Vector Int +range = V.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 - +default_initializer name shp = + case length shp of + 1 -> I.zeros name shp + 2 -> I.xavier 2.0 I.XavierGaussian I.XavierIn name shp + _ -> I.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 +main = runFeiM'nept "jiasen/lenet" () $ do + net <- runLayerBuilder Model.symbol + initSession @"lenet" net (Config { + _cfg_data = M.singleton "x" (STensor [batch_size, 1, 28, 28]), + _cfg_label = ["y"], + _cfg_initializers = M.empty, + _cfg_default_initializer = default_initializer, + _cfg_fixed_params = S.fromList [], + _cfg_context = contextGPU0 }) + optm <- 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 := batch_size .& Nil) + let testingData = mnistIter (#image := "data/t10k-images-idx3-ubyte" + .& #label := "data/t10k-labels-idx1-ubyte" + .& #batch_size := 16 .& Nil) - 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) + total <- sizeD trainingData - total1 <- sizeD trainingData - total2 <- sizeD testingData + logInfo . display $ sformat "[Train] " - 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" + let acc_metric = Accuracy Nothing PredByArgmax 0 + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") + ce_metric = CrossEntropy Nothing True + (\_ p -> p ^?! ix 0) + (\b _ -> b ^?! ix "y") - 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 "" + forM_ (range 5) $ \ind -> do + logInfo .display $ sformat ("iteration " % int) ind + metrics <- newMetric "train" (ce_metric :* acc_metric :* MNil) + void $ forEachD_i trainingData $ \(i, (x, y)) -> askSession $ do + fitAndEval optm (M.fromList [("x", x), ("y", y)]) metrics - -- 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) + kv <- metricsToList metrics + lift $ mapM_ (uncurry neptLog) kv + + when (i `mod` 100 == 0) $ do + eval <- metricFormat metrics + logInfo . display $ sformat (int % "/" % int % " " % stext) i total eval + + metrics <- newMetric "val" (acc_metric :* MNil) + void $ forEachD_i testingData $ \(_, (x, y)) -> askSession $ do + pred <- forwardOnly (M.singleton "x" x) + void $ metricUpdate metrics (M.singleton "y" y) pred + + kv <- metricsToList metrics + mapM_ (uncurry neptLog) kv + eval <- metricFormat metrics + logInfo $ display eval + diff --git a/src/rcnn.hs b/src/rcnn.hs deleted file mode 100644 --- a/src/rcnn.hs +++ /dev/null @@ -1,214 +0,0 @@ -{-# 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