MMDetectionでWeights & Biasesを使用する方法

# Install PyTorch.
!pip install -qq torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

# Install mmcv-full.
!pip install -qq mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html

# Install mmdetection
!git clone https://github.com/open-mmlab/mmdetection
%cd mmdetection
!pip install -e .

!pip install -q --upgrade wandb

import wandb
wandb.login()

!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
!unzip -q balloon_dataset.zip

!mkdir checkpoints
!wget -c https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth \
      -O checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth

config_file = 'mmdetection/configs/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'
cfg = Config.fromfile(config_file)


# Define type and path to the images.
cfg.dataset_type = 'COCODataset'

cfg.data.test.ann_file = 'balloon/val/annotation_coco.json'
cfg.data.test.img_prefix = 'balloon/val/'
cfg.data.test.classes = ('balloon',)

cfg.data.train.ann_file = 'balloon/train/annotation_coco.json'
cfg.data.train.img_prefix = 'balloon/train/'
cfg.data.train.classes = ('balloon',)

cfg.data.val.ann_file = 'balloon/val/annotation_coco.json'
cfg.data.val.img_prefix = 'balloon/val/'
cfg.data.val.classes = ('balloon',)

# modify num classes of the model in box head and mask head
cfg.model.roi_head.bbox_head.num_classes = 1
cfg.model.roi_head.mask_head.num_classes = 1

# Use the pretrained model.
cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'

# The original learning rate (LR) is set for 8-GPU training.
# We divide it by 8 since we only use one GPU.
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# Epochs
cfg.runner.max_epochs = 12

# Set seed thus the results are more reproducible
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# ⭐️ Set the checkpoint interval.
cfg.checkpoint_config.interval = 4

# Set up working dir to save files and logs.
cfg.work_dir = './tutorial_exps'
cfg.device = 'cuda'

# ⭐️ Set the evaluation interval.
cfg.evaluation.interval = 2

cfg.log_config.hooks = [
    dict(type='TextLoggerHook'),
    dict(type='MMDetWandbHook',
         init_kwargs={'project': 'MMDetection-tutorial'},
         interval=10,
         log_checkpoint=True,
         log_checkpoint_metadata=True,
         num_eval_images=100)]

meta = dict()
meta['exp_name'] = os.path.basename(config_file)

def train(cfg, meta):
    # Build dataset
    datasets = [mmdet.datasets.build_dataset(cfg.data.train)]

    # Build the detector
    model = mmdet.models.build_detector(
        cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
    model.CLASSES = datasets[0].CLASSES

    # Dump config for logging
    mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))
    cfg.dump(os.patg.join(cfg.work_dir, meta['exp_name']))

    # Train the model.
    mmdet.apis.train_detector(model, datasets, cfg, distributed=False, validate=True, meta=meta)

train(cfg, meta)