<jupyter_start><jupyter_text>Train the MobileNet-based Mask-RCNN on the Coco DatasetThis notebook provides a fast way to start training on Coco dataset based on MobileNet-based Mask-RCNN.<jupyter_code>"""
Based on the work of <NAME> (Matterport)
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

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

# assign the GPU 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_LOGS_DIR, "MobileNet-v1", "my-training-1", "training20181203T1453", "mask_rcnn_training_0164.h5")
INIT_MODEL_PATH = os.path.join(HOME_DIR, MODEL_PATH_UNDER_HOME)
LOG_DIR = os.path.join(DEFAULT_LOGS_DIR, "MobileNet-v1", "my-training-1")
# training config
class TrainingConfig(coco.CocoConfig):
    NAME = "training"
    
    # 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
    BACKBONE = "mobilenetv1"
    BACKBONE_STRIDES = [4, 8, 16, 32, 64]
    
    # heads
    TOP_DOWN_PYRAMID_SIZE = 256
    DETECTION_HEAD = "original" 
    MASK_HEAD = "original"
#     MASK_HEAD = "mobile"
    FPN_CLASSIF_FC_LAYERS_SIZE = 1024
    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, 'mrcnn_class_loss': 1.0, 'mrcnn_bbox_loss': 1.5, 'mrcnn_mask_loss': 1.0}
    STEPS_PER_EPOCH = 1000
    VALIDATION_STEPS = 50
    TRAIN_BN = False
    TRAIN_ROIS_PER_IMAGE = 200
    ROI_POSITIVE_RATIO = 0.33
    
config = TrainingConfig()
config.display()
model = modellib.MaskRCNN(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)
# # # layers: "heads", "5+", "4+", "3+", "2+", all"
model.train(train_dataset=train_dataset,
                val_dataset=val_dataset,
                learning_rate=config.LEARNING_RATE / 10.0,
                epochs=194,
                layers='2+',
                augmentation=augmentation)
# Refer to the trained ResNet-101-based Mask-RCNN, loss should be 
# Epoch 1/30
#  28/1000 [..............................] - ETA: 1:37:39 
# - loss: 0.6337 
# - rpn_class_loss: 0.0112 
# - rpn_bbox_loss: 0.1044 
# - mrcnn_class_loss: 0.2009 
# - mrcnn_bbox_loss: 0.0904 
# - mrcnn_mask_loss: 0.2268<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12167952013710369, 'fraction_numerical': 0.04655812624964296, 'mean_word_length': 3.8104395604395602, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, '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': '131494', 'license': 'MIT', 'n_tokens_mistral': 1463, 'n_tokens_neox': 1337, 'n_words': 282, 'path': 'notebooks/trainval/MobileNet-v1/train_on_coco.ipynb', 'repo_name': 'SayanSahu99/my-Mobile-Mask-RCNN', 'size': '3509'}