# Copyright (c) 2021 Graphcore Ltd. All rights reserved.

imagenet_config: &imagenet
  dataset: imagenet
  dataset-path: /localdata/datasets/imagenet-data

cifar10_config: &cifar10
  dataset: cifar10
  dataset-path: /localdata/datasets

resnet8_test:
  model-name: cifar_resnet8
  dataset: cifar10
  micro-batch-size: 8
  num-epochs: 1
  eight-bit-transfer: False
  gradient-accumulation-count: 1
  num-replicas: 1
  validation: False
  dataset-path: /localdata/datasets

resnet8_test_lars:
  model-name: cifar_resnet8
  dataset: cifar10
  micro-batch-size: 8
  num-epochs: 50
  eight-bit-transfer: False
  gradient-accumulation-count: 1
  num-replicas: 1
  validation: True
  dataset-path: /localdata/datasets
  precision: 32.32
  accelerator-side-preprocess: False
  l2-regularization: 1e-5
  optimizer: lars
  optimizer-params:
    'eeta': 0.001
    'momentum': 0.9
    'epsilon': 0
    'weight_decay': 2.5e-6
  lr-schedule: 'const'
  lr-schedule-params:
    'initial_learning_rate': 0.000125

resnet8_16ipus_bn: &resnet8_16ipus_bn
  <<: *cifar10
  num-epochs: 70
  validation: True
  num-replicas: 16
  logs-per-epoch: 1
  precision: '16.16'
  half-partials: True
  model: cifar_resnet8
  micro-batch-size: 25
  global-batch-size: 2000
  norm-layer:
    'name': 'batch_norm'
  lr-schedule: 'cosine'
  lr-warmup-params:
    'warmup_mode': 'shift'
    'warmup_epochs': 5
  lr-schedule-params:
    'initial_learning_rate': 0.0001220703125
    'epochs_to_total_decay': 65
  stochastic-rounding: 'ON'
  eight-bit-transfer: True
  l2-regularization: 0.0001
  accelerator-side-preprocess: True

resnet50_16ipus_16k_bn_pipeline:
  <<: *imagenet
  model-name: resnet50
  half-partials: True
  recomputation: True
  logs-per-epoch: 1
  eight-bit-transfer: True
  accelerator-side-preprocess: True
  num-epochs: 100
  validation: True
  micro-batch-size: 16
  global-batch-size: 16384
  num-replicas: 4
  optimizer-params:
    'momentum': 0.9
  pipeline-schedule: Grouped
  pipeline-splits: [
    conv2_block2_out,
    conv3_block1_out,
    conv4_block2_out]
  precision: '16.32'
  available-memory-proportion: [15]
  loss-scaling: 128
  l2-regularization: 0.0001
  stochastic-rounding: 'ON'
  lr-schedule: 'cosine'
  lr-warmup-params:
    'warmup_mode': 'shift'
    'warmup_epochs': 5
  lr-schedule-params:
    'initial_learning_rate': 0.0001220703125
    'epochs_to_total_decay': 95
  lr-staircase: True

resnet50_16ipus_8k_bn_pipeline: &resnet50_16ipus_8k_bn_pipeline
  <<: *imagenet
  model-name: resnet50
  micro-batch-size: 16
  global-batch-size: 8192
  num-replicas: 4
  num-epochs: 100
  validation: True
  precision: '16.16'
  half-partials: True
  recomputation: True
  stable-norm: True
  optimizer-state-offloading: False
  logs-per-epoch: 1
  eight-bit-transfer: True
  accelerator-side-preprocess: True
  optimizer-params:
    'momentum': 0.9
  loss-scaling: 128
  l2-regularization: 0.0001
  pipeline-schedule: Grouped
  pipeline-splits: [
    conv2_block2_out,
    conv3_block1_out,
    conv4_block2_out]
  available-memory-proportion: [15]
  stochastic-rounding: 'ON'
  lr-schedule: 'cosine'
  lr-schedule-params:
    'initial_learning_rate': 0.0001220703125
    'epochs_to_total_decay': 100
  lr-warmup-params:
    'warmup_mode': 'mask'
    'warmup_epochs': 5
  lr-staircase: True

resnet56_16ipus_8k_bn_pipeline:
  <<: *resnet50_16ipus_8k_bn_pipeline
  model: cifar_resnet56
  available-memory-proportion: [13.5]
  micro-batch-size: 16
  pipeline-validation-model: True
  pipeline-splits: [
    conv2_block4_post_addition_relu,
    conv2_block7_post_addition_relu,
    conv3_block4_post_addition_relu]

resnet50_16ipus_8k_gn_pipeline:
  <<: *resnet50_16ipus_8k_bn_pipeline
  norm-layer:
    'name': 'group_norm'
    'num_groups': 32

