coco-centernet-hourglass-tpu.yaml

# COCO AP 40.6% for float16 precision is achieved with the configuration below.
# Expected COCO AP for float32 from OD API is 41.92 +/- 0.16.
runtime:
  distribution_strategy: 'tpu'
  mixed_precision_dtype: 'bfloat16'
task:
  model:
    num_classes: 90
    max_num_instances: 128
    input_size: [512, 512, 3]
    backbone:
      type: hourglass
      hourglass:
        model_id: 52
        num_hourglasses: 2
    head:
      heatmap_bias: -2.19
      input_levels: ['2_0', '2']
    detection_generator:
      max_detections: 100
      peak_error: 0.000001
      peak_extract_kernel_size: 3
      use_nms: false
      nms_pre_thresh: 0.1
      nms_thresh: 0.4
      class_offset: 1
    norm_activation:
      norm_epsilon: 0.00001
      norm_momentum: 0.1
      use_sync_bn: true
  losses:
    detection:
      offset_weight: 1.0
      scale_weight: 0.1
    gaussian_iou: 0.7
    class_offset: 1
  per_category_metrics: false
  weight_decay: 0.0005
  gradient_clip_norm: 10.0
  annotation_file: '/readahead/200M/placer/prod/home/tensorflow-performance-data/datasets/coco/instances_val2017.json'
  init_checkpoint: gs://tf_model_garden/vision/centernet/extremenet_hg104_512x512_coco17
  init_checkpoint_modules: 'backbone'
  train_data:
    input_path: '/readahead/200M/placer/prod/home/tensorflow-performance-data/datasets/coco/train*'
    drop_remainder: true
    dtype: 'bfloat16'
    global_batch_size: 128
    is_training: true
    parser:
      aug_rand_hflip: true
      aug_scale_min: 0.6
      aug_scale_max: 1.3
      aug_rand_saturation: true
      aug_rand_brightness: true
      aug_rand_hue: true
      aug_rand_contrast: true
      odapi_augmentation: true
  validation_data:
    input_path: '/readahead/200M/placer/prod/home/tensorflow-performance-data/datasets/coco/val*'
    drop_remainder: false
    dtype: 'bfloat16'
    global_batch_size: 64
    is_training: false
trainer:
  train_steps: 140000
  validation_steps: 78  # 5000 / 64
  steps_per_loop: 924  # 118287 / 128
  validation_interval: 924
  summary_interval: 924
  checkpoint_interval: 924
  optimizer_config:
    learning_rate:
      type: 'cosine'
      cosine:
        initial_learning_rate: 0.001
        decay_steps: 140000
    optimizer:
      type: adam
      adam:
        epsilon: 0.0000001
    warmup:
      type: 'linear'
      linear:
        warmup_steps: 2000