Update readme with ShapeMask instructions.

PiperOrigin-RevId: 317896878

official/README.md

@@ -43,6 +43,7 @@ In the near future, we will add:
 |-------|-------------------|
 | [MNIST](vision/image_classification) | A basic model to classify digits from the [MNIST dataset](http://yann.lecun.com/exdb/mnist/) |
 | [ResNet](vision/image_classification) | [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385) |
+| [EfficientNet](vision/image_classification) | [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946) |
 
 #### Object Detection and Segmentation
 
@@ -50,6 +51,7 @@ In the near future, we will add:
 |-------|-------------------|
 | [RetinaNet](vision/detection) | [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002) |
 | [Mask R-CNN](vision/detection) | [Mask R-CNN](https://arxiv.org/abs/1703.06870) |
+| [ShapeMask](vision/detection) | [ShapeMask: Learning to Segment Novel Objects by Refining Shape Priors](https://arxiv.org/abs/1904.03239) |
 
 ### Natural Language Processing
 

official/vision/detection/README.md

@@ -123,8 +123,6 @@ predict:
  predict_batch_size: 8
 architecture:
  use_bfloat16: False
-retinanet_parser:
- use_bfloat16: False
 train:
  total_steps: 1
  batch_size: 8
@@ -245,8 +243,6 @@ predict:
  predict_batch_size: 8
 architecture:
  use_bfloat16: False
-maskrcnn_parser:
- use_bfloat16: False
 train:
  total_steps: 1000
  batch_size: 8
@@ -255,6 +251,140 @@ use_tpu: False
 "
 ```
 
+## Train ShapeMask on TPU
+
+### Train a ResNet-50 based ShapeMask.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+RESNET_CHECKPOINT="<path to the pre-trained Resnet-50 checkpoint>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+SHAPE_PRIOR_PATH="<path to shape priors>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --params_override="{train: { checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN} } shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}} }"
+```
+
+The pre-trained ResNet-50 checkpoint can be downloaded [here](https://storage.cloud.google.com/cloud-tpu-checkpoints/model-garden-vision/detection/resnet50-2018-02-07.tar.gz).
+
+The shape priors can be downloaded [here]
+(https://storage.googleapis.com/cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy)
+
+
+### Train a custom ShapeMask using the config file.
+
+First, create a YAML config file, e.g. *my_shapemask.yaml*.
+This file specifies the parameters to be overridden:
+
+```YAML
+# my_shapemask.yaml
+train:
+  train_file_pattern: <path to the TFRecord training data>
+  total_steps: <total steps to train>
+  batch_size: <training batch size>
+eval:
+  eval_file_pattern: <path to the TFRecord validation data>
+  val_json_file: <path to the validation annotation JSON file>
+  batch_size: <evaluation batch size>
+shapemask_head:
+  shape_prior_path: <path to shape priors>
+```
+
+Once the YAML config file is created, you can launch the training using the
+following command.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+## Train ShapeMask on GPU
+
+Training on GPU is similar to that on TPU. The major change is the strategy type
+(use
+"[mirrored](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy)"
+for multiple GPU and
+"[one_device](https://www.tensorflow.org/api_docs/python/tf/distribute/OneDeviceStrategy)"
+for single GPU).
+
+Multi-GPUs example (assuming there are 8GPU connected to the host):
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=mirrored \
+  --num_gpus=8 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+A single GPU example
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+
+An example with inline configuration (YAML or JSON format):
+
+```
+python3 ~/models/official/vision/detection/main.py \
+  --model_dir=<model folder> \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --mode=train \
+  --model=shapemask \
+  --params_override="eval:
+ eval_file_pattern: <Eval TFRecord file pattern>
+ batch_size: 8
+ val_json_file: <COCO format groundtruth JSON file>
+train:
+ total_steps: 1000
+ batch_size: 8
+ train_file_pattern: <Eval TFRecord file pattern>
+use_tpu: False
+"
+```
+
+
+### Run the evaluation (after training)
+
+```
+python3 /usr/share/models/official/vision/detection/main.py \
+   --strategy_type=tpu \
+   --tpu=${TPU_NAME} \
+   --model_dir=${MODEL_DIR} \
+   --mode=eval \
+   --model=shapemask \
+   --params_override="{eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 } }"
+```
+
+`MODEL_DIR` needs to point to the trained path of ShapeMask model.
+Change `strategy_type=mirrored` and `num_gpus=1` to run on a GPU.
+
 Note: The JSON groundtruth file is useful for [COCO dataset](http://cocodataset.org/#home) and can be
 downloaded from the [COCO website](http://cocodataset.org/#download). For custom dataset, it is unncessary because the groundtruth can be included in the TFRecord files.