#!/usr/bin/bash
MODEL=${1:-"r100"}
BZ_PER_DEVICE=${2:-64}
ITER_NUM=${3:-120}
GPUS=${4:-0}
NODE_NUM=${5:-1}
SAMPLE_RATIO=${6:-1.0}
DTYPE=${7:-"fp32"}
TEST_NUM=${8:-1}

a=`expr ${#GPUS} + 1`
gpu_num_per_node=`expr ${a} / 2`
gpu_num=`expr ${gpu_num_per_node} \* ${NODE_NUM}`
total_bz=`expr ${BZ_PER_DEVICE} \* ${gpu_num}`

export CUDA_VISIBLE_DEVICES=${GPUS}
export HOROVOD_GPU_ALLREDUCE=NCCL
export HOROVOD_GPU_ALLGATHER=NCCL
export HOROVOD_GPU_BROADCAST=NCLL
export MXNET_CPU_WORKER_NTHREADS=3

if [ ${NODE_NUM} -eq 1 ] ; then
    node_ip=localhost:${gpu_num_per_node}
elif [ ${NODE_NUM} -eq 2 ] ; then
	node_ip=${NODE1}:${gpu_num_per_node},${NODE2}:${gpu_num_per_node}
elif [ ${NODE_NUM} -eq 4 ] ; then
	node_ip=${NODE1}:${gpu_num_per_node},${NODE2}:${gpu_num_per_node},${NODE3}:${gpu_num_per_node},${NODE4}:${gpu_num_per_node}
else
    echo "Not a valid node."
fi

if [ "$DTYPE" = "fp16" ] ; then
    sed -i "s/\(config.fp16 = \)\S*/config.fp16 = True/" default.py
else
    sed -i "s/\(config.fp16 = \)\S*/config.fp16 = False/" default.py
fi
sed -i "s/\(config.batch_size = \)\S*/config.batch_size = ${BZ_PER_DEVICE}/" default.py
sed -i "s/\(config.max_update = \)\S*/config.max_update = ${ITER_NUM}/" default.py
sed -i "s/\(config.sample_ratio = \)\S*/config.sample_ratio = ${SAMPLE_RATIO}/" default.py


log_folder=./logs-20210222-sample-ratio-${SAMPLE_RATIO}/mxnet/partial_fc/bz${BZ_PER_DEVICE}/${NODE_NUM}n${gpu_num_per_node}g
mkdir -p $log_folder
log_file=$log_folder/${MODEL}_b${BZ_PER_DEVICE}_${DTYPE}_$TEST_NUM.log

# use `which python` to get the absolute path of your python interpreter
# dataset: emore webface glint360k_8GPU ; loss : arcface; cosface

dataset=emore
loss=arcface


PYTHON_EXEC=/home/leinao/anaconda3/envs/mxnet/bin/python
FOLDER=$(dirname $(readlink -f "$0"))

if [ ${NODE_NUM} -eq 1 ] ; then
    export HOROVOD_CACHE_CAPACITY=0
    horovodrun -np ${gpu_num} -H ${node_ip} ${PYTHON_EXEC} \
    ${FOLDER}/train_memory.py \
        --dataset ${dataset}  \
        --loss ${loss} \
        --network ${MODEL}  2>&1 | tee ${log_file}
else
    export HOROVOD_CACHE_CAPACITY=1024
    mpirun --allow-run-as-root -oversubscribe \
    -np ${gpu_num} -H ${node_ip} \
    -x HOROVOD_CACHE_CAPACITY=1024  \
    -bind-to none -map-by slot \
    -x LD_LIBRARY_PATH -x PATH \
    -mca pml ob1 -mca btl ^openib \
    -mca plm_rsh_args "-p 22  -q -o StrictHostKeyChecking=no" \
    -mca btl_tcp_if_include ib0 \
    -x OMP_NUM_THREADS=2 \
    -x MXNET_USE_OPERATOR_TUNING=1 \
    -x MXNET_USE_NUM_CORES_OPERATOR_TUNING=1 \
    -x MXNET_CUDNN_AUTOTUNE_DEFAULT=1 \
    ${PYTHON_EXEC} ${FOLDER}/train_memory.py \
        --dataset ${dataset}  \
        --loss ${loss} \
        --network ${MODEL}  2>&1 | tee ${log_file}
fi
