fei-modelzoo (empty) → 1.0.0
raw patch · 11 files changed
+1997/−0 lines, 11 filesdep +attoparsecdep +basedep +fei-base
Dependencies added: attoparsec, base, fei-base, fei-nn, formatting, lens, random-fu, repa, rio, text, transformers-base, vector
Files
- LICENSE +29/−0
- fei-modelzoo.cabal +58/−0
- src/MXNet/NN/ModelZoo/Lenet.hs +45/−0
- src/MXNet/NN/ModelZoo/RCNN/FPN.hs +76/−0
- src/MXNet/NN/ModelZoo/RCNN/FasterRCNN.hs +587/−0
- src/MXNet/NN/ModelZoo/RCNN/MaskRCNN.hs +119/−0
- src/MXNet/NN/ModelZoo/RCNN/RCNN.hs +325/−0
- src/MXNet/NN/ModelZoo/Resnet.hs +364/−0
- src/MXNet/NN/ModelZoo/Resnext.hs +199/−0
- src/MXNet/NN/ModelZoo/Utils/Box.hs +80/−0
- src/MXNet/NN/ModelZoo/VGG.hs +115/−0
+ LICENSE view
@@ -0,0 +1,29 @@+BSD 3-Clause License++Copyright (c) 2020, Jiasen Wu+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++1. Redistributions of source code must retain the above copyright notice, this+ list of conditions and the following disclaimer.++2. Redistributions in binary form must reproduce the above copyright notice,+ this list of conditions and the following disclaimer in the documentation+ and/or other materials provided with the distribution.++3. Neither the name of the copyright holder nor the names of its+ contributors may be used to endorse or promote products derived from+ this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ fei-modelzoo.cabal view
@@ -0,0 +1,58 @@+cabal-version: 2.4+name: fei-modelzoo+version: 1.0.0+synopsis: A collection of standard models+description: A collection of standard models+homepage: http://github.com/pierric/fei-modelzoo+license: BSD-3-Clause+license-file: LICENSE+author: Jiasen Wu+maintainer: jiasenwu@hotmail.com+copyright: 2020 - Jiasen Wu+category: Machine Learning, AI+build-type: Simple++Library+ exposed-modules: MXNet.NN.ModelZoo.Lenet+ MXNet.NN.ModelZoo.VGG+ MXNet.NN.ModelZoo.Resnet+ MXNet.NN.ModelZoo.Resnext+ MXNet.NN.ModelZoo.RCNN.FPN+ MXNet.NN.ModelZoo.RCNN.RCNN+ MXNet.NN.ModelZoo.RCNN.MaskRCNN+ MXNet.NN.ModelZoo.RCNN.FasterRCNN+ MXNet.NN.ModelZoo.Utils.Box+ hs-source-dirs: src+ ghc-options: -Wall+ default-language: Haskell2010+ default-extensions: GADTs,+ TypeFamilies,+ OverloadedLabels,+ OverloadedStrings,+ OverloadedLists,+ FlexibleContexts,+ FlexibleInstances,+ StandaloneDeriving,+ RecordWildCards,+ DataKinds,+ TypeOperators,+ TypeApplications,+ PartialTypeSignatures,+ LambdaCase,+ MultiWayIf,+ DoAndIfThenElse,+ TemplateHaskell,+ NoImplicitPrelude+ build-depends: base >= 4.7 && < 5.0+ , lens+ , transformers-base+ , repa+ , lens+ , random-fu+ , vector+ , text+ , rio+ , formatting+ , attoparsec+ , fei-base+ , fei-nn
+ src/MXNet/NN/ModelZoo/Lenet.hs view
@@ -0,0 +1,45 @@+module MXNet.NN.ModelZoo.Lenet where++import MXNet.Base+import MXNet.NN.Layer+import RIO++-- # 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 :: Layer SymbolHandle+symbol = do+ x <- variable "x"+ y <- variable "y"++ logit <- sequential "features" $ 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 <- 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 <- flatten p2++ v3 <- fullyConnected (#data := fl .& #num_hidden := 500 .& Nil)+ a3 <- activation (#data := v3 .& #act_type := #tanh .& Nil)++ fullyConnected (#data := a3 .& #num_hidden := 10 .& Nil)++ named "output" $ softmaxoutput (#data := logit .& #label := y .& Nil)
+ src/MXNet/NN/ModelZoo/RCNN/FPN.hs view
@@ -0,0 +1,76 @@+module MXNet.NN.ModelZoo.RCNN.FPN where++import RIO+import qualified RIO.NonEmpty as NE (reverse, unzip, zip, (<|))++import MXNet.Base (ArgOf (..), HMap (..),+ SymbolHandle, at', internals,+ prim, (.&))+import MXNet.Base.Operators.Tensor (_UpSampling)+import MXNet.NN.Layer++-- TODO+-- no_bias ?+-- batchnorm args ?++fpnFeatureExpander :: SymbolHandle -> NonEmpty (Text, Int) -> Layer (NonEmpty SymbolHandle)+fpnFeatureExpander sym output_layers = do+ sym <- internals sym+ layers <- mapM (at' sym) layer_names++ outputs <- liftIO $ newIORef (error "empty")+ sequential "fpn" $ do+ foldM_ (topDownPass outputs) Nothing (NE.zip layer_filters layers)+ -- return features bottom-up (from big to small)+ liftIO $ readIORef outputs++ where+ (layer_names, layer_filters) = NE.unzip $ NE.reverse output_layers+ topDownPass outputs Nothing (nflt, layer) = subscope_next_name $ unique' $ do+ y <- named "conv1" $+ convolution (#data := layer+ .& #num_filter := nflt+ .& #kernel := [1,1]+ .& #pad := [0, 0]+ .& #stride := [1,1]+ .& #no_bias := True .& Nil)+ y <- named "bn0" $+ batchnorm (#data := y .& Nil)+ out <- named "conv2" $+ convolution (#data := y+ .& #num_filter := nflt+ .& #kernel := [3, 3]+ .& #pad := [1, 1]+ .& #stride := [1,1]+ .& #no_bias := True .& Nil)+ out <- named "bn1" $+ batchnorm (#data := out .& Nil)+ writeIORef outputs [out]+ return (Just y)+ topDownPass outputs (Just prev) (nflt, layer) = subscope_next_name $ unique' $ do+ y <- named "conv1" $+ convolution (#data := layer+ .& #num_filter := nflt+ .& #kernel := [1,1]+ .& #pad := [0, 0]+ .& #stride := [1,1]+ .& #no_bias := True .& Nil)+ y <- named "bn0" $+ batchnorm (#data := y .& Nil)+ prev_up <- prim _UpSampling+ (#data := [prev]+ .& #num_args := 1+ .& #scale := 2+ .& #sample_type := #nearest .& Nil)+ y <- add_ prev_up y+ out <- named "conv2" $+ convolution (#data := y+ .& #num_filter := nflt+ .& #kernel := [3, 3]+ .& #pad := [1, 1]+ .& #stride := [1,1]+ .& #no_bias := True .& Nil)+ out <- named "bn1" $+ batchnorm (#data := out .& Nil)+ modifyIORef outputs (out NE.<|)+ return (Just y)
+ src/MXNet/NN/ModelZoo/RCNN/FasterRCNN.hs view
@@ -0,0 +1,587 @@+module MXNet.NN.ModelZoo.RCNN.FasterRCNN where++import RIO+import RIO.List (unzip3, zip3, zip4)+import RIO.List.Partial (head, last)+import qualified RIO.NonEmpty as NE (toList)++import MXNet.Base+import MXNet.Base.Operators.Tensor (_Custom, _MakeLoss, __arange,+ __contrib_AdaptiveAvgPooling2D,+ __contrib_ROIAlign,+ __contrib_box_decode,+ __contrib_box_nms, __zeros,+ _add_n, _clip,+ _repeat, _sigmoid,+ _smooth_l1, _transpose)+import MXNet.NN.Layer+import MXNet.NN.ModelZoo.RCNN.FPN+import MXNet.NN.ModelZoo.RCNN.RCNN+import qualified MXNet.NN.ModelZoo.Resnet as Resnet+import qualified MXNet.NN.ModelZoo.VGG as VGG++data Backbone = VGG16+ | RESNET50+ | RESNET101+ | RESNET50FPN+ deriving (Show, Read, Eq)++data RcnnConfiguration = RcnnConfigurationTrain+ { backbone :: Backbone+ , batch_size :: Int+ , feature_strides :: [Int]+ , pretrained_weights :: String+ , bbox_reg_std :: (Float, Float, Float, Float)+ , rpn_anchor_scales :: [Int]+ , rpn_anchor_ratios :: [Float]+ , rpn_anchor_base_size :: Int+ , rpn_pre_topk :: Int+ , rpn_post_topk :: Int+ , rpn_nms_thresh :: Float+ , rpn_min_size :: Int+ , rpn_batch_rois :: Int+ , rpn_fg_fraction :: Float+ , rpn_fg_overlap :: Float+ , rpn_bg_overlap :: Float+ , rpn_allowd_border :: Int+ , rcnn_num_classes :: Int+ , rcnn_pooled_size :: Int+ , rcnn_batch_rois :: Int+ , rcnn_fg_fraction :: Float+ , rcnn_fg_overlap :: Float+ , rcnn_max_num_gt :: Int+ }+ | RcnnConfigurationInference+ { backbone :: Backbone+ , batch_size :: Int+ , feature_strides :: [Int]+ , checkpoint :: String+ , bbox_reg_std :: (Float, Float, Float, Float)+ , rpn_anchor_scales :: [Int]+ , rpn_anchor_ratios :: [Float]+ , rpn_anchor_base_size :: Int+ , rpn_pre_topk :: Int+ , rpn_post_topk :: Int+ , rpn_nms_thresh :: Float+ , rpn_min_size :: Int+ , rcnn_num_classes :: Int+ , rcnn_pooled_size :: Int+ , rcnn_batch_rois :: Int+ , rcnn_force_nms :: Bool+ , rcnn_nms_thresh :: Float+ , rcnn_topk :: Int+ }+ deriving Show++data FasterRCNN = FasterRCNN+ { _rpn_loss :: (SymbolHandle, SymbolHandle, SymbolHandle)+ , _box_loss :: (SymbolHandle, SymbolHandle)+ , _cls_targets :: SymbolHandle+ , _roi_boxes :: SymbolHandle+ , _gt_matches :: SymbolHandle+ , _positive_indices :: SymbolHandle+ , _top_feature :: SymbolHandle+ }+ | FasterRCNNInferenceOnly+ { _top_feature :: SymbolHandle+ , _cls_ids :: SymbolHandle+ , _scores :: SymbolHandle+ , _boxes :: SymbolHandle+ }++stageList :: Backbone -> [Int]+stageList RESNET50FPN = [2..5]+stageList _ = [3]++resnet50Args = (#num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,6,3]+ .& #bottle_neck := True+ .& #workspace := 256+ .& Nil)++resnet101Args = (#num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,23,3]+ .& #bottle_neck := True+ .& #workspace := 256+ .& Nil)++features1 :: Backbone -> SymbolHandle -> Layer (NonEmpty SymbolHandle)+features1 VGG16 dat = fmap (:| []) $ VGG.getFeature dat [2, 2, 3, 3, 3] [64, 128, 256, 512, 512] False False+features1 RESNET50 dat = fmap (:| []) $ Resnet.getFeature dat resnet50Args+features1 RESNET101 dat = fmap (:| []) $ Resnet.getFeature dat resnet101Args+features1 RESNET50FPN dat = do+ sym <- Resnet.getFeature dat resnet50Args+ sym <- Resnet.getTopFeature sym resnet50Args+ fpnFeatureExpander sym+ [ ("features.5.2.plus_output", 256)+ , ("features.6.3.plus_output", 256)+ , ("features.7.5.plus_output", 256)+ , ("features.8.2.plus_output", 256) ]++features2 :: Backbone -> SymbolHandle -> Layer SymbolHandle+features2 VGG16 dat = VGG.getTopFeature dat+features2 RESNET50 dat = Resnet.getTopFeature dat resnet50Args+features2 RESNET101 dat = Resnet.getTopFeature dat resnet101Args+features2 RESNET50FPN dat = return dat++rpn :: RcnnConfiguration+ -> NonEmpty SymbolHandle -> SymbolHandle+ -> Layer (SymbolHandle, SymbolHandle, SymbolHandle, SymbolHandle)+rpn conf convFeats imInfo = unique "rpn" $ do+ conv3x3_feat <- named "rpn_conv_3x3" $+ convolutionShared (#kernel := [3,3]+ .& #pad := [1,1]+ .& #num_filter := 512 .& Nil)+ conv1x1_cls <- named "rpn_cls_score" $+ convolutionShared (#kernel := [1,1]+ .& #pad := [0,0]+ .& #num_filter := num_variation .& Nil)+ conv1x1_reg <- named "rpn_bbox_pred" $+ convolutionShared (#kernel := [1,1]+ .& #pad := [0,0]+ .& #num_filter := 4 * num_variation .& Nil)+ layers <- zipWithM (rpn_head conv3x3_feat conv1x1_cls conv1x1_reg)+ (NE.toList convFeats)+ (feature_strides conf)++ let (rpn_pres, rpn_raw_scores, rpn_raw_boxregs) = unzip3 layers+ -- concat the list of RPN ROI boxes, in which boxes are decoded, and suppressed items are -1s+ -- result shape: (batch_size, Σ(feat_H_i*feat_W_i*num_variation), 5)+ rpn_pres <- concat_ 1 rpn_pres+ -- concat the list of RPN raw scores of all predictions,+ -- result shape: (batch_size, Σ(feat_H_i*feat_W_i*num_variation), 1)+ rpn_raw_scores <- concat_ 1 rpn_raw_scores+ -- concat the list of RPN raw box regression of all predictions,+ -- result shape: (batch_size, Σ(feat_H_i*feat_W_i*num_variation), 4)+ rpn_raw_boxregs <- concat_ 1 rpn_raw_boxregs++ -- non-maximum suppress the rois, and split the score and box part+ (rpn_roi_scores, rpn_roi_boxes) <- nms rpn_pres++ return (rpn_roi_scores, rpn_roi_boxes, rpn_raw_scores, rpn_raw_boxregs)++ where+ num_variation = length (rpn_anchor_scales conf) * length (rpn_anchor_ratios conf)+ rpn_head conv3x3_feat conv1x1_cls conv1x1_reg feat stride = do+ x <- conv3x3_feat feat+ x <- activation (#data := x+ .& #act_type := #relu .& Nil)++ anchors <- prim _Custom (#op_type := "anchor_generator"+ .& #data := [x]+ .& #stride :≅ stride+ .& #scales :≅ rpn_anchor_scales conf+ .& #ratios :≅ rpn_anchor_ratios conf+ .& #base_size :≅ rpn_anchor_base_size conf+ .& #alloc_size :≅ ((128, 128) :: (Int, Int))+ .& Nil)++ rpn_raw_score <- conv1x1_cls x+ -- (batch_size, num_variation, H, W) ==> (batch_size, H, W, num_variation)+ rpn_raw_score <- prim _transpose (#data := rpn_raw_score .& #axes := [0, 2, 3, 1] .& Nil)+ -- (batch_size, H, W, num_variation) ==> (batch_size, H*W*num_variation, 1)+ rpn_raw_score <- reshape [0, -1, 1] rpn_raw_score++ rpn_cls_score <- blockGrad =<< prim _sigmoid (#data := rpn_raw_score .& Nil)+ rpn_cls_score <- reshape [0, -1, 1] rpn_cls_score++ rpn_raw_boxreg <- conv1x1_reg x+ -- (batch_size, num_variation * 4, H, W) ==> (batch_size, H, W, num_variation * 4)+ rpn_raw_boxreg <- prim _transpose (#data := rpn_raw_boxreg .& #axes := [0, 2, 3, 1] .& Nil)+ -- (batch_size, H, W, num_variation * 4) ==> (batch_size, H*W*num_variation, 4)+ rpn_raw_boxreg <- reshape [0, -1, 4] rpn_raw_boxreg++ rpn_pre <- region_proposer (fromIntegral $ rpn_min_size conf)+ anchors+ rpn_raw_boxreg+ rpn_cls_score+ (1, 1, 1, 1)++ return (rpn_pre, rpn_raw_score, rpn_raw_boxreg)++ region_proposer min_size anchors boxregs scores stds = do+ let (std0, std1, std2, std3) = stds+ rois <- prim __contrib_box_decode (#data := boxregs+ .& #anchors := anchors+ .& #format := #corner+ .& #std0 := std0+ .& #std1 := std1+ .& #std2 := std2+ .& #std3 := std3+ .& Nil)+ (xmin, ymin, xmax, ymax) <- bbox_clip_to_image rois imInfo+ width <- sub_ xmax xmin >>= addScalar 1+ height <- sub_ ymax ymin >>= addScalar 1+ invalid <- ltScalar min_size width >>= \c1 ->+ ltScalar min_size height >>= \c2 ->+ or_ c1 c2+ mask <- onesLike invalid >>= mulScalar (-1)+ scores <- where_ invalid mask scores+ invalid <- broadcastAxis [2] [4] invalid+ mask <- onesLike invalid >>= mulScalar (-1)+ rois <- concat_ (-1) [xmin, ymin, xmax, ymax]+ rois <- where_ invalid mask rois+ blockGrad =<< named "proposals" (concat_ (-1) [scores, rois])++ nms rpn_pre = do+ -- rpn_pre shape: (batch_size, num_anchors, 5)+ -- in dim 3: [score, xmin, ymin, xmax, ymax]+ tmp <- prim __contrib_box_nms (#data := rpn_pre+ .& #overlap_thresh := rpn_nms_thresh conf+ .& #topk := rpn_pre_topk conf+ .& #coord_start := 1+ .& #score_index := 0+ .& #id_index := (-1)+ .& #force_suppress := True .& Nil)+ tmp <- slice_axis tmp 1 0 (Just $ rpn_post_topk conf)+ rpn_roi_scores <- blockGrad =<<+ slice_axis tmp (-1) 0 (Just 1)+ rpn_roi_boxes <- blockGrad =<<+ slice_axis tmp (-1) 1 Nothing+ return (rpn_roi_scores, rpn_roi_boxes)++ bbox_clip_to_image rois info = do+ -- rois: (B, N, 4)+ -- info: (B, 3)+ -- return: (B,N), (B,N), (B,N), (B,N)+ [xmin, ymin, xmax, ymax] <- splitBySections 4 (-1) False rois+ [height, width, _] <- splitBySections 3 (-1) False info+ height <- expandDims (-1) height+ width <- expandDims (-1) width+ w_ub <- subScalar 1 width+ h_ub <- subScalar 1 height+ z <- zerosLike xmin+ w <- onesLike xmin >>= mulBroadcast w_ub+ h <- onesLike xmin >>= mulBroadcast h_ub+ cond <- ltScalar 0 xmin+ xmin <- where_ cond z xmin+ cond <- ltScalar 0 ymin+ ymin <- where_ cond z ymin+ cond <- gtBroadcast xmax w_ub+ xmax <- where_ cond w xmax+ cond <- gtBroadcast ymax h_ub+ ymax <- where_ cond h ymax+ return (xmin, ymin, xmax, ymax)++alignROIs :: NonEmpty SymbolHandle -> SymbolHandle -> [Int] -> Int -> [Int] -> Layer _+alignROIs features rois stage_indices roi_pooled_size strides = do+ -- rois: (N, 5), batch_index, min_x, min_y, max_x, max_y+ let min_stage = head stage_indices+ max_stage = last stage_indices+ [_, xmin, ymin, xmax, ymax] <- splitBySections 5 (-1) False rois+ w <- xmax `sub_` xmin >>= addScalar 1+ h <- ymax `sub_` ymin >>= addScalar 1+ -- heuristic to compute the stage where each rois box fits+ -- bigger box in higher stage+ -- smaller box in lower stage+ roi_level_raw <- w `mul_` h >>=+ sqrt_ >>=+ divScalar 224 >>=+ addScalar 1e-6 >>=+ log2_ >>=+ addScalar 4 >>=+ floor_+ roi_level <- prim _clip (#data := roi_level_raw+ .& #a_min := fromIntegral min_stage+ .& #a_max := fromIntegral max_stage .& Nil)+ >>= squeeze Nothing+ let align (lvl, feat, stride) = do+ cond <- eqScalar (fromIntegral lvl) roi_level+ omit <- onesLike rois >>= mulScalar (-1)+ masked <- where_ cond rois omit+ prim __contrib_ROIAlign (#data := feat+ .& #rois := masked+ .& #pooled_size := [roi_pooled_size, roi_pooled_size]+ .& #spatial_scale := 1 / fromIntegral stride+ .& #sample_ratio := 2 .& Nil)+ features <- mapM align $ zip3 [max_stage,max_stage-1..min_stage] (NE.toList features) strides+ prim _add_n (#args := features .& Nil)+++graphT :: RcnnConfiguration -> Layer (FasterRCNN, SymbolHandle)+graphT conf@RcnnConfigurationTrain{..} = do+ -- dat: (B, image_height, image_width)+ dat <- variable "data"+ -- imInfo: (B, 3,)+ imInfo <- variable "im_info"+ -- gt_boxes: (B, M, 5), the last dim: min_x, min_y, max_x, max_y, class_id (background class: 0)+ gt_boxes <- variable "gt_boxes"+ rpn_cls_targets <- variable "rpn_cls_targets"+ rpn_box_targets <- variable "rpn_box_targets"+ rpn_box_masks <- variable "rpn_box_masks"++ gt_labels <- unique' $ slice_axis gt_boxes (-1) 4 Nothing+ gt_boxes <- unique' $ slice_axis gt_boxes (-1) 0 (Just 4)++ let (std0, std1, std2, std3) = bbox_reg_std+ bbox_reg_mean <- named "bbox_reg_mean" $ prim __zeros (#shape := [4] .& Nil)+ bbox_reg_std <- named "bbox_reg_std" $ constant [4] [std0, std1, std2, std3]++ sequential "features" $ do+ feats <- features1 backbone dat++ (rois_scores, roi_boxes, rpn_raw_scores, rpn_raw_boxregs) <- rpn conf feats imInfo++ -- total number of ROIs in a batch+ -- batch_size: number of images+ -- rcnn_batch_rois: number of rois per image+ (feat_aligned, roi_boxes, samples, matches) <- unique "rcnn" $ do+ (rois_boxes, samples, matches) <- rcnnSampler batch_size+ rpn_post_topk+ rcnn_batch_rois+ rcnn_fg_overlap+ rcnn_fg_fraction+ rcnn_max_num_gt+ roi_boxes+ rois_scores+ gt_boxes++ roi_batchid <- prim __arange (#start := 0 .& #stop := Just (fromIntegral batch_size) .& Nil)+ roi_batchid <- prim _repeat (#data := roi_batchid .& #repeats := rcnn_batch_rois .& Nil)+ roi_batchid <- reshape [-1, 1] roi_batchid+ -- rois: (B * rcnn_batch_rois, 4)+ rois <- reshape [-1, 4] rois_boxes+ rois <- concat_ 1 [roi_batchid, rois] >>= blockGrad+ feat <- alignROIs feats rois (stageList backbone) rcnn_pooled_size feature_strides+ return (feat, rois_boxes, samples, matches)++ -- feat_aligned: (batch_size * rcnn_batch_rois, num_channels, feature_height, feature_width)+ -- TODO num_channels is set to rcnn_batch_rois, it is coincidance or on purpose?+ -- Apply the remaining feature extraction layers+ top_feat <- features2 backbone feat_aligned++ unique "rcnn" $ do+ (cls_targets, bbox_targets, bbox_masks, positive_indices) <-+ bboxTargetGenerator batch_size+ (rcnn_num_classes-1)+ (floor $ rcnn_fg_fraction * fromIntegral rcnn_batch_rois)+ samples+ matches+ roi_boxes+ gt_labels+ gt_boxes+ bbox_reg_mean+ bbox_reg_std++ -- sigmoid + binary-cross-entropy+ -- rpn_raw_scores: (B, num_rois, 1)+ -- rpn_cls_targets: (B, num_rois, 1)+ rpn_cls_prob <- prim _sigmoid (#data := rpn_raw_scores .& Nil)+ sample_mask <- geqScalar 0 rpn_cls_targets+ rpn_cls_loss <- sigmoidBCE rpn_raw_scores rpn_cls_targets (Just sample_mask) AggSum+ -- a <- log2_ rpn_cls_prob+ -- ra <- log2_ =<< rsubScalar 1 rpn_cls_prob+ -- b <- identity rpn_cls_targets+ -- rb <- rsubScalar 1 rpn_cls_targets+ -- rpn_cls_loss <- (join $ liftM2 add_ (mul_ a b) (mul_ ra rb)) >>= rsubScalar 0++ -- average number of targets per batch example+ cls_mask <- geqScalar 0 rpn_cls_targets+ num_pos_avg <- sum_ cls_mask Nothing False >>= divScalar (fromIntegral batch_size) >>= addScalar 1e-14++ rpn_cls_loss <- divBroadcast rpn_cls_loss num_pos_avg+ rpn_cls_loss <- prim _MakeLoss (#data := rpn_cls_loss .& #grad_scale := 1.0 .& Nil)++ rpn_bbox_reg <- sub_ rpn_raw_boxregs rpn_box_targets+ rpn_bbox_reg <- prim _smooth_l1 (#data := rpn_bbox_reg .& #scalar := 3.0 .& Nil)+ rpn_bbox_loss <- mul_ rpn_bbox_reg rpn_box_masks >>= flip divBroadcast num_pos_avg+ rpn_bbox_loss <- prim _MakeLoss+ (#data := rpn_bbox_loss .& #grad_scale := 1.0 .& Nil)++ box_feat <- prim __contrib_AdaptiveAvgPooling2D (#data := top_feat .& #output_size := [7, 7] .& Nil)+ -- box_feat <- pooling (#data := top_feat+ -- .& #kernel := [3,3]+ -- .& #stride := [2,2]+ -- .& #pad := [1,1]+ -- .& #pool_type := #avg .& Nil)+ -- box_feat <- named "rcnn_cls_score_fc" $+ -- fullyConnected (#data := box_feat .& #num_hidden := 1024 .& Nil)+ box_feat <- activation (#data := box_feat .& #act_type := #relu .& Nil)++ -- rcnn class prediction+ -- cls_score: (batch_size * rcnn_batch_rois, rcnn_num_classes)+ cls_score <- named "rcnn_cls_score" $+ fullyConnected (#data := box_feat .& #num_hidden := rcnn_num_classes .& Nil)+ cls_score <- reshape [batch_size, rcnn_batch_rois, rcnn_num_classes] cls_score++ -- `preserve_shape = True` makes softmax on the last dim.+ -- `normalization = valid` divides the loss by the number of valid items.+ -- we actually want to divide by average number of valid items in the batch,+ -- so scale up by the size batch_size+ cls_prob <- named "rcnn_cls_prob" $+ softmaxoutput (#data := cls_score+ .& #label := cls_targets+ .& #preserve_shape := True+ .& #use_ignore := True+ .& #ignore_label := -1+ .& #normalization := #valid+ .& #grad_scale := fromIntegral batch_size .& Nil)++ ---------------------------+ -- bbox_loss part+ --+ bbox_feature <- named "rcnn_bbox_feature" $ fullyConnected (#data := top_feat .& #num_hidden := 1024 .& Nil)+ bbox_feature <- activation (#data := bbox_feature .& #act_type := #relu .& Nil)+ -- bbox_feature: (B * rcnn_batch_rois, num_hidden) ==> (B, rcnn_batch_rois, num_hidden)+ bbox_feature <- expandDims 0 bbox_feature >>= reshape [batch_size, rcnn_batch_rois, 1024]+ -- select only feature that has foreground gt for each batch example+ -- positive_indices: (B, rcnn_fg_fraction * num_sample)+ bbox_feature <- forM ([0..batch_size-1] :: [_]) $ \i -> do+ ind <- slice_axis positive_indices 0 i (Just (i+1)) >>= squeeze Nothing+ bat <- slice_axis bbox_feature 0 i (Just (i+1)) >>= squeeze Nothing+ takeI ind bat+ bbox_feature <- concat_ 0 bbox_feature++ -- for each foreground ROI, predict boxes (reg) for each foreground class+ avg_valid_pred <- gtScalar (-1) cls_targets+ >>= \s -> sum_ s Nothing False+ >>= divScalar (fromIntegral batch_size)+ >>= addScalar 1e-14++ bbox_pred <- named "rcnn_bbox_pred" $+ fullyConnected (#data := bbox_feature .& #num_hidden := 4 * (rcnn_num_classes - 1) .& Nil)+ -- bbox_pred: (B * rcnn_fg_fraction * num_sample, num_fg_classes * 4)+ -- ==> (B, rcnn_fg_fraction * num_sample, num_fg_classes, 4)+ bbox_pred <- reshape [batch_size, -1, rcnn_num_classes - 1, 4] bbox_pred+ bbox_reg <- sub_ bbox_pred bbox_targets+ bbox_reg <- prim _smooth_l1 (#data := bbox_reg .& #scalar := 1.0 .& Nil)+ bbox_loss <- mul_ bbox_reg bbox_masks >>= flip divBroadcast avg_valid_pred+ bbox_loss <- prim _MakeLoss (#data := bbox_loss .& #grad_scale := 1.0 .& Nil)++ cls_targets <- reshape [batch_size, -1] cls_targets >>= blockGrad+ box_targets <- blockGrad bbox_targets+ rpn_cls_prob <- blockGrad rpn_cls_prob++ result_sym <- group $ [rpn_cls_prob, rpn_cls_loss, rpn_bbox_loss, cls_prob, bbox_loss, cls_targets]++ return $ (FasterRCNN {+ _rpn_loss = (rpn_cls_prob, rpn_cls_loss, rpn_bbox_loss),+ _box_loss = (cls_prob, bbox_loss),+ _cls_targets = cls_targets,+ _roi_boxes = roi_boxes,+ _gt_matches = matches,+ _positive_indices = positive_indices,+ _top_feature = top_feat+ }, result_sym)+++graphI :: RcnnConfiguration -> Layer (FasterRCNN, SymbolHandle)+graphI conf@RcnnConfigurationInference{..} = do+ -- dat: (B, image_height, image_width)+ dat <- variable "data"+ -- imInfo: (B, 3,)+ imInfo <- variable "im_info"++ sequential "features" $ do+ feats <- features1 backbone dat++ -- roi_boxes: (B, rpn_post_topk, 4), rpn predicted boxes (x0,y0,x1,y1)+ (roi_scores, roi_boxes, _, _) <- rpn conf feats imInfo++ roi_batchid <- prim __arange (#start := 0 .& #stop := Just (fromIntegral batch_size) .& Nil)+ roi_batchid <- prim _repeat (#data := roi_batchid .& #repeats := rpn_post_topk .& Nil)+ roi_batchid <- reshape [-1, 1] roi_batchid+ -- rois: (B * rpn_post_topk, 4)+ rois <- reshape [-1, 4] roi_boxes+ rois <- concat_ 1 [roi_batchid, rois]+ feat_aligned <- alignROIs feats rois (stageList backbone) rcnn_pooled_size feature_strides++ top_feat <- features2 backbone feat_aligned++ sequential "rcnn" $ do+ box_feat <- prim __contrib_AdaptiveAvgPooling2D (#data := top_feat .& #output_size := [7, 7] .& Nil)+ box_feat <- activation (#data := box_feat .& #act_type := #relu .& Nil)++ -- rcnn class prediction+ -- cls_score: (batch_size * rpn_post_topk, rcnn_num_classes)+ cls_score <- named "rcnn_cls_score" $+ fullyConnected (#data := box_feat .& #num_hidden := rcnn_num_classes .& Nil)+ cls_score <- reshape [batch_size, rpn_post_topk, rcnn_num_classes] cls_score++ cls_prob <- softmax (#data := cls_score .& #axis := (-1) .& Nil)++ bbox_feature <- named "rcnn_bbox_feature" $ fullyConnected (#data := top_feat .& #num_hidden := 1024 .& Nil)+ -- bbox_feature: (B * rpn_post_topk, num_hidden)+ bbox_feature <- activation (#data := bbox_feature .& #act_type := #relu .& Nil)++ bbox_pred <- named "rcnn_bbox_pred" $+ fullyConnected (#data := bbox_feature+ .& #num_hidden := 4 * (rcnn_num_classes - 1) .& Nil)+ -- bbox_pred: (B * rpn_post_topk, num_fg_classes * 4)+ -- ==> (B, rpn_post_topk, num_fg_classes, 4)+ bbox_pred <- reshape [batch_size, -1, rcnn_num_classes - 1, 4] bbox_pred++ -------------------+ -- decode the classes and boxes+ --++ -- `cls_prob` predicts `rcnn_num_classes` classes, background class at 0.+ (clsids, scores) <- multiClassDecodeWithClsId rcnn_num_classes (-1) 0.01 cls_prob++ -- tranpose the bbox_pred, clsids, scores, because bbox_nms does suppresion on+ -- the last two dimensions+ bbox_pred <- transpose bbox_pred [0, 2, 1, 3]+ clsids <- transpose clsids [0, 2, 1]+ scores <- transpose scores [0, 2, 1]++ -- `roi_boxes` B x (1, rpn_post_topk, 4)+ roi_boxes <- splitBySections batch_size 0 False roi_boxes+ -- `bbox_pred` B x (num_fg_classes, rpn_post_topk, 4)+ bbox_pred <- splitBySections batch_size 0 True bbox_pred+ -- `bbox_clsids` B x (num_fg_classes, rpn_post_topk)+ bbox_clsids <- splitBySections batch_size 0 True clsids+ -- `bbox_scores` B x (num_fg_classes, rpn_post_topk)+ bbox_scores <- splitBySections batch_size 0 True scores+++ let (std0, std1, std2, std3) = bbox_reg_std++ results <- forM (zip4 roi_boxes bbox_pred bbox_clsids bbox_scores) $ \ (roi, pred, clsid, score) -> do+ bbox_decoded <- prim __contrib_box_decode (#data := pred+ .& #anchors := roi+ .& #format := #corner+ .& #std0 := std0+ .& #std1 := std1+ .& #std2 := std2+ .& #std3 := std3+ .& Nil)++ -- concatenate all clsid, score, and box+ -- res: (num_fg_classes, rpn_post_topk, 6)+ clsid <- expandDims (-1) clsid+ score <- expandDims (-1) score+ res <- concat_ (-1) [clsid, score, bbox_decoded]++ -- if force_nms, we do nms all boxes from all classes+ res <- if rcnn_force_nms+ then reshape [1, -1, 0] res+ else return res++ res <- prim __contrib_box_nms (#data := res+ .& #overlap_thresh := rcnn_nms_thresh+ .& #valid_thresh := 0.001+ .& #topk := rcnn_topk+ .& #coord_start := 2+ .& #score_index := 1+ .& #id_index := 0+ .& #force_suppress := rcnn_force_nms .& Nil)+ res <- slice_axis res 1 0 (Just rcnn_topk)+ -- final result: (num_fg_classes * rcnn_topk, 6)+ reshape [-3, 0] res++ results <- stack 0 results+ result_cls_ids <- slice_axis results (-1) 0 (Just 1)+ result_scores <- slice_axis results (-1) 1 (Just 2)+ result_boxes <- slice_axis results (-1) 2 Nothing++ res_sym <- group [result_cls_ids, result_scores, result_boxes]+ let res_data = FasterRCNNInferenceOnly+ { _top_feature = top_feat+ , _cls_ids = result_cls_ids+ , _scores = result_scores+ , _boxes = result_boxes+ }+ return $ (res_data, res_sym)+
+ src/MXNet/NN/ModelZoo/RCNN/MaskRCNN.hs view
@@ -0,0 +1,119 @@+module MXNet.NN.ModelZoo.RCNN.MaskRCNN where++import RIO++import MXNet.Base+import qualified MXNet.Base.Operators.Tensor as T+import MXNet.NN.Layer+import qualified MXNet.NN.ModelZoo.RCNN.FasterRCNN as FasterRCNN+import MXNet.NN.ModelZoo.RCNN.RCNN+++data MaskRCNN = MaskRCNN+ { _faster_rcnn_result :: FasterRCNN.FasterRCNN+ , _masks_loss :: SymbolHandle+ }+ | MaskRCNNInferenceOnly+ { _faster_rcnn_result :: FasterRCNN.FasterRCNN+ , _masks :: SymbolHandle+ }++maskHead :: SymbolHandle -> Int -> Int -> Int -> Int -> Layer SymbolHandle+maskHead top_feat num_fcn_conv num_fg_classes batch_size num_mask_channels = do+ -- top_feat: The network input tensor of shape (B * N, fC, fH, fW).+ --+ -- returns:+ -- Mask prediction of shape (B, N, C, MS, MS)+ feat <- sequential "conv" $ foldM one_conv top_feat ([1..num_fcn_conv] :: [_])+ feat <- named "conv-transposed" $+ prim T._Deconvolution (#data := feat+ .& #num_filter := num_mask_channels+ .& #kernel := [2, 2]+ .& #stride := [2, 2]+ .& #pad := [0, 0] .& Nil)+ feat <- activation (#data := feat .& #act_type := #relu .& Nil)+ mask <- named "conv-last" $+ convolution (#data := feat+ .& #kernel := [1, 1]+ .& #num_filter := num_fg_classes+ .& #stride := [1, 1]+ .& #pad := [0, 0] .& Nil)+ reshape [-4, batch_size, -1, 0, 0, 0] mask++ where+ one_conv x _ = do+ x <- convolution (#data := x+ .& #num_filter := num_mask_channels+ .& #kernel := [3, 3]+ .& #stride := [1, 1]+ .& #pad := [1, 1] .& Nil)+ activation (#data := x .& #act_type := #relu .& Nil)+++graphT :: FasterRCNN.RcnnConfiguration -> Layer (MaskRCNN, SymbolHandle)+graphT conf@(FasterRCNN.RcnnConfigurationTrain{..}) = do+ gt_masks <- variable "gt_masks"+ (fr@FasterRCNN.FasterRCNN{..}, fr_outputs) <- FasterRCNN.graphT conf+ unique "mask" $ do+ let take_pos t = do+ ts <- forM ([0..batch_size-1] :: [_]) $ \i -> do+ ind <- slice_axis _positive_indices 0 i (Just (i+1)) >>= squeeze Nothing+ bat <- slice_axis t 0 i (Just (i+1)) >>= squeeze Nothing+ takeI ind bat+ concat_ 0 ts++ -- like box_feature, we select only layers of the feature/... that have+ -- foreground gt, for each example in the batch.+ -- positive_indices: (B, rcnn_fg_fraction * rcnn_batch_rois)+ -- _top_feature: (B * rcnn_batch_rois, num_channels, rcnn_pooled_size, rcnn_pooled_size)+ -- => (B * rcnn_fg_fraction * rcnn_batch_rois, .., .., ..)+ -- _roi_boxes: (B, rcnn_batch_rois, 4) => (B, rcnn_fg_fraction * rcnn_batch_rois, 4)+ -- _gt_matches: (B, rcnn_batch_rois) => (B, rcnn_fg_fraction * rcnn_batch_rois)+ -- _cls_targets: (B, rcnn_batch_rois) => (B, rcnn_fg_fraction * rcnn_batch_rois)+ feature <- take_pos =<< reshape [batch_size, -1, 0, 0, 0] =<< expandDims 0 _top_feature+ roi_boxes <- reshape [batch_size, -1, 4] =<< take_pos _roi_boxes+ gt_matches <- reshape [batch_size, -1] =<< take_pos _gt_matches+ cls_targets <- reshape [batch_size, -1] =<< take_pos _cls_targets++ let num_fcn_conv = case backbone of+ FasterRCNN.RESNET50FPN -> 4+ _ -> 0+ num_fg_classes = rcnn_num_classes-1+ -- mask_size should be twice the final feature size+ -- becuase there is only one Conv2DTranspose layer+ mask_size = rcnn_pooled_size * 2+ masks <- unique "mask_head" $ maskHead feature num_fcn_conv num_fg_classes batch_size 256++ (mask_targets, mask_weights) <- unique "target_gen" $+ maskTargetGenerator batch_size+ num_fg_classes+ mask_size+ gt_masks+ roi_boxes+ gt_matches+ cls_targets+ masks_loss <- unique "loss" $ do+ masks_loss <- sigmoidBCE masks mask_targets (Just mask_weights) AggSum+ num_pos_avg <- sum_ mask_weights Nothing False >>= divScalar (fromIntegral batch_size) >>= addScalar 1e-14+ masks_loss <- divBroadcast masks_loss num_pos_avg+ prim T._MakeLoss (#data := masks_loss .& #grad_scale := 1.0 .& Nil)++ result_sym <- group $ [fr_outputs, masks_loss]+ return $ (MaskRCNN {+ _faster_rcnn_result = fr,+ _masks_loss = masks_loss+ }, result_sym)++graphI :: FasterRCNN.RcnnConfiguration -> Layer (MaskRCNN, SymbolHandle)+graphI conf@(FasterRCNN.RcnnConfigurationInference{..}) = do+ (fr@FasterRCNN.FasterRCNNInferenceOnly{..}, fr_outputs) <- FasterRCNN.graphI conf+ feature <- reshape [batch_size, -1, 0, 0, 0] =<< expandDims 0 _top_feature+ let num_fcn_conv = case backbone of+ FasterRCNN.RESNET50FPN -> 4+ _ -> 0+ num_fg_classes = rcnn_num_classes-1+ masks <- unique "mask_head" $ maskHead feature num_fcn_conv num_fg_classes batch_size 256+ masks <- prim T._sigmoid (#data := masks .& Nil)+ res_sym <- group [fr_outputs, masks]+ let res_data = MaskRCNNInferenceOnly fr masks+ return (res_data, res_sym)
+ src/MXNet/NN/ModelZoo/RCNN/RCNN.hs view
@@ -0,0 +1,325 @@+module MXNet.NN.ModelZoo.RCNN.RCNN where++import RIO+import RIO.List (unzip, unzip3, unzip4, zip4)++import MXNet.Base+import qualified MXNet.Base.Operators.Tensor as T+import MXNet.NN.Layer+++rcnnSampler :: Int -> Int -> Int -> Float -> Float -> Int+ -> SymbolHandle -> SymbolHandle -> SymbolHandle+ -> Layer _+rcnnSampler batch_size num_proposal num_sample fg_overlap fg_fraction max_num_gt+ rois scores gt_boxes = do+ -- B: batch_size, N: num_proposal (post-topk), S: num_sample (rcnn_batch_rois)+ -- rois: (B,N,4), min_x, min_y, max_x, max_y+ -- scores: (B,N,1), value range [0,1], -1 for being ignored+ -- gt_boxes: (B,M,4), min_x, min_y, max_x, max_y+ -- return:+ -- rois: (B,S,4)+ -- samples:(B,S), value -1 (negative), 0 (ignore), 1 (positive)+ -- matches:(B,S), value [0, M)+ (rois, samples, matches) <- unzip3 <$> mapM sampler [0..batch_size-1]++ rois <- stack 0 rois+ samples <- stack 0 samples+ matches <- stack 0 matches++ return (rois, samples, matches)++ where+ sampler batch_index = do+ roi <- getBatch rois batch_index+ score <- getBatch scores batch_index+ gt_box <- getBatch gt_boxes batch_index++ -- why sum up the coordinates as score?+ -- because of padding gt are coded as all -1+ gt_score <- addScalar 1 =<< sum_ gt_box (Just [(-1)]) True+ gt_score <- prim T._sign (#data := gt_score .& Nil)++ -- all_rois (N+M, 4)+ -- all_scores (N+M,)+ all_rois <- concat_ 0 [roi, gt_box]+ all_scores <- concat_ 0 [score, gt_score] >>= squeeze (Just [-1])++ ious <- prim T.__contrib_box_iou (#lhs := all_rois+ .& #rhs := gt_box+ .& #format := #corner .& Nil)+ -- iou of the best gt box of each roi+ ious_max <- prim T._max (#data := ious .& #axis := Just [-1] .& Nil)+ -- index of the best gt box of each roi+ ious_argmax <- argmax ious (Just (-1)) False++ class_0 <- zerosLike ious_max+ class_2 <- onesLike ious_max >>= mulScalar 2+ class_3 <- onesLike ious_max >>= mulScalar 3++ ignore_indices <- ltScalar 0 all_scores+ pos_indices <- gtScalar fg_overlap ious_max++ -- mask (mark the class of each roi)+ -- score == -1 ==> ignore (class 0)+ -- iou <= fg_overlap ==> neg sample (class 2)+ -- iou > fg_overlap ==> pos sample (class 3)+ mask <- where_ ignore_indices class_0 class_2+ mask <- where_ pos_indices class_3 mask++ -- shuffle mask and ious_argmax+ rand <- prim T.__random_uniform (#low := 0 .& #high := 1+ .& #shape := [num_proposal + max_num_gt] .& Nil)+ rand <- prim T._slice_like (#data := rand .& #shape_like := ious_max .& Nil)+ index<- prim T._argsort (#data := rand .& Nil)+ mask <- takeI index mask+ ious_argmax <- takeI index ious_argmax++ -- sort in order of pos, neg, ignore+ let max_pos = floor $ fromIntegral num_sample * fg_fraction+ topk <- prim T._topk (#data := mask .& #k := max_pos .& #is_ascend := False .& Nil)+ topk_indices <- takeI topk index+ topk_samples <- takeI topk mask+ topk_matches <- takeI topk ious_argmax++ -- sample the positive class+ pos_class <- onesLike topk_samples+ neg_class <- onesLike topk_samples >>= mulScalar (-1)+ -- class 3 ==> label 1+ -- class 2 ==> label -1+ -- class 0 ==> label 0+ cond <- eqScalar 3 topk_samples+ topk_samples <- where_ cond pos_class topk_samples+ cond <- eqScalar 2 topk_samples+ topk_samples <- where_ cond neg_class topk_samples++ -- sample the negative class+ index <- slice_axis index 0 max_pos Nothing+ mask <- slice_axis mask 0 max_pos Nothing+ ious_argmax <- slice_axis ious_argmax 0 max_pos Nothing+ -- class 2 ==> class 4+ class_4 <- onesLike mask >>= mulScalar 4+ cond <- eqScalar 2 mask+ mask <- where_ cond class_4 mask++ let num_neg = num_sample - max_pos+ bottomk <- prim T._topk (#data := mask .& #k := num_neg .& #is_ascend := False .& Nil)+ bottomk_indices <- takeI bottomk index+ bottomk_samples <- takeI bottomk mask+ bottomk_matches <- takeI bottomk ious_argmax++ -- class 4 ==> label -1+ -- class 3 ==> label 1+ -- class 0 ==> label 0+ cond <- eqScalar 3 bottomk_samples+ pos_class <- onesLike bottomk_samples+ bottomk_samples <- where_ cond pos_class bottomk_samples+ cond <- eqScalar 4 bottomk_samples+ neg_class <- onesLike bottomk_samples >>= mulScalar (-1)+ bottomk_samples <- where_ cond neg_class bottomk_samples++ -- concat+ indices <- concat_ 0 [topk_indices, bottomk_indices]+ samples <- concat_ 0 [topk_samples, bottomk_samples]+ matches <- concat_ 0 [topk_matches, bottomk_matches]++ sampled_rois <- takeI indices all_rois+ [x1, y1, x2, y2] <- splitBySections 4 (-1) True sampled_rois+ rois_areas <- join $ liftM2 mul_ (sub_ x2 x1) (sub_ y2 y1)+ ind <- prim T._argsort (#data := rois_areas .& Nil)+ r <- takeI ind sampled_rois+ s <- takeI ind samples+ m <- takeI ind matches+ return (r, s, m)++ getBatch s i = squeeze (Just [0]) =<<+ slice_axis s 0 i (Just (i + 1))+++bboxTargetGenerator :: Int -> Int -> Int+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> Layer (SymbolHandle, SymbolHandle, SymbolHandle, SymbolHandle)+bboxTargetGenerator batch_size num_fg_classes max_pos samples matches anchors gt_label gt_boxes means stds = do+ -- B: batch_size, N: num_rois, M: num_gt, N_pos: max_pos, C: num_fg_classes+ --+ -- samples: (B, N), value -1 (negative), 0 (ignore), 1 (positive)+ -- matches: (B, N), value range [0, M), the best-matched gt of each roi+ -- anchors: (B, N, 4), anchor boxes, min_x, min_y, max_x, max_y+ -- gt_label: (B, M), value range [0, num_fg_classes), excluding background class+ -- gt_boxes: (B, N, 4), gt boxes, min_x, min_y, max_x, max_y+ --+ -- returns:+ -- cls_targets: (B, N_pos), value [0, num_classes], -1 to be ignored+ -- box_targets: (B, N_pos, C, 4)+ -- box_masks: (B, N_pos, C, 4)+ -- mask_sel: (B, N_pos)+ --+ (fg_cls_targets, cls_targets) <- multiClassEncode gt_label samples matches++ ret <- prim T.__contrib_box_encode (#samples := samples+ .& #matches := matches+ .& #anchors := anchors+ .& #refs := gt_boxes+ .& #means := means+ .& #stds := stds .& Nil)+ [box_targets, box_masks] <- mapM (ret `at`) ([0, 1] :: [Int])++ fg_cls_targets <- expandDims 2 fg_cls_targets+ class_ids_fg <- prim T.__arange (#start := 0 .& #stop := Just (fromIntegral num_fg_classes) .& Nil)+ class_ids_fg <- reshape [1,1,-1] class_ids_fg+ -- (B, N, C), one hot indicator for the best gt class id for each roi of each batch+ target_class_fg_onehot <- eqBroadcast fg_cls_targets class_ids_fg++ masks_sel <- slice_axis box_masks (-1) 0 (Just 1)+ masks_sel <- prim T._argsort (#data := masks_sel .& #axis := Just 1 .& #is_ascend := False .& Nil)+ masks_sel <- reshape [batch_size, -1] masks_sel+ -- mask indices of those positive ones (take at most max_pos items)+ masks_sel <- slice_axis masks_sel 1 0 (Just max_pos)++ (box_targets, box_masks, clsid_ohs) <- fmap unzip3 $ forM [0..batch_size-1] $ \i -> do+ ind <- slice_axis masks_sel 0 i (Just (i+1)) >>= squeeze (Just [0])+ target <- slice_axis box_targets 0 i (Just (i+1)) >>= squeeze (Just [0])+ mask <- slice_axis box_masks 0 i (Just (i+1)) >>= squeeze (Just [0])+ clsid_oh <- slice_axis target_class_fg_onehot 0 i (Just (i+1)) >>= squeeze (Just [0])++ target <- takeI ind target >>= expandDims 0+ mask <- takeI ind mask >>= expandDims 0+ clsid_oh <- takeI ind clsid_oh >>= expandDims 0++ return (target, mask, clsid_oh)++ box_targets <- concat_ 0 (box_targets :: [SymbolHandle]) >>= expandDims 2+ box_masks <- concat_ 0 (box_masks :: [SymbolHandle]) >>= expandDims 2+ -- broadcast the one-hot indicator+ clsid_ohs <- concat_ 0 (clsid_ohs:: [SymbolHandle]) >>= expandDims 3 >>= broadcastAxis [3] [4]++ box_targets <- broadcastAxis [2] [num_fg_classes] box_targets+ box_masks <- mulBroadcast box_masks clsid_ohs+ -- return the index of positive masks because we will calculate box loss only on those items+ return (cls_targets, box_targets, box_masks, masks_sel)+++maskTargetGenerator :: Int -> Int -> Int+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> SymbolHandle+ -> Layer (SymbolHandle, SymbolHandle)+maskTargetGenerator batch_size num_fg_classes mask_size gt_masks rois matches cls_targets = do+ -- rois: (B, N, 4), input proposals+ -- gt_masks: (B, M, H, W), input masks of full image size+ -- matches: (B, N), value [0, M), index to gt_label and gt_box.+ -- cls_targets: (B, N), value [0, num_class), excluding background class.+ --+ -- returns:+ -- mask_targets: (B, N, C, MS, MS), sampled masks.+ -- box_weight: (B, N, C, MS, MS), only foreground class has nonzero weight.++ -- gt_masks (B, M, H, W) -> (B, M, 1, H, W) -> B * (M, 1, H, W)+ gt_masks <- reshape [0, -4, -1, 1, 0, 0] gt_masks+ gt_masks <- splitBySections batch_size 0 True gt_masks++ -- rois (B, N, 4) -> B * (N, 4)+ rois <- splitBySections batch_size 0 True rois++ -- remove all -1 (setting to 0), (B, N) -> B * (N,)+ matches <- prim T._relu (#data := matches .& Nil)+ matches <- splitBySections batch_size 0 True matches++ -- (B, N) -> B * (N,)+ cls_targets <- splitBySections batch_size 0 True cls_targets++ class_ids_fg <- prim T.__arange (#start := 0 .& #stop := Just (fromIntegral num_fg_classes) .& Nil)+ -- (C,) -> (1, C)+ class_ids_fg <- reshape [1, -1] class_ids_fg++ masks <- unique "make" $ mapM (make_target class_ids_fg) $ zip4 rois gt_masks matches cls_targets+ let (mask_targets, mask_weights) = unzip masks++ mask_targets <- stack 0 mask_targets+ mask_weights <- stack 0 mask_weights+ return (mask_targets, mask_weights)++ where+ make_target cids (roi, gt, match, cls_targets) = do+ -- gt: (M, 1, H, W)+ -- padded_rois: (N, 5), along the dim-2, gt_index (1) and rois_box (4)+ match <- reshape [-1, 1] match+ padded_rois <- concat_ (-1) [match, roi]+ -- (N, 1, mask_size, mask_size)+ pooled_mask <- prim T.__contrib_ROIAlign (#data := gt+ .& #rois := padded_rois+ .& #pooled_size := [mask_size, mask_size]+ .& #spatial_scale := 1+ .& #sample_ratio := 2 .& Nil)+ -- (N,) -> (N,1)+ cls_targets <- expandDims 1 cls_targets+ -- (N,1) (1,C) -> (N,C)+ cid_onehot <- eqBroadcast cls_targets cids++ cid_onehot <- reshape [-2, 1, 1] cid_onehot+ -- (N, C, mask_size, mask_size)+ mask_weights <- prim T._broadcast_like+ (#lhs := cid_onehot+ .& #rhs := pooled_mask+ .& #lhs_axes := Just [2, 3]+ .& #rhs_axes := Just [2, 3] .& Nil)+ -- (N, 1, mask_size, mask_size) -> (N, C, mask_size, mask_size)+ mask_targets <- broadcastAxis [1] [num_fg_classes] pooled_mask+ return (mask_targets, mask_weights)+++multiClassEncode gt_label samples matches = do+ -- gt_label: (B, M), value range [0, num_fg_classes), excluding background class+ -- samples: (B, N), value -1 (negative), 0 (ignore), 1 (positive)+ -- matches: (B, N), value range [0, M), the best-matched gt of each roi+ labels <- reshape [0, 1, -1] gt_label+ labels <- prim T._broadcast_like (#lhs := labels .& #rhs := matches .& #lhs_axes := Just [1] .& #rhs_axes := Just [1] .& Nil)+ -- labels: (B,N,M) forall batch, roi, gt. class id+ -- fg_cls_targets: (B,N) forall batch, roi, class id of the best gt+ fg_cls_targets <- pick (#data := labels .& #index := matches .& #axis := Just 2 .& Nil)+ -- shift by 1, reserve 0 for the background class+ cls_targets <- addScalar 1 fg_cls_targets++ ign <- onesLike cls_targets >>= mulScalar (-1)+ bck <- zerosLike cls_targets+ pos <- gtScalar 0.5 samples+ neg <- ltScalar (-0.5) samples+ -- [1, num_fg_classes] for fg, 0 for background, -1 for being ignored+ cls_targets <- where_ pos cls_targets ign+ cls_targets <- where_ neg bck cls_targets++ return (fg_cls_targets, cls_targets)+++multiClassDecodeWithClsId num_classes axis threshold prediction = do+ -- num_classes: number of classes, including the background class+ -- axis: the axis where the class prediction is+ -- threshold: prediction under the threshold will be masked+ -- prediction: (B, N, num_classes), predicated probablities+ -- return:+ -- cls_ids: (B, N, num_classes-1)+ -- pred_fg: (B, N, num_classes-1)+ let num_fg_classes = num_classes - 1+ pred_fg <- slice_axis prediction axis 1 Nothing++ -- make a (B, N, num_fg_classes) of values [0..num_fg_classes-1]+ zero <- zerosLike =<< slice_axis prediction axis 0 (Just 1)+ cls_ids <- reshape [1, 1, num_fg_classes]+ =<< prim T.__arange (#start := 0+ .& #stop := Just (fromIntegral num_fg_classes) .& Nil)+ cls_ids <- addBroadcast zero cls_ids++ mask <- gtScalar threshold pred_fg+ ign1 <- zerosLike pred_fg+ ign2 <- mulScalar (-1) =<< onesLike cls_ids+ pred_fg <- where_ mask pred_fg ign1+ cls_ids <- where_ mask cls_ids ign2++ return (cls_ids, pred_fg)
+ src/MXNet/NN/ModelZoo/Resnet.hs view
@@ -0,0 +1,364 @@+module MXNet.NN.ModelZoo.Resnet where++import Data.Typeable (Typeable)+import RIO+import RIO.List (zip3)+import qualified RIO.NonEmpty as RNE++import MXNet.Base+import MXNet.NN.Layer++data NoKnownExperiment = NoKnownExperiment Int+ deriving (Typeable, Show)+instance Exception NoKnownExperiment++-------------------------------------------------------------------------------+-- ResNet++resnet50Args = (#num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,6,3]+ .& #bottle_neck := True+ .& #workspace := 256 .& Nil)++resnet50 num_classes x = do+ flt <- sequential "features" $ do+ u0 <- getFeature x resnet50Args+ u1 <- getTopFeature u0 resnet50Args+ flatten u1+ named "output" $ fullyConnected (#data := flt .& #num_hidden := num_classes .& Nil)++resnet101Args = (#num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,23,3]+ .& #bottle_neck := True+ .& #workspace := 256+ .& Nil)++resset101 num_classes x = do+ flt <- sequential "features" $ do+ u0 <- getFeature x resnet101Args+ u1 <- getTopFeature u0 resnet101Args+ flatten u1+ named "dense0" $ fullyConnected (#data := flt .& #num_hidden := num_classes .& Nil)++symbol :: Int -> Int -> Int -> Layer SymbolHandle+symbol num_classes num_layers image_size = do+ let args = if image_size <= 28 then args_small_image else args_large_image++ x <- variable "x"+ y <- variable "y"++ flt <- sequential "features" $ do+ u0 <- getFeature x args+ u1 <- getTopFeature u0 args+ flatten u1++ logits <- named "output" $ fullyConnected (#data := flt .& #num_hidden := num_classes .& Nil)+ ret <- named "softmax" $ softmaxoutput (#data := logits .& #label := y .& Nil)+ return ret++ where+ args_common = #workspace := 256 .& Nil+ unit0 = (num_layers - 2) `div` 9+ unit1 = (num_layers - 2) `div` 6+ args_small_image+ | (num_layers - 2) `mod` 9 == 0 && num_layers >= 164 = #num_stages := 3+ .& #filter_list := [64, 64, 128, 256]+ .& #units := [unit0, unit0, unit0]+ .& #bottle_neck := True+ .& args_common+ | (num_layers - 2) `mod` 6 == 0 && num_layers < 164 = #num_stages := 3+ .& #filter_list := [64, 64, 32, 64]+ .& #units := [unit1, unit1, unit1]+ .& #bottle_neck := False+ .& args_common++ args_large_image+ | num_layers == 18 = #num_stages := 4+ .& #filter_list := [64, 64, 128, 256, 512]+ .& #units := [2,2,2,2]+ .& #bottle_neck := False+ .& args_common+ | num_layers == 34 = #num_stages := 4+ .& #filter_list := [64, 64, 128, 256, 512]+ .& #units := [3,4,6,3]+ .& #bottle_neck := False+ .& args_common+ | num_layers == 50 = #num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,6,3]+ .& #bottle_neck := True+ .& args_common+ | num_layers == 101 = #num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,4,23,3]+ .& #bottle_neck := True+ .& args_common+ | num_layers == 152 = #num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,8,36,3]+ .& #bottle_neck := True+ .& args_common+ | num_layers == 200 = #num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,24,36,3]+ .& #bottle_neck := True+ .& args_common+ | num_layers == 269 = #num_stages := 4+ .& #filter_list := [64, 256, 512, 1024, 2048]+ .& #units := [3,30,48,8]+ .& #bottle_neck := True+ .& args_common++eps :: Double+eps = 2e-5++bn_mom :: Float+bn_mom = 0.9++type instance ParameterList "resnet" t =+ '[ '("num_stages" , 'AttrReq Int)+ , '("filter_list", 'AttrReq (NonEmpty Int))+ , '("units" , 'AttrReq (NonEmpty Int))+ , '("bottle_neck", 'AttrReq Bool)+ , '("workspace" , 'AttrReq Int)]++getFeature :: (Fullfilled "resnet" () args)+ => SymbolHandle+ -> ArgsHMap "resnet" () args+ -> Layer SymbolHandle+getFeature inp args = do+ bnx <- batchnorm (#data := inp+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := True .& Nil)++ bdy <- convolution (#data := bnx+ .& #kernel := [7,7]+ .& #num_filter:= filter0+ .& #stride := [2,2]+ .& #pad := [3,3]+ .& #workspace := conv_workspace+ .& #no_bias := True .& Nil)++ bdy <- batchnorm (#data := bdy+ .& #fix_gamma := False+ .& #eps := eps+ .& #momentum := bn_mom .& Nil)++ bdy <- activation (#data := bdy+ .& #act_type := #relu .& Nil)++ bdy <- pooling (#data := bdy+ .& #kernel := [3,3]+ .& #stride := [2,2]+ .& #pad := [1,1]+ .& #pool_type := #max+ .& Nil)++ foldM (buildLayer bottle_neck conv_workspace) bdy (zip3 [0::Int ..2] filter_list units)++ where+ filter0 :| filter_list = args ! #filter_list+ units = RNE.toList $ args ! #units+ bottle_neck = args ! #bottle_neck+ conv_workspace = args ! #workspace++getTopFeature :: (Fullfilled "resnet" () args)+ => SymbolHandle -> ArgsHMap "resnet" () args -> Layer SymbolHandle+getTopFeature inp args = do+ bdy <- buildLayer bottle_neck conv_workspace inp (3, filter, unit)+ bn1 <- batchnorm (#data := bdy -- 9+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ ac1 <- unique' $+ activation (#data := bn1 -- 10+ .& #act_type := #relu .& Nil)+ unique' $ pooling (#data := ac1 -- 11+ .& #kernel := [7,7]+ .& #pool_type := #avg+ .& #global_pool := True .& Nil)+ where+ filter = RNE.last $ args ! #filter_list+ unit = RNE.last $ args ! #units+ bottle_neck = args ! #bottle_neck+ conv_workspace = args ! #workspace++buildLayer :: Bool -> Int -> SymbolHandle -> (Int, Int, Int) -> Layer SymbolHandle+buildLayer bottle_neck workspace bdy (stage_id, filter_size, unit) =+ -- unique (sformat ("stage" % int) (stage_id + 1)) $ do+ subscope_next_name $ sequential' $ do+ bdy <- residual (0,0)+ (#data := bdy+ .& #num_filter := filter_size+ .& #stride := stride0+ .& #dim_match := False+ .& resargs)+ let conv_id = if bottle_neck then 4 else 3+ bn_id = 3+ foldM (\bdy unit_id ->+ residual (conv_id + (unit_id - 1) * 3, bn_id + (unit_id - 1) * 3)+ (#data := bdy+ .& #num_filter := filter_size+ .& #stride := [1,1]+ .& #dim_match := True+ .& resargs)) -- unit_id+ bdy+ ([1..unit-1] :: [Int])+ where+ stride0 = if stage_id == 0 then [1,1] else [2,2]+ -- name unit_id = sformat ("features." % int % "." % int) (stage_id+5) unit_id+ resargs = #bottle_neck := bottle_neck .& #workspace := workspace .& #memonger := False .& Nil++type instance ParameterList "_residual_layer(resnet)" t =+ '[ '("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)+ => (Int, Int)+ -> ArgsHMap "_residual_layer(resnet)" () args+ -> Layer SymbolHandle+residual (conv_id, bn_id) args = subscope_next_name $ 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+ if bottle_neck+ then do+ bn1 <- -- named (sformat ("batchnorm" % int) bn_id) $+ named "bn1" $+ batchnorm (#data := dat+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act1 <- unique' $+ activation (#data := bn1+ .& #act_type := #relu .& Nil)+ conv1<- -- named (sformat ("conv" % int) conv_id) $+ named "conv1" $+ convolution (#data := act1+ .& #kernel := [1,1]+ .& #num_filter := num_filter `div` 4+ .& #stride := [1,1]+ .& #pad := [0,0]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)++ bn2 <- -- named (sformat ("batchnorm" % int) (bn_id + 1)) $+ named "bn2" $+ batchnorm (#data := conv1+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act2 <- unique' $+ activation (#data := bn2+ .& #act_type := #relu .& Nil)+ conv2<- -- named (sformat ("conv" % int) (conv_id + 1)) $+ named "conv2" $+ convolution (#data := act2+ .& #kernel := [3,3]+ .& #num_filter := (num_filter `div` 4)+ .& #stride := stride+ .& #pad := [1,1]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)++ bn3 <- -- named (sformat ("batchnorm" % int) (bn_id + 2)) $+ named "bn3" $+ batchnorm (#data := conv2+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act3 <- unique' $+ activation (#data := bn3+ .& #act_type := #relu .& Nil)+ conv3<- -- named (sformat ("conv" % int) (conv_id + 2)) $+ named "conv3" $+ convolution(#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 -- named (sformat ("conv" % int) (conv_id + 3)) $+ named "downsample" $+ convolution (#data := act1+ .& #kernel := [1,1]+ .& #num_filter := num_filter+ .& #stride := stride+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ when memonger $+ liftIO $ void $ mxSymbolSetAttr shortcut "mirror_stage" "true"++ named "plus" $ add_ conv3 shortcut++ else do+ bn1 <- -- named (sformat ("batchnorm" % int) bn_id) $+ named "bn1" $+ batchnorm (#data := dat+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act1 <- unique' $+ activation (#data := bn1+ .& #act_type := #relu .& Nil)+ conv1<- --named (sformat ("conv" % int) conv_id) $+ named "conv1" $+ convolution (#data := act1+ .& #kernel := [3,3]+ .& #num_filter:= num_filter+ .& #stride := stride+ .& #pad := [1,1]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)++ bn2 <- --named (sformat ("batchnorm" % int) (bn_id + 1)) $+ named "bn2" $+ batchnorm (#data := conv1+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act2 <- unique' $+ activation (#data := bn2+ .& #act_type := #relu .& Nil)+ conv2<- --named (sformat ("conv" % int) (conv_id + 1)) $+ named "conv2" $+ convolution (#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 -- named (sformat ("conv" % int) (conv_id + 2)) $+ named "downsample" $+ convolution (#data := act1+ .& #kernel := [1,1]+ .& #num_filter:= num_filter+ .& #stride := stride+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ when memonger $+ liftIO $ void $ mxSymbolSetAttr shortcut "mirror_stage" "true"++ named "plus" $ add_ conv2 shortcut+
+ src/MXNet/NN/ModelZoo/Resnext.hs view
@@ -0,0 +1,199 @@+module MXNet.NN.ModelZoo.Resnext where++import Formatting+import RIO++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++rootName = "resnext0"++symbol :: SymbolHandle -> Layer SymbolHandle+symbol dat = unique rootName $ do+ bnx <- named "batchnorm0" $+ batchnorm (#data := dat+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := True .& Nil)++ cvx <- named "conv0" $+ convolution (#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, stage_id, unit_id) ->+ unique (sformat ("stage" % int) stage_id) $ residual unit_id+ (#data := layer+ .& #num_filter := num_filter+ .& #stride := stride+ .& #dim_match := dim_match .& resargs))+ cvx+ residual'parms++ pool1 <- pooling (#data := bdy+ .& #kernel := [7,7]+ .& #pool_type := #avg+ .& #global_pool := True .& Nil)+ flat <- flatten pool1+ named "dense0" $ fullyConnected (#data := flat+ .& #num_hidden := 10 .& Nil)+ where+ bn_mom = 0.9 :: Float+ conv_workspace = 256 :: Int+ eps = 2e-5 :: Double+ residual'parms =+ [(64, [1,1], False, 1::Int, 1::Int)] ++ [(64, [1,1], True, 1, i) | i <- [2..18]]+ ++ [(128, [2,2], False, 2, 1)] ++ [(128, [1,1], True, 2, i) | i <- [2..18]]+ ++ [(256, [2,2], False, 3, 1)] ++ [(256, [1,1], True, 3, i) | i <- [2..18]]+ resargs = #bottle_neck := True .& #workspace := conv_workspace .& #memonger := False .& Nil++type instance ParameterList "_residual_layer(resnext)" t =+ '[ '("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)+ => Int -> ArgsHMap "_residual_layer(resnext)" () args -> Layer SymbolHandle+residual _id 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 <- named (sformat ("conv" % int) _id) $+ convolution (#data := dat+ .& #kernel := [1,1]+ .& #num_filter:= num_filter `div` 2+ .& #stride := [1,1]+ .& #pad := [0,0]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ bn1 <- named (sformat ("batchnorm" % int) _id) $+ batchnorm (#data := conv1+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act1 <- activation (#data := bn1+ .& #act_type := #relu .& Nil)+ conv2 <- named (sformat ("conv" % int) (_id + 1)) $+ convolution (#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 <- named (sformat ("batchnorm" % int) (_id + 1)) $+ batchnorm (#data := conv2+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act2 <- activation (#data := bn2+ .& #act_type := #relu .& Nil)+ conv3 <- named (sformat ("conv" % int) (_id + 2)) $+ convolution (#data := act2+ .& #kernel := [1,1]+ .& #num_filter:= num_filter+ .& #stride := [1,1]+ .& #pad := [0,0]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ bn3 <- named (sformat ("batchnorm" % int) (_id + 2)) $+ batchnorm (#data := conv3+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ shortcut <-+ if dim_match+ then return dat+ else do+ shortcut_conv <- named (sformat ("conv" % int) (_id + 3)) $+ convolution (#data := dat+ .& #kernel := [1,1]+ .& #num_filter := num_filter+ .& #stride := stride+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ named (sformat ("conv" % int) (_id + 3)) $+ batchnorm (#data := shortcut_conv+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ when memonger $+ liftIO $ void $ mxSymbolSetAttr shortcut "mirror_stage" "true"+ eltwise <- add_ bn3 shortcut+ activation (#data := eltwise .& #act_type := #relu .& Nil)+ else do+ conv1 <- named (sformat ("conv" % int) _id) $+ convolution (#data := dat+ .& #kernel := [3,3]+ .& #num_filter := num_filter+ .& #stride := stride+ .& #pad := [1,1]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ bn1 <- named (sformat ("batchnorm" % int) _id) $+ batchnorm (#data := conv1+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ act1 <- activation (#data := bn1+ .& #act_type := #relu .& Nil)+ conv2 <- named (sformat ("conv" % int) (_id + 1)) $+ convolution (#data := act1+ .& #kernel := [3,3]+ .& #num_filter := num_filter+ .& #stride := [1,1]+ .& #pad := [1,1]+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ bn2 <- named (sformat ("batchnorm" % int) (_id + 1)) $+ batchnorm (#data := conv2+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ shortcut <-+ if dim_match+ then return dat+ else do+ shortcut_conv <- named (sformat ("conv" % int) (_id + 2)) $+ convolution (#data := act1+ .& #kernel := [1,1]+ .& #num_filter := num_filter+ .& #stride := stride+ .& #workspace := workspace+ .& #no_bias := True .& Nil)+ named (sformat ("batchnorm" % int) (_id + 2)) $+ batchnorm (#data := shortcut_conv+ .& #eps := eps+ .& #momentum := bn_mom+ .& #fix_gamma := False .& Nil)+ when memonger $+ liftIO $ void $ mxSymbolSetAttr shortcut "mirror_stage" "true"+ eltwise <- add_ bn2 shortcut+ activation (#data := eltwise .& #act_type := #relu .& Nil)
+ src/MXNet/NN/ModelZoo/Utils/Box.hs view
@@ -0,0 +1,80 @@+module MXNet.NN.ModelZoo.Utils.Box where++import RIO+import Data.Array.Repa (Array, U, DIM1, Z(..), (:.)(..))+import qualified Data.Array.Repa as Repa++import MXNet.NN.Utils.Repa+++type RBox = Array U DIM1 Float++bboxArea :: RBox -> Float+bboxArea box = (box ^#! 2 - box ^#! 0 + 1) * (box ^#! 3 - box ^#! 1 + 1)++bboxIntersect :: RBox -> RBox -> Maybe RBox+bboxIntersect box1 box2 | not valid = Nothing+ | otherwise = Just $ Repa.fromListUnboxed (Z:.4) [x1, y1, x2, y2]+ where+ valid = x2 - x1 > 0 && y2 - y1 > 0+ x1 = max (box1 ^#! 0) (box2 ^#! 0)+ x2 = min (box1 ^#! 2) (box2 ^#! 2)+ y1 = max (box1 ^#! 1) (box2 ^#! 1)+ y2 = min (box1 ^#! 3) (box2 ^#! 3)++bboxIOU :: RBox -> RBox -> Float+bboxIOU box1 box2 = case bboxIntersect box1 box2 of+ Nothing -> 0+ Just boxI -> let areaI = bboxArea boxI+ areaU = bboxArea box1 + bboxArea box2 - areaI+ in areaI / areaU++whctr :: RBox -> RBox+whctr box1 = Repa.fromListUnboxed (Z:.4) [w, h, x, y]+ where+ [x0, y0, x1, y1] = Repa.toList box1+ w = x1 - x0 + 1+ h = y1 - y0 + 1+ x = x0 + 0.5 * (w - 1)+ y = y0 + 0.5 * (h - 1)++bboxTransform :: RBox -> RBox -> RBox -> RBox+bboxTransform stds box1 box2 =+ let [w1, h1, cx1, cy1] = Repa.toList $ whctr box1+ [w2, h2, cx2, cy2] = Repa.toList $ whctr box2+ dx = (cx2 - cx1) / (w1 + 1e-14)+ dy = (cy2 - cy1) / (h1 + 1e-14)+ dw = log (w2 / w1)+ dh = log (h2 / h1)+ in Repa.computeS $ Repa.fromListUnboxed (Z:.4) [dx, dy, dw, dh] Repa./^ stds++ctrwh :: RBox -> RBox+ctrwh box1 = Repa.fromListUnboxed (Z:.4) [x0, y0, x1, y1]+ where+ [w, h, cx, cy] = Repa.toList box1+ x0 = cx - 0.5 * (w - 1)+ y0 = cy - 0.5 * (h - 1)+ x1 = w + x0 - 1+ y1 = h + y0 - 1++bboxTransInv :: RBox -> RBox -> RBox -> RBox+bboxTransInv stds box delta =+ let [dx, dy, dw, dh] = Repa.toList $ delta Repa.*^ stds+ [w1, h1, cx1, cy1] = Repa.toList $ whctr box+ w2 = exp dw * w1+ h2 = exp dh * w2+ cx2 = dx * w1 + cx1+ cy2 = dy * h1 + cy1+ in ctrwh $ Repa.fromListUnboxed (Z:.4) [w2, h2, cx2, cy2]+++bboxClip :: Float -> Float -> RBox -> RBox+bboxClip height width box = Repa.fromListUnboxed (Z:.4) [x0', y0', x1', y1']+ where+ [x0, y0, x1, y1] = Repa.toList box+ w' = width - 1+ h' = height - 1+ x0' = max 0 (min x0 w')+ y0' = max 0 (min y0 h')+ x1' = max 0 (min x1 w')+ y1' = max 0 (min y1 h')
+ src/MXNet/NN/ModelZoo/VGG.hs view
@@ -0,0 +1,115 @@+{-# LANGUAGE ViewPatterns #-}+module MXNet.NN.ModelZoo.VGG where++import RIO+import RIO.List (scanl, zip3)++import MXNet.Base+import MXNet.NN.Layer++{-+VGG(+ (features): HybridSequential(+ (0): Conv2D(3 -> 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (1): Activation(relu)+ (2): Conv2D(64 -> 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (3): Activation(relu)+ (4): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False, global_pool=False, pool_type=max, layout=NCHW)+ (5): Conv2D(64 -> 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (6): Activation(relu)+ (7): Conv2D(128 -> 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (8): Activation(relu)+ (9): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False, global_pool=False, pool_type=max, layout=NCHW)+ (10): Conv2D(128 -> 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (11): Activation(relu)+ (12): Conv2D(256 -> 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (13): Activation(relu)+ (14): Conv2D(256 -> 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (15): Activation(relu)+ (16): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False, global_pool=False, pool_type=max, layout=NCHW)+ (17): Conv2D(256 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (18): Activation(relu)+ (19): Conv2D(512 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (20): Activation(relu)+ (21): Conv2D(512 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (22): Activation(relu)+ (23): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False, global_pool=False, pool_type=max, layout=NCHW)+ (24): Conv2D(512 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (25): Activation(relu)+ (26): Conv2D(512 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (27): Activation(relu)+ (28): Conv2D(512 -> 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))+ (29): Activation(relu)+ ** (30): MaxPool2D(size=(2, 2), stride=(2, 2), padding=(0, 0), ceil_mode=False, global_pool=False, pool_type=max, layout=NCHW)+ (31): Dense(25088 -> 4096, Activation(relu))+ (32): Dropout(p = 0.5, axes=())+ (33): Dense(4096 -> 4096, Activation(relu))+ (34): Dropout(p = 0.5, axes=())+ )+ (output): Dense(4096 -> 1000, linear)+)++** It appears only if `with_last_pooling` is True.+ -}+++getFeature :: SymbolHandle -> [Int] -> [Int] -> Bool -> Bool -> Layer SymbolHandle+getFeature dat layers filters with_batch_norm with_last_pooling = do+ sym <- foldM build1 dat 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 sym $ zip [idx..] $ replicate num filter+ if not with_last_pooling+ then return sym+ else pooling (#data := sym+ .& #pool_type := #max+ .& #kernel := [2,2]+ .& #stride := [2,2] .& Nil)++ where+ idxes = scanl (+) 0 layers+ last_group:groups = reverse $ zip3 idxes layers filters+ specs = reverse groups++ build1 sym (idx, num, filter) = do+ sym <- foldM build2 sym $ zip [idx..] $ replicate num filter+ pooling (#data := sym+ .& #pool_type := #max+ .& #kernel := [2,2]+ .& #stride := [2,2] .& Nil)++ build2 sym (idx, filter) = do+ sym <- convolution (#data := sym+ .& #kernel := [3,3]+ .& #pad := [1,1]+ .& #num_filter := filter+ .& #workspace := 2048 .& Nil)+ sym <- if with_batch_norm+ then batchnorm (#data := sym .& Nil)+ else return sym+ activation (#data := sym .& #act_type := #relu .& Nil)++getTopFeature :: SymbolHandle -> Layer SymbolHandle+getTopFeature input = do+ sym <- unique' $ flatten input+ sym <- fullyConnected (#data := sym .& #num_hidden := 4096 .& Nil)+ -- sym <- activation (#data := sym .& #act_type := #relu .& Nil)+ sym <- dropout sym 0.5+ sym <- fullyConnected (#data := sym .& #num_hidden := 4096 .& Nil)+ -- sym <- activation (#data := sym .& #act_type := #relu .& Nil)+ dropout sym 0.5++symbol :: SymbolHandle -> Int -> Bool -> Layer SymbolHandle+symbol dat num_layers with_batch_norm =+ getFeature dat layers filters with_batch_norm True >>= getTopFeature+ 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])++vgg16 dat num_classes = do+ sym <- sequential "features" $ symbol dat 16 False+ named "output" $ fullyConnected (#data := sym .& #num_hidden := num_classes .& Nil)