runtime:
  distribution_strategy: tpu
  mixed_precision_dtype: 'bfloat16'
task:
  init_checkpoint: ''
  model:
    backbone:
      resnet:
        model_id: 50
      type: resnet
    projection_head:
      ft_proj_idx: 1
      num_proj_layers: 3
      proj_output_dim: 128
    backbone_trainable: true
    heads: !!python/tuple
    # Define heads for the PRETRAIN networks here
    - task_name: pretrain_imagenet
      mode: pretrain
    # # Define heads for the FINETUNE networks here
    - task_name: finetune_imagenet_10percent
      mode: finetune
      supervised_head:
        num_classes: 1001
        zero_init: true
    input_size: [224, 224, 3]
    l2_weight_decay: 0.0
    norm_activation:
      norm_epsilon: 1.0e-05
      norm_momentum: 0.9
      use_sync_bn: true
  task_routines: !!python/tuple
  # Define TASK CONFIG for the PRETRAIN networks here
  - task_name: pretrain_imagenet
    task_weight: 30.0
    task_config:
      evaluation:
        one_hot: true
        top_k: 5
      loss:
        l2_weight_decay: 0.0
        projection_norm: true
        temperature: 0.1
      model:
        input_size: [224, 224, 3]
        mode: pretrain
      train_data:
        input_path: /readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*
        input_set_label_to_zero: true    # Set labels to zeros to double confirm that no label is used during pretrain
        is_training: true
        global_batch_size: 4096
        dtype: 'bfloat16'
        parser:
          aug_rand_hflip: true
          mode: pretrain
        decoder:
          decode_label: true
      validation_data:
        input_path: /readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*
        is_training: false
        global_batch_size: 2048
        dtype: 'bfloat16'
        drop_remainder: false
        parser:
          mode: pretrain
        decoder:
          decode_label: true
  # Define TASK CONFIG for the FINETUNE Networks here
  - task_name: finetune_imagenet_10percent
    task_weight: 1.0
    task_config:
      evaluation:
        one_hot: true
        top_k: 5
      loss:
        l2_weight_decay: 0.0
        label_smoothing: 0.0
        one_hot: true
      model:
        input_size: [224, 224, 3]
        mode: finetune
        supervised_head:
          num_classes: 1001
          zero_init: true
      train_data:
        tfds_name: 'imagenet2012_subset/10pct'
        tfds_split: 'train'
        input_path: ''
        is_training: true
        global_batch_size: 1024
        dtype: 'bfloat16'
        parser:
          aug_rand_hflip: true
          mode: finetune
        decoder:
          decode_label: true
      validation_data:
        tfds_name: 'imagenet2012_subset/10pct'
        tfds_split: 'validation'
        input_path: ''
        is_training: false
        global_batch_size: 2048
        dtype: 'bfloat16'
        drop_remainder: false
        parser:
          mode: finetune
        decoder:
          decode_label: true
trainer:
  trainer_type: interleaving
  task_sampler:
    proportional:
      alpha: 1.0
    type: proportional
  train_steps: 32000   # 100 epochs
  validation_steps: 24  # NUM_EXAMPLES (50000) // global_batch_size
  validation_interval: 625
  steps_per_loop: 625  # NUM_EXAMPLES (1281167) // global_batch_size
  summary_interval: 625
  checkpoint_interval: 625
  max_to_keep: 3
  optimizer_config:
    learning_rate:
      cosine:
        decay_steps: 32000
        initial_learning_rate: 4.8
      type: cosine
    optimizer:
      lars:
        exclude_from_weight_decay: [batch_normalization, bias]
        momentum: 0.9
        weight_decay_rate: 1.0e-06
      type: lars
    warmup:
      linear:
        name: linear
        warmup_steps: 3200
      type: linear
