diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/fei-modelzoo.cabal b/fei-modelzoo.cabal new file mode 100644 --- /dev/null +++ b/fei-modelzoo.cabal @@ -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 diff --git a/src/MXNet/NN/ModelZoo/Lenet.hs b/src/MXNet/NN/ModelZoo/Lenet.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/Lenet.hs @@ -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) diff --git a/src/MXNet/NN/ModelZoo/RCNN/FPN.hs b/src/MXNet/NN/ModelZoo/RCNN/FPN.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/RCNN/FPN.hs @@ -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) diff --git a/src/MXNet/NN/ModelZoo/RCNN/FasterRCNN.hs b/src/MXNet/NN/ModelZoo/RCNN/FasterRCNN.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/RCNN/FasterRCNN.hs @@ -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) + diff --git a/src/MXNet/NN/ModelZoo/RCNN/MaskRCNN.hs b/src/MXNet/NN/ModelZoo/RCNN/MaskRCNN.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/RCNN/MaskRCNN.hs @@ -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) diff --git a/src/MXNet/NN/ModelZoo/RCNN/RCNN.hs b/src/MXNet/NN/ModelZoo/RCNN/RCNN.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/RCNN/RCNN.hs @@ -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) diff --git a/src/MXNet/NN/ModelZoo/Resnet.hs b/src/MXNet/NN/ModelZoo/Resnet.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/Resnet.hs @@ -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 + diff --git a/src/MXNet/NN/ModelZoo/Resnext.hs b/src/MXNet/NN/ModelZoo/Resnext.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/Resnext.hs @@ -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) diff --git a/src/MXNet/NN/ModelZoo/Utils/Box.hs b/src/MXNet/NN/ModelZoo/Utils/Box.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/Utils/Box.hs @@ -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') diff --git a/src/MXNet/NN/ModelZoo/VGG.hs b/src/MXNet/NN/ModelZoo/VGG.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/ModelZoo/VGG.hs @@ -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)