# --- RESNET50 --- #
resnet50: &resnet50
  model: resnet50
  data: imagenet
  precision: "16.16"
  enable_stochastic_rounding: true
  label_smoothing: 0.1
  weight_decay: 0.0001
  optimizer: sgd_combined
  lr_schedule: cosine
  warmup_epoch: 5
  epoch: 100
  loss_scaling: 128
  lr_scheduler_freq: 16
  half_partial: true
  seed: 42
  momentum: 0.9
  lr: 2.0
  micro_batch_size: 18
  pipeline_splits: []
  replicas: 16
  gradient_accumulation: 57
  device_iterations: 1
  norm_type: batch
  disable_stable_batchnorm: true
  offload_optimizer: false
  enable_optimizer_rts: false
  eight_bit_io: true
  normalization_location: ipu
  use_bbox_info: true
  exchange_memory_target: memory
  recompute_mode: manual
  recompute_checkpoints: [add, add_1, add_2, add_3, add_5, add_6, add_7, add_9, add_11, add_12, layer4/0/conv1, layer4/0/bn1, layer4/0/conv2, layer4/0/bn2, layer4/0/conv3, layer4/0/bn3, layer4/0/downsample_0, layer4/0/downsample/1, add_13, layer4/1/conv1, layer4/1/bn1, layer4/1/conv2, layer4/1/bn2, layer4/1/conv3, layer4/1/bn3, add_14, layer4/2/conv1, layer4/2/bn1, layer4/2/conv2, layer4/2/bn2, layer4/2/conv3, layer4/2/bn3, add_15, avgpool, fc]
  available_memory_proportion: [0.16]

resnet50-pod4:
  <<: *resnet50
  replicas: 4
  gradient_accumulation: 228

resnet50-pod64:
  <<: *resnet50
  replicas: 64
  gradient_accumulation: 14


# --- EfficientNet --- #
efficientnet-g16-gn-pod16-base: &efficientnet-g16-gn-16ipu-base
  data: imagenet
  seed: 0
  precision: "16.16"
  eight_bit_io: true
  normalization_location: ipu
  use_bbox_info: true
  enable_stochastic_rounding: false
  efficientnet_group_dim: 16
  efficientnet_expand_ratio: 4
  norm_type: group
  norm_eps: 0.001
  norm_num_groups: 4
  epoch: 350
  warmup_epoch: 5
  optimizer: rmsprop_tf
  offload_optimizer: false
  enable_optimizer_rts: false
  optimizer_eps: 0.001
  momentum: 0.9
  rmsprop_decay: 0.953125
  lr_schedule: exponential
  lr_scheduler_freq: 0.424
  lr_decay: 0.97
  loss_scaling: 128
  weight_decay: 0.00001
  label_smoothing: 0.1
  weight_avg_strategy: exponential
  weight_avg_exp_decay: 0.97
  available_memory_proportion: [0.15]
  exchange_memory_target: balanced
  half_partial: true
  enable_fast_groupnorm: true


efficientnet-b0-g16-gn-pod16:
  <<: *efficientnet-g16-gn-16ipu-base
  model: efficientnet-b0
  pipeline_splits: [blocks/3/0/conv_pw]
  lr: 0.05
  micro_batch_size: 20
  replicas: 8
  gradient_accumulation: 5
  mixup_alpha: 0.1


efficientnet-b4-g16-gn-pod16: &efficientnet-b4-g16-gn-pod16
  <<: *efficientnet-g16-gn-16ipu-base
  model: efficientnet-b4
  precision: "16.16"
  enable_stochastic_rounding: true
  micro_batch_size: 4
  available_memory_proportion: [0.1]
  pipeline_splits: [blocks/3/0/bn2]
  recompute_checkpoints: [add_1, add_3]
  lr: 0.1
  replicas: 8
  gradient_accumulation: 64
  device_iterations: 4
  mixup_alpha: 0.2
  cutmix_lambda_low: 0.65
  cutmix_lambda_high: 1.0
  cutmix_disable_prob: 0.35
  precision: "16.16"
  exchange_memory_target: memory
  enable_optimizer_rts: true


efficientnet-b4-g16-gn-pod64:
  <<: *efficientnet-b4-g16-gn-pod16
  replicas: 32
  gradient_accumulation: 16


# Converges to 82.14% val accuracy
efficientnet-b4-g1-gn-pod16: &efficientnet-b4-g1-gn-pod16
  <<: *efficientnet-b4-g16-gn-pod16
  efficientnet_group_dim: 1
  efficientnet_expand_ratio: 6
  replicas: 4
  micro_batch_size: 4
  gradient_accumulation: 48
  lr: 0.0483
  pipeline_splits: [blocks/1/1/conv_pw, blocks/3/0/conv_pw, blocks/4/5/conv_pw]
  recompute_checkpoints: []


# --- MobileNet v3--- #
mobilenet-v3-base: &mobilenet-v3-base
  data: imagenet
  label_smoothing: 0.1
  weight_decay: 0.00003
  optimizer: sgd_combined
  lr_schedule: cosine



  warmup_epoch: 5
  initial_loss_scaling: 128
  loss_scaling: 128
  seed: 42
  norm_type: batch
  offload_optimizer: false
  enable_optimizer_rts: false
  half_partial: true
  eight_bit_io: true
  normalization_location: ipu
  epoch: 360
  momentum: 0.9
  lr: 0.52
  device_iterations: 1
  validation_mode: after
  micro_batch_size: 32


mobilenet-v3-small: &mobilenet-v3-small
  <<: *mobilenet-v3-base
  model: mobilenet-v3-small
  precision: "16.16"
  recompute_mode: auto
  gradient_accumulation: 32
  replicas: 1
  enable_stochastic_rounding: true


mobilenet-v3-small-pod4:
  <<: *mobilenet-v3-small
  replicas: 4
  gradient_accumulation: 8
  enable_stochastic_rounding: false


mobilenet-v3-large: &mobilenet-v3-large
  <<: *mobilenet-v3-base
  model: mobilenet-v3-large
  precision: "16.16"
  pipeline_splits: [mb_block/2/seq/0, mb_block/7/seq/0, mb_block/11/seq/0]
  gradient_accumulation: 32
  replicas: 1
  enable_stochastic_rounding: true


mobilenet-v3-large-pod16:
  <<: *mobilenet-v3-large
  replicas: 4
  gradient_accumulation: 8
  enable_stochastic_rounding: false