resnet50_mlperf_bs16: &resnet50_mlperf_bs16
  <<: *imagenet
  model-name: resnet50
  micro-batch-size: 16
  global-batch-size: 2304
  loss-scaling: 128
  num-replicas: 16
  optimizer: sgd
  optimizer-params:
    'momentum': 0.9
  num-epochs: 44
  l2-regularization: 2.5e-05
  lr-schedule: 'polynomial'
  lr-schedule-params:
    'initial_learning_rate': 0.001991319444444
    'epochs_to_total_decay': 39
    'end_learning_rate_ratio': 0.0000218
    'power': 2
  lr-warmup-params:
    'warmup_mode': 'shift'
    'warmup_epochs': 5
  norm-layer:
    'name': 'custom_batch_norm'
    'momentum': 0.9
  label-smoothing: 0.1
  precision: '16.16'
  half-partials: True
  optimizer-state-offloading: False
  pipeline-schedule: Sequential
  pipeline-splits: [
    conv2_block1_out,
    conv2_block3_out,
    conv3_block4_out]
  device-mapping: [0, 0, 0, 0]
  recomputation: True
  available-memory-proportion: [20]
  eight-bit-transfer: True
  accelerator-side-preprocess: True
  internal-exchange-optimization-target: memory
  max-cross-replica-buffer-size: 52428800
  gather-conv-output: False
  min-remote-tensor-size: 128
  dbn-replica-group-size: 2
  fused-preprocessing: True
  ckpts-per-epoch: 1/4
  validation-micro-batch-size: 108

resnet50_mlperf_pod16_bs20: &resnet50_mlperf_pod16_bs20
  <<: *resnet50_mlperf_bs16
  micro-batch-size: 20
  global-batch-size: 1920
  available-memory-proportion: [15]
  gather-conv-output: True
  num-epochs: 40
  pipeline-splits: [
    conv2_block1_out,
    conv2_block3_out,
    conv3_block1_out,
    conv3_block3_out,
    conv4_block1_out,
    conv4_block4_out]
  device-mapping: [0, 0, 0, 0, 0, 0, 0]
  validation-micro-batch-size: 115

resnet50_mlperf_pod4_bs20:
  <<: *resnet50_mlperf_pod16_bs20
  num-replicas: 4
  global-batch-size: 480

resnet50_mlperf_pod16_lars: &resnet50_mlperf_pod16_lars
  <<: [*imagenet, *resnet50_mlperf_bs16]
  micro-batch-size: 18
  global-batch-size: 3520
  pipeline-splits: [
    conv2_block1_out,
    conv2_block3_out,
    conv3_block1_out,
    conv3_block3_out,
    conv4_block1_out,
    conv4_block4_out]
  device-mapping: [0, 0, 0, 0, 0, 0, 0]
  optimizer: lars
  optimizer-params:
    'momentum': 0.9
    'weight_decay': 0.00005
    'eeta': 0.001
    'epsilon': 0
  available-memory-proportion: [15]
  l2-regularization: 5e-05
  num-epochs: 38
  lr-schedule-params:
    'initial_learning_rate': 0.003045455
    'epochs_to_total_decay': 36
    'end_learning_rate_ratio': 9.32836e-06
    'power': 2
  lr-warmup-params:
    'warmup_mode': 'shift'
    'warmup_epochs': 2
  first-ckpt-epoch: 2
  max-reduce-many-buffer-size: 216

resnet50_infer_base: &resnet50_infer_base
  model-name: resnet50
  micro-batch-size: 1
  global-batch-size: 1
  optimizer: sgd
  optimizer-params:
    'momentum': 0.9
  num-epochs: 1
  l2-regularization: 2.5e-05
  lr-schedule: 'polynomial'
  lr-schedule-params:
    'initial_learning_rate': 0.001991319444444
    'epochs_to_total_decay': 39
    'end_learning_rate_ratio': 0.0000218
    'power': 2
  lr-warmup-params:
    'warmup_mode': 'shift'
    'warmup_epochs': 5
  norm-layer:
    'name': 'custom_batch_norm'
    'momentum': 0.9
  label-smoothing: 0.1
  precision: '16.16'
  half-partials: True
  optimizer-state-offloading: False
  recomputation: False
  available-memory-proportion: [20]
  eight-bit-transfer: True
  accelerator-side-preprocess: True
  internal-exchange-optimization-target: memory
  gather-conv-output: False
  min-remote-tensor-size: 128
  fused-preprocessing: True

resnet50_infer_bs1:
  <<: [*imagenet, *resnet50_infer_base]

resnet50_infer_bs16:
  <<: [*imagenet, *resnet50_infer_base]
  micro-batch-size: 16
  global-batch-size: 16

resnet50_infer_bs32:
  <<: [*imagenet, *resnet50_infer_base]
  micro-batch-size: 32
  global-batch-size: 32

resnet50_infer_test:
  <<: [*imagenet, *resnet50_infer_base]
