#!/bin/bash

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

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

pushd $ILSVRC2012_BASEDIR/Tensorflow1.8/Code/tf_cnn_benchmarks

if [ $3 -eq 8 ]
then
####################### 2 nodes #############################################

blaunch $4 -no-wait "export CUDA_VISIBLE_DEVICES=0,1,2,3; python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device gpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 0 --worker_hosts=$4:15000,$5:15000 --ps_hosts=$4:15001,$5:15001 2>&1 > slave_0.txt"

blaunch $4 -no-wait "export CUDA_VISIBLE_DEVICES=0,1,2,3; python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device gpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4  --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name ps --task_index 0 --worker_hosts=$4:15000,$5:15000 --ps_hosts=$4:15001,$5:15001"

blaunch $5 -no-wait "export CUDA_VISIBLE_DEVICES=0,1,2,3; python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device gpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 1 --worker_hosts=$4:15000,$5:15000 --ps_hosts=$4:15001,$5:15001 2>&1 > slave_1.txt"

blaunch $5 "export CUDA_VISIBLE_DEVICES=0,1,2,3; python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device gpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4  --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name ps --task_index 1 --worker_hosts=$4:15000,$5:15000 --ps_hosts=$4:15001,$5:15001"

elif [ $3 -eq 16 ]
then
####################### 4 nodes #############################################
blaunch $4 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device cpu  --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name ps --task_index 0 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000 --ps_hosts=$4:5001

sleep 5

blaunch $4 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 0 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000 --ps_hosts=$4:5001

blaunch $5 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 1 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000 --ps_hosts=$4:5001

blaunch $6 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 2 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000 --ps_hosts=$4:5001

blaunch $7 python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 3 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000 --ps_hosts=$4:5001


elif [ $3 -eq 32 ]
then
####################### 8 nodes #############################################
blaunch $4 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device cpu  --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name ps --task_index 0 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

sleep 5

blaunch $4 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 0 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $5 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 1 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $6 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 2 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $7 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 3 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $8 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 4 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $9 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 5 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $10 -no-wait python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 6 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

blaunch $11 python -u tf_cnn_benchmarks.py --model resnet50 --batch_size=$BATCH_SIZE --num_epochs=$NUM_EPOCH --data_format NCHW --local_parameter_device cpu --num_intra_threads 160 --num_inter_threads 2 --device gpu  --num_gpus 4 --data_name imagenet --data_dir $ILSVRC2012_BASEDIR/Dataset/PostProcessedTFRecords/ --variable_update distributed_replicated --job_name worker --task_index 7 --worker_hosts=$4:5000,$5:5000,$6:5000,$7:5000,$8:5000,$9:5000,$10:5000,$11:5000 --ps_hosts=$4:5001

else
	echo "Invalid number of nodes."
fi

popd

mark=`date +%s`
let dif=mark-beginning

sleep 10

blaunch "pkill python"
blaunch "pkill python"
blaunch "pkill python"

echo "=========== Finished $dif ============"

