#!/bin/bash

BATCH_SIZE=$1
NUM_EPOCH=$2
NUM_GPUS=$3

beginning=`date +%s`
mark=`date +%s`
let dif=mark-beginning
echo "=========== Setting TF Horovod benchmark 1 GPU $dif ============"

pushd $ILSVRC2012_BASEDIR/KerasHorovod/Code

mpirun -np $NUM_GPUS \
-bind-to none \
-map-by slot \
-x NCCL_DEBUG=INFO \
-x LD_LIBRARY_PATH \
-x PATH \
-x NCCL_SOCKET_IFNAME=^docker0 \
-x HOROVOD_MPI_THREADS_DISABLE=1 \
--mca orte_tmpdir_base "/tmp" \
--mca btl_tcp_if_exclude lo,docker0 \
python \
$ILSVRC2012_BASEDIR/KerasHorovod/Code/keras_imagenet_resnet50.py \
--batch-size=$BATCH_SIZE \
--epochs=$NUM_EPOCH \
--train-dir="$ILSVRC2012_BASEDIR/Dataset/Decompressed/train" \
--val-dir="$ILSVRC2012_BASEDIR/Dataset/Decompressed/val" \
--log-dir="$TMPDIR/logs" \
--checkpoint-format="$TMPDIR/checkpoint-{epoch}.h5"

popd

mark=`date +%s`
let dif=mark-beginning
echo "=========== Finished $dif ============"
