Mask Rcnn Scriptmode Fsx
Distributed Training of Mask-RCNN on Amazon SageMaker using FSx
This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.
This notebook is a step-by-step tutorial on distributed training of Mask R-CNN implemented in TensorFlow framework.
Concretely, we will describe the steps for training TensorPack Faster-RCNN/Mask-RCNN and AWS Samples Mask R-CNN on Amazon SageMaker using Amazon S3 and Amazon FSx for Lustre file-system as data sources.
The outline of steps is as follows:
- Stage COCO 2017 dataset on Amazon S3
- Create Amazon FSx Lustre file-system and import data into the file-system from S3
- Build Docker training image and push it to Amazon ECR
- Configure data input channels
- Configure hyper-prarameters
- Define training metrics
- Define training job and start training
Initialize SageMaker Session
First, let us specify the s3_bucket that we will use throughout the notebook. The s3_bucket must be located in the region of this notebook instance. If you do not specify S3 bucket name in s3_bucket, default SageMaker bucket is used, if it exists. We also initialize the SageMaker session.
Check for Attached FSx for Lustre File-system
We check to see if an FSx for Lustre is attached, and if it is, we use the attached FSx for Lustre file-system for data input, otherwise, we use Amazon S3.
Note: If you created this notebook instance using the stack-sm.sh script, an FSx for Lustre file-system is automatically created and attached to this notebook.
Stage COCO 2017 dataset on Amazon S3
We use COCO 2017 dataset. This step downloads COCO 2017 training and validation dataset to this notebook instance, extracts the files from the dataset, and uploads the extracted files to your Amazon S3 bucket. Expected time to execute this step is 30 minutes.
Stage COCO 2017 dataset on Amazon FSx for Lustre
Next, we stage COCO 2017 dataset on Amazon FSx for Lustre file-system, if such a file-system is attached. The prepare-efs.sh script executes this step. The expected time to execute this step is 30 minutes, if the data isn't already imported from Amazon S3 bucket.
Specify Model Type
We have a choice of two different models:
-
TensorPack Faster-RCNN/Mask-RCNN implementation supports a maximum per-GPU batch size of 1.
-
AWS Samples Mask R-CNN is an optimized implementation that supports a maximum per GPU batch size of 4, assuming per GPU memory of 32 GB.
Below, set the model_type to "aws-samples-mask-rcnn", or "tensorpack-mask-rcnn".
Build and push SageMaker Training Image to ECR
Next, we build and push the training image to Amazon ECR, based on the selected model type. This may take several minutes on first-time build on this notebook. We also set the training_script based on the selected model type.
Note: For this step, the IAM Role attached to this notebook instance needs full access to Amazon ECR service. If you created this notebook instance using the stack-sm.sh script, the IAM Role attached to this notebook instance is already setup with full access to ECR service.
Define SageMaker Data Channels
We define train data channels for Amazon S3, and Amazon FSx, if FSx Lustre file-system is available.
For the training job, S3 data channel is used only if the FSx Lustre file-system is not available.
Define S3 Train Data Channel
We first define S3 train data channel below.
Define Amazon FSx Lustre Train Data Channel
Next, we define the train data channel using FSx Lustre file-system, if Amazon FSx Lustre file-system is available.
Define Model Output Location
Next, we define the model output location in S3 bucket.
Define Security Group and Subnets
If an EFS file-system is attached to this notebook, we retrieve the security groups and subnets associated with the EFS file-system mount-targets, and use them in defining the training job.
Note: For this step, the IAM Role attached to this notebook instance needs permission to describe EFS mount targets, and mount target security groups. If you created this notebook instance using the stack-sm.sh script, the IAM Role attached to this notebook instance is already setup with required permissions.
Configure Hyper-parameters
Next, we define the hyper-parameters.
Note, some hyper-parameters are different between the two implementations. The batch size per GPU in TensorPack Faster-RCNN/Mask-RCNN is fixed at 1, but is configurable in AWS Samples Mask-RCNN. The learning rate schedule is specified in units of steps in TensorPack Faster-RCNN/Mask-RCNN, but in epochs in AWS Samples Mask-RCNN.
The default learning rate schedule values shown below correspond to training for a total of 24 epochs, at 120,000 images per epoch.
TensorPack Faster-RCNN/Mask-RCNN Hyper-parameters
| Hyper-parameter | Description | Default |
|---|---|---|
| backbone_weights | ResNet backbone pre-trained weights file | 'ImageNet-R50-AlignPadding.npz' |
| batch_norm | Batch normalization option ('FreezeBN', 'SyncBN', 'GN', 'None') | 'FreezeBN' |
| config: | Any hyper-parameter prefixed with config: is set as a model config parameter | - |
| data_train | Training data | 'coco_train2017' |
| data_val | Validation data | 'coco_val2017' |
| eval_period | Number of epochs period for evaluation during training | 1 |
| images_per_epoch | Images per epoch | 120000 |
| load_model | Pre-trained model to load | - |
| lr_schedule | Learning rate schedule in training steps | '[240000, 320000, 360000]' |
| mode_fpn | Use Feature Pyramid Network (FPN) mode | True |
| mode_mask | Compute masks | True |
| resnet_arch | Must be 'resnet50' or 'resnet101' | 'resnet50' |
AWS Samples Mask-RCNN Hyper-parameters
| Hyper-parameter | Description | Default |
|---|---|---|
| backbone_weights | ResNet backbone pre-trained weights file | 'ImageNet-R50-AlignPadding.npz' |
| batch_norm | Batch normalization option ('FreezeBN', 'SyncBN', 'GN', 'None') | 'FreezeBN' |
| batch_size_per_gpu | Batch size per gpu, 1 - 6 | 16 GB: 2, 32 GB: 4, > 32 GB : 6 |
| config: | Any hyper-parameter prefixed with config: is set as a model config parameter | - |
| data_train | Training data | 'train2017' |
| data_val | Validation data | 'val2017' |
| eval_period | Number of epochs period for evaluation during training | 1 |
| lr_schedule | Learning rate schedule in training steps | '[(16, 0.1), (20, 0.01), (24, None)]' |
| images_per_epoch | Images per epoch | 120000 |
| load_model | Pre-trained model to load | - |
| mode_fpn | Use Feature Pyramid Network (FPN) mode. Must be True. | True |
| mode_mask | Compute masks | True |
| resnet_arch | Must be 'resnet50' or 'resnet101' | 'resnet50' |
Define Training Metrics
Next, we define the regular expressions that SageMaker uses to extract algorithm metrics from training logs and send them to AWS CloudWatch metrics. These algorithm metrics are visualized in SageMaker console.
Define SageMaker Training Job
Next, we use SageMaker Estimator API to define a SageMaker Training Job.
Define training job
We recommned using 16 GPUs for the training job, so we set instance_count=2. We recommend using 100 GB Amazon EBS storage volume with each training instance, so we set volume_size = 100.
Launch Training Job
Finally, we launch the SageMaker training job. See Training Jobs in SageMaker console to monitor the training job.
Notebook CI Test Results
This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.