single_process_xla.sh

#!/bin/bash

lrank=$OMPI_COMM_WORLD_LOCAL_RANK
drank=$OMPI_COMM_WORLD_RANK

APP="python3 ./official/vision/image_classification/resnet/resnet_ctl_imagenet_main.py   --num_gpus=1  --skip_eval=true   --batch_size=512 --train_epochs=90  --use_synthetic_data=false  --distribution_strategy=multi_worker_mirrored  --all_reduce_alg=nccl --dtype=fp32  --data_dir=${data_dir}   --task_index=${drank}"
case ${lrank} in
[0])
  export HIP_VISIBLE_DEVICES=0
  TF_XLA_FLAGS="--tf_xla_auto_jit=1" numactl --cpunodebind=0 --membind=0 ${APP}
  ;;
[1])
  export HIP_VISIBLE_DEVICES=1
  TF_XLA_FLAGS="--tf_xla_auto_jit=1" numactl --cpunodebind=1 --membind=1 ${APP}
  ;;
[2])
  export HIP_VISIBLE_DEVICES=2
  TF_XLA_FLAGS="--tf_xla_auto_jit=1" numactl --cpunodebind=2 --membind=2 ${APP}
  ;;
[3])
  export HIP_VISIBLE_DEVICES=3
  TF_XLA_FLAGS="--tf_xla_auto_jit=1" numactl --cpunodebind=3 --membind=3 ${APP}
  ;;
esac