#!/usr/bin/env bash
# SimCLR ResNet-18 pretraining on CIFAR-10 (50 epochs, fp16), then the repo's linear evaluation (100 epochs).
set -o pipefail
cd "$(dirname "$0")"
export TQDM_DISABLE=1
PY=../env/bin/python
SEED=$(sed -n 's/^torch.manual_seed(\([0-9]*\)).*/\1/p' simclr.py)
echo "KVERITAS_INPUT src=seed:$SEED"
echo "KVERITAS_PHASE name=pretrain"
CKPT=""
while IFS= read -r line; do
  echo "$line"
  case "$line" in
    Epoch:*) set -- $line
      echo "KVERITAS_METRIC name=contrastive_loss value=$4 step=$2"
      echo "KVERITAS_METRIC name=contrastive_top1 value=$7 step=$2" ;;
    Checkpoint:*) CKPT=${line#Checkpoint: } ;;
  esac
done < <($PY -u run.py -data ../data/datasets -dataset-name cifar10 --epochs 50 --fp16-precision 2>&1)
[ -n "$CKPT" ] || { echo "no checkpoint"; exit 1; }
echo "KVERITAS_PHASE name=linear_eval"
$PY -u feature_eval/linear_eval.py "$CKPT" ../data/datasets 2>&1 | awk -F'\t' '
{ print; fflush() }
/^Epoch [0-9]+\t/ {
  split($1,e," "); split($2,a," "); split($3,b,": "); split($4,c,": ")
  printf "KVERITAS_METRIC name=linear_train_top1 value=%s step=%s\n", a[4], e[2]
  printf "KVERITAS_METRIC name=linear_test_top1 value=%s step=%s\n", b[2], e[2]
  last1=b[2]; last5=c[2]; fflush()
}
END { if (last1!="") printf "KVERITAS_CLAIM metric=linear_eval_test_top1 value=%s\nKVERITAS_CLAIM metric=linear_eval_test_top5 value=%s\n", last1, last5 }'
