Update 2 files

- /scripts/benchmark_sbatch_submit.sh
- /README.md

README.md

@@ -83,9 +83,14 @@ bash scripts/copy_data_for_benchmark.sh
  
 You can choose either singularity or enroot in the scripts ```benchmark_single_node.sbatch``` and ```benchmark_multi_node.sbatch```.
 
-The input argument is the number of nodes used for benchmarking.
+The input arguments are: 
+1. Number of nodes,
+2. Number of gpus per node,
+3. Number of iterations for each parameter setting.
+
+We will average the benchmark performance over the iterations. We use a batch size of 128 which is the maximum usable batch size without a OOM error.
 ```
-bash scripts/benchmark_sbatch_submit.sh 1
+bash scripts/benchmark_sbatch_submit.sh 1 8 100
 ```
 
 ### Results  

scripts/benchmark_sbatch_submit.sh

+#!/bin/bash
+
 mkdir -p sbatch_out
 
 dim=2
 NUM_NODES=$1
+NUM_GPUS=$2
+NUM_ITERATIONS=$3
 if [ $NUM_NODES -eq 1 ]; then
     for nodes in {1..1}; do
-        for gpus in {1..8}; do
+        for gpus in {1..$NUM_GPUS}; do
             for batch_size in 128; do
-                for iteration in {1..100}; do
+                for iteration in {1..$NUM_ITERATIONS}; do
 
                     echo dim ${dim}, nodes ${nodes}, gpus ${gpus}, batch_size ${batch_size}, iteration ${iteration}
                     # For single node
@@ -20,9 +24,9 @@ if [ $NUM_NODES -eq 1 ]; then
     done
 else
     for nodes in {2..$NUM_NODES}; do
-        for gpus in {8}; do
+        for gpus in {$NUM_GPUS}; do
             for batch_size in 128; do
-                for iteration in {1..100}; do
+                for iteration in {1..$NUM_ITERATIONS}; do
 
                     # For multi node
                     sbatch -o sbatch_out/benchmark_dim${dim}_nodes${nodes}_gpus${gpus}_batchsize${batch_size}_iteration${iteration}.out scripts/benchmark_multi_node.sbatch ${dim} ${nodes} ${gpus} ${batch_size} ${iteration}