#!/bin/bash

BATCH_SIZE=$1
NUM_ITER=$2
NUM_GPUS=$3

beginning=`date +%s`
mark=`date +%s`
let dif=mark-beginning
let enum=NUM_GPUS-1
echo "=========== Setting Caffe benchmark $dif ============"

data_train=$ILSVRC2012_BASEDIR/Dataset/PostProcessedLMDB/train/
data_val=$ILSVRC2012_BASEDIR/Dataset/PostProcessedLMDB/val/
modeldir=$ILSVRC2012_BASEDIR/Caffe/Model
solverfile=$modeldir/VGG_ILSVRC_19_layers_solver.prototxt
netfile=$modeldir/VGG_ILSVRC_19_layers_train_val.prototxt

cp $modeldir/VGG_ILSVRC_19_layers_solver.prototxt.template $solverfile
cp $modeldir/VGG_ILSVRC_19_layers_train_val.prototxt.template $netfile


sed -i -e 's@<num-iterations>@'"$NUM_ITER"'@' $solverfile
sed -i -e 's@<path-to-net-file>@'"$netfile"'@' $solverfile
sed -i -e 's@<path-to-snapshot-dir>@'"$modeldir/Snapshots"'@' $solverfile

sed -i -e 's@<batch-size>@'"$BATCH_SIZE"'@' $netfile
sed -i -e 's@<dataset-train>@'"$data_train"'@' $netfile
sed -i -e 's@<dataset-val>@'"$data_val"'@' $netfile
sed -i -e 's@<mean-file>@'"$data_train/imagenet_mean.binaryproto"'@' $netfile

mark=`date +%s`
let dif=mark-beginning
echo "=========== Checking CUDA state $dif ============="

$CAFFE_BIN/caffe device_query -gpu all

mark=`date +%s`
let dif=mark-beginning
echo "=========== Training VGG-19 $dif $NUM_GPUS GPU  ============="

$CAFFE_BIN/caffe train --solver=$solverfile -gpu `seq -s , 0 $enum`

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