Spaces:
Runtime error
Runtime error
| import detectron2.data.transforms as T | |
| from detectron2.config.lazy import LazyCall as L | |
| from detectron2.layers.batch_norm import NaiveSyncBatchNorm | |
| from detectron2.solver import WarmupParamScheduler | |
| from fvcore.common.param_scheduler import MultiStepParamScheduler | |
| from ..common.data.coco import dataloader | |
| from ..common.models.mask_rcnn_fpn import model | |
| from ..common.optim import SGD as optimizer | |
| from ..common.train import train | |
| # train from scratch | |
| train.init_checkpoint = "" | |
| train.amp.enabled = True | |
| train.ddp.fp16_compression = True | |
| model.backbone.bottom_up.freeze_at = 0 | |
| # SyncBN | |
| # fmt: off | |
| model.backbone.bottom_up.stem.norm = \ | |
| model.backbone.bottom_up.stages.norm = \ | |
| model.backbone.norm = "SyncBN" | |
| # Using NaiveSyncBatchNorm becase heads may have empty input. That is not supported by | |
| # torch.nn.SyncBatchNorm. We can remove this after | |
| # https://github.com/pytorch/pytorch/issues/36530 is fixed. | |
| model.roi_heads.box_head.conv_norm = \ | |
| model.roi_heads.mask_head.conv_norm = lambda c: NaiveSyncBatchNorm(c, | |
| stats_mode="N") | |
| # fmt: on | |
| # 2conv in RPN: | |
| # https://github.com/tensorflow/tpu/blob/b24729de804fdb751b06467d3dce0637fa652060/models/official/detection/modeling/architecture/heads.py#L95-L97 # noqa: E501, B950 | |
| model.proposal_generator.head.conv_dims = [-1, -1] | |
| # 4conv1fc box head | |
| model.roi_heads.box_head.conv_dims = [256, 256, 256, 256] | |
| model.roi_heads.box_head.fc_dims = [1024] | |
| # resize_and_crop_image in: | |
| # https://github.com/tensorflow/tpu/blob/b24729de804fdb751b06467d3dce0637fa652060/models/official/detection/utils/input_utils.py#L127 # noqa: E501, B950 | |
| image_size = 1024 | |
| dataloader.train.mapper.augmentations = [ | |
| L(T.ResizeScale)( | |
| min_scale=0.1, max_scale=2.0, target_height=image_size, target_width=image_size | |
| ), | |
| L(T.FixedSizeCrop)(crop_size=(image_size, image_size)), | |
| L(T.RandomFlip)(horizontal=True), | |
| ] | |
| # recompute boxes due to cropping | |
| dataloader.train.mapper.recompute_boxes = True | |
| # larger batch-size. | |
| dataloader.train.total_batch_size = 64 | |
| # Equivalent to 100 epochs. | |
| # 100 ep = 184375 iters * 64 images/iter / 118000 images/ep | |
| train.max_iter = 184375 | |
| lr_multiplier = L(WarmupParamScheduler)( | |
| scheduler=L(MultiStepParamScheduler)( | |
| values=[1.0, 0.1, 0.01], | |
| milestones=[163889, 177546], | |
| num_updates=train.max_iter, | |
| ), | |
| warmup_length=500 / train.max_iter, | |
| warmup_factor=0.067, | |
| ) | |
| optimizer.lr = 0.1 | |
| optimizer.weight_decay = 4e-5 | |