<jupyter_start><jupyter_text>Mask-RCNN: mimic training on Coco Dataset<jupyter_code>"""
written by wozhouh
"""

# Import Python Packages
import os
import imgaug

# Root directory of the project
ROOT_DIR = os.path.abspath("../../../")
import sys

# Import Mask RCNN
sys.path.append(ROOT_DIR)
from mrcnn import model as modellib
from optimization import mimicking

# Import COCO config
sys.path.append(os.path.join(ROOT_DIR, "samples", "coco"))  # To find local version
import coco

# assign the GPUs for training
os.environ['CUDA_VISIBLE_DEVICES'] = '5,6,7'
# Default setting
HOME_DIR = os.getenv('HOME')
DEFAULT_WEIGHTS_DIR = os.path.join(ROOT_DIR, "weights")
DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")
DEFAULT_DATASET_YEAR = "2017"
DEFAULT_COCO_PATH = os.path.join(HOME_DIR, "data", "Coco")

# weights to load
MODEL_PATH_UNDER_HOME = os.path.join(DEFAULT_WEIGHTS_DIR, "ResNet-101", "mask_rcnn_coco.h5")
INIT_MODEL_PATH = os.path.join(HOME_DIR, MODEL_PATH_UNDER_HOME)
LOG_DIR = os.path.join(DEFAULT_LOGS_DIR, "Mimicking")
# Mimicking config
class MimicConfig(coco.CocoConfig):
    NAME = "Mimicking"
    
    # GPU
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 2
    IMAGES_PER_GPU = 2
    
    # data
    IMAGE_MIN_DIM = 800
    IMAGE_MAX_DIM = 1024
    RPN_ANCHOR_SCALES = (32, 64, 128, 256, 512)
    
    # model
    TEACHER_BACKBONE = "resnet101"
    STUDENT_BACKBONE = "resnet50"
    BACKBONE_STRIDES = [4, 8, 16, 32, 64]
    
    # heads
    TOP_DOWN_PYRAMID_SIZE = 256
    RPN_TRAIN_ANCHORS_PER_IMAGE = 256
    
    # training
    LEARNING_RATE = 0.001
    WEIGHT_DECAY = 0.0001
    LOSS_WEIGHTS = {'rpn_class_loss': 1.0, 'rpn_bbox_loss': 1.0, 'rpn_mimic_loss': 0.1}
    STEPS_PER_EPOCH = 1000
    VALIDATION_STEPS = 50
    
config = MimicConfig()
config.display()
model = mimicking.MimicMaskRCNN(mode="training", model_dir=LOG_DIR, config=config)
# print the model summary
model.keras_model.summary()
# Load the weights
model.load_weights(INIT_MODEL_PATH, by_name=True)
# Load the dataset
train_dataset = coco.CocoDataset()
train_dataset.load_coco(DEFAULT_COCO_PATH, "train", year=DEFAULT_DATASET_YEAR)
train_dataset.prepare()
val_dataset = coco.CocoDataset()
val_dataset.load_coco(DEFAULT_COCO_PATH, "val", year=DEFAULT_DATASET_YEAR)
val_dataset.prepare()

# Image Augmentation
# Right/Left flip 50% of the time
augmentation = imgaug.augmenters.Fliplr(0.5)
model.train(train_dataset=train_dataset,
                val_dataset=val_dataset,
                learning_rate=config.LEARNING_RATE,
                epochs=1,
                layers='all',
                augmentation=augmentation)<jupyter_output>Starting at epoch 0. LR=0.001

Checkpoint Path: /home/processyuan/code/HumanMask/my-Mobile-Mask-RCNN/logs/Mimicking/mimicking20181211T1543/mask_rcnn_mimicking_{epoch:04d}.h5
Selecting layers to train
s_conv1                (Conv2D)
s_bn_conv1             (BatchNorm)
s_res2a_branch2a       (Conv2D)
s_bn2a_branch2a        (BatchNorm)
s_res2a_branch2b       (Conv2D)
s_bn2a_branch2b        (BatchNorm)
s_res2a_branch2c       (Conv2D)
s_res2a_branch1        (Conv2D)
s_bn2a_branch2c        (BatchNorm)
s_bn2a_branch1         (BatchNorm)
s_res2b_branch2a       (Conv2D)
s_bn2b_branch2a        (BatchNorm)
s_res2b_branch2b       (Conv2D)
s_bn2b_branch2b        (BatchNorm)
s_res2b_branch2c       (Conv2D)
s_bn2b_branch2c        (BatchNorm)
s_res2c_branch2a       (Conv2D)
s_bn2c_branch2a        (BatchNorm)
s_res2c_branch2b       (Conv2D)
s_bn2c_branch2b        (BatchNorm)
s_res2c_branch2c       (Conv2D)
s_bn2c_branch2c        (BatchNorm)
s_res3a_branch2a       (Conv2D)
s_bn3a_branch2a        (BatchN[...]
{'fraction_non_alphanumeric': 0.09586466165413533, 'fraction_numerical': 0.04215896885069818, 'mean_word_length': 3.7634271099744248, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '511848', 'license': 'MIT', 'n_tokens_mistral': 1547, 'n_tokens_neox': 1461, 'n_words': 268, 'path': 'notebooks/trainval/Mimicking/train_on_coco.ipynb', 'repo_name': 'SayanSahu99/my-Mobile-Mask-RCNN', 'size': '3724'}