{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!rm -rf detectron2/\n!python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datetime import datetime\nimport logging\n\nfrom detectron2.checkpoint import DetectionCheckpointer\nfrom detectron2.config import LazyConfig, instantiate\nfrom detectron2.engine import (\n    AMPTrainer,\n    SimpleTrainer,\n    default_argument_parser,\n    default_setup,\n    default_writers,\n    hooks,\n    launch\n)\nfrom detectron2.engine.defaults import _try_get_key\nfrom detectron2.engine.defaults import create_ddp_model\nfrom detectron2.evaluation import inference_on_dataset, print_csv_format\nfrom detectron2.utils import comm\n\nfrom functools import partial\nfrom detectron2.utils.file_io import PathManager\nfrom omegaconf import OmegaConf\nimport torch.nn as nn\nfrom fvcore.common.param_scheduler import MultiStepParamScheduler\n\nfrom detectron2 import model_zoo\nfrom detectron2.config import LazyCall as L\nfrom detectron2.solver import WarmupParamScheduler\nfrom detectron2.modeling import MViT\nfrom detectron2.layers import ShapeSpec\nfrom detectron2.modeling.box_regression import Box2BoxTransform\nfrom detectron2.modeling.matcher import Matcher\nfrom detectron2.modeling.roi_heads import (\n    FastRCNNOutputLayers,\n    FastRCNNConvFCHead,\n    CascadeROIHeads,\n)\nfrom detectron2.utils.env import seed_all_rng\nfrom detectron2.data.datasets import register_coco_instances\nimport detectron2.data.transforms as T\nfrom detectron2 import model_zoo\nfrom detectron2.config import LazyCall as L\nimport detectron2.data.transforms as T\nfrom detectron2.config import LazyCall as L\nfrom detectron2.data import (\n    DatasetMapper,\n    build_detection_test_loader,\n    build_detection_train_loader,\n    get_detection_dataset_dicts,\n)\nfrom detectron2.evaluation import COCOEvaluator\nfrom detectron2.utils.logger import setup_logger\nfrom detectron2.utils.collect_env import collect_env_info","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def setup():\n    model = model_zoo.get_config(\"common/models/mask_rcnn_fpn.py\").model\n    constants = model_zoo.get_config(\"common/data/constants.py\").constants\n    model.pixel_mean = constants.imagenet_rgb256_mean\n    model.pixel_std = constants.imagenet_rgb256_std\n    model.input_format = \"RGB\"\n    model.backbone.bottom_up = L(MViT)(\n        embed_dim=96,\n        depth=24,\n        num_heads=1,\n        last_block_indexes=(1, 4, 20, 23),\n        residual_pooling=True,\n        drop_path_rate=0.4,\n        norm_layer=partial(nn.LayerNorm, eps=1e-6),\n        out_features=(\"scale2\", \"scale3\", \"scale4\", \"scale5\"),\n    )\n    model.backbone.in_features = \"${.bottom_up.out_features}\"\n    model.backbone.square_pad = 1024\n\n    # New heads and LN\n    model.backbone.norm = \"LN\"  # Use LN in FPN\n    model.roi_heads.box_head.conv_norm = model.roi_heads.mask_head.conv_norm = \"LN\"\n\n    # 2conv in RPN:\n    model.proposal_generator.head.conv_dims = [-1, -1]\n\n    # arguments that don't exist for Cascade R-CNN\n    [model.roi_heads.pop(k) for k in [\"box_head\", \"box_predictor\", \"proposal_matcher\"]]\n    model.roi_heads.update(\n        _target_=CascadeROIHeads,\n        box_heads=[\n            L(FastRCNNConvFCHead)(\n                input_shape=ShapeSpec(channels=256, height=7, width=7),\n                conv_dims=[256, 256, 256, 256],\n                fc_dims=[1024],\n                conv_norm=\"LN\",\n            )\n            for _ in range(3)\n        ],\n        box_predictors=[\n            L(FastRCNNOutputLayers)(\n                input_shape=ShapeSpec(channels=1024),\n                test_score_thresh=0.05,\n                box2box_transform=L(Box2BoxTransform)(weights=(w1, w1, w2, w2)),\n                cls_agnostic_bbox_reg=True,\n                num_classes=\"${...num_classes}\",\n            )\n            for (w1, w2) in [(10, 5), (20, 10), (30, 15)]\n        ],\n        proposal_matchers=[\n            L(Matcher)(thresholds=[th], labels=[0, 1], allow_low_quality_matches=False)\n            for th in [0.5, 0.6, 0.7]\n        ],\n    )\n    model.roi_heads.num_classes = 4\n    model.roi_heads.batch_size_per_image = 256\n\n    register_coco_instances(\"badlad_train\", {}, '../dataset/badlad/labels/coco_format/train/badlad-train-coco.json', '../dataset/badlad/images/train')\n    register_coco_instances(\"badlad_test\", {}, '../dataset/badlad/badlad-test-metadata.json', '../dataset/badlad/images/test')\n\n    dataloader = OmegaConf.create()\n\n    image_size = 1024\n    dataloader.train = L(build_detection_train_loader)(\n        dataset=L(get_detection_dataset_dicts)(names=\"badlad_train\"),\n        mapper=L(DatasetMapper)(\n            is_train=True,\n            augmentations=[\n                L(T.RandomBrightness)(intensity_min=0.8,intensity_max=1.2),\n                L(T.RandomContrast)(intensity_min=0.5,intensity_max=1.5),\n                L(T.RandomSaturation)(intensity_min=0.5,intensity_max=1.5),\n                L(T.RandomFlip)(prob=0.5,horizontal=True, vertical=False), \n                L(T.RandomFlip)(prob=0.5,horizontal=False, vertical=True),\n                L(T.RandomRotation)(angle=[0,90,-90,180], sample_style=\"choice\"),\n                L(T.RandomRotation)(angle=[-30, 30], sample_style=\"range\"),\n                L(T.ResizeScale)(\n                    min_scale=0.1, max_scale=2.0, target_width=image_size, target_height=image_size\n                ),\n                L(T.FixedSizeCrop)(crop_size=(image_size, image_size)),\n            ],\n            image_format=\"RGB\",\n            use_instance_mask=True,\n        ),\n        total_batch_size=16,\n        num_workers=4,\n    )\n\n    dataloader.test = L(build_detection_test_loader)(\n        dataset=L(get_detection_dataset_dicts)(names=\"badlad_test\", filter_empty=False),\n        mapper=L(DatasetMapper)(\n            is_train=False,\n            augmentations=[\n                L(T.ResizeShortestEdge)(short_edge_length=800, max_size=1333),\n            ],\n            image_format=\"RGB\",\n        ),\n        num_workers=4,\n    )\n\n    dataloader.evaluator = L(COCOEvaluator)(\n        dataset_name=\"${..test.dataset.names}\",\n    )\n\n    dataloader.train.num_workers = 2\n    dataloader.train.total_batch_size = 8\n    # recompute boxes due to cropping\n    dataloader.train.mapper.recompute_boxes = True\n\n    # Initialization and trainer settings\n    train = model_zoo.get_config(\"common/train.py\").train\n    train.amp.enabled = True\n    train.ddp.fp16_compression = True\n    train.init_checkpoint = \"./output/mvit2b/model_0004999.pth\"\n    train.output_dir = \"./output/mvit2b\"\n    # Schedule\n    # 100 ep = 184375 iters * 64 images/iter / 118000 images/ep\n    train.max_iter = 10000\n    train.eval_period = 500\n    train.log_period = 20\n    train.checkpointer.period = 1000\n    train.device = \"cuda\"\n\n    lr_multiplier = L(WarmupParamScheduler)(\n        scheduler=L(MultiStepParamScheduler)(\n            values=[1.0, 0.1, 0.01],\n            milestones=[9000, 9500],\n            num_updates=train.max_iter,\n        ),\n        warmup_length=50 / train.max_iter,\n        warmup_factor=0.001,\n    )\n\n    optimizer = model_zoo.get_config(\"common/optim.py\").AdamW\n    optimizer.params.overrides = {\"pos_embed\": {\"weight_decay\": 0.0}}\n    optimizer.lr = 0.0001\n\n    dataloader.evaluator.output_dir = train.output_dir\n\n    cfg = OmegaConf.create()\n    cfg.model = model\n    cfg.dataloader = dataloader\n    cfg.train = train\n    cfg.optimizer = optimizer\n    cfg.lr_multiplier = lr_multiplier\n\n    return cfg","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def do_test(cfg, model):\n    if \"evaluator\" in cfg.dataloader:\n        ret = inference_on_dataset(\n            model, instantiate(cfg.dataloader.test), instantiate(cfg.dataloader.evaluator)\n        )\n        print_csv_format(ret)\n        return ret","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def do_train(resume, cfg):\n    model = instantiate(cfg.model)\n    logger = logging.getLogger(\"detectron2\")\n    logger.info(\"Model:\\n{}\".format(model))\n    model.to(cfg.train.device)\n\n    cfg.optimizer.params.model = model\n    optim = instantiate(cfg.optimizer)\n\n    train_loader = instantiate(cfg.dataloader.train)\n\n    model = create_ddp_model(model, **cfg.train.ddp)\n    trainer = (AMPTrainer if cfg.train.amp.enabled else SimpleTrainer)(model, train_loader, optim)\n    checkpointer = DetectionCheckpointer(\n        model,\n        cfg.train.output_dir,\n        trainer=trainer,\n    )\n    trainer.register_hooks(\n        [\n            hooks.IterationTimer(),\n            hooks.LRScheduler(scheduler=instantiate(cfg.lr_multiplier)),\n            hooks.PeriodicCheckpointer(checkpointer, **cfg.train.checkpointer)\n            if comm.is_main_process()\n            else None,\n            # hooks.EvalHook(cfg.train.eval_period, lambda: do_test(cfg, model)),\n            hooks.PeriodicWriter(\n                default_writers(cfg.train.output_dir, cfg.train.max_iter),\n                period=cfg.train.log_period,\n            )\n            if comm.is_main_process()\n            else None,\n        ]\n    )\n\n    checkpointer.resume_or_load(cfg.train.init_checkpoint, resume=resume)\n    if resume and checkpointer.has_checkpoint():\n        start_iter = trainer.iter + 1\n    else:\n        start_iter = 0\n    trainer.train(start_iter, cfg.train.max_iter)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main():\n    cfg=setup()\n    output_dir = cfg.train.output_dir\n    if comm.is_main_process() and output_dir:\n        PathManager.mkdirs(output_dir)\n\n    rank = comm.get_rank()\n    setup_logger(output_dir, distributed_rank=rank, name=\"fvcore\")\n    logger = setup_logger(output_dir, distributed_rank=rank)\n\n    logger.info(\"Rank of current process: {}. World size: {}\".format(rank, comm.get_world_size()))\n    logger.info(\"Environment info:\\n\" + collect_env_info())\n\n\n    # make sure each worker has a different, yet deterministic seed if specified\n    cfg.train.seed = int(datetime.now().timestamp()) + rank\n\n    do_train(False, cfg)\n\nif __name__ == \"__main__\":\n    launch(\n        main,\n        4,\n        num_machines=1,\n        args=(),\n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}