fixed some args fields in train.py

mines-opt-ml · Jan 16, 2024 · 97886dd · 97886dd
1 parent 05fc386
commit 97886dd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/shortest_path_experiment.sh b/shortest_path_experiment.sh
@@ -2,7 +2,7 @@
 
 # Define the model types
 model_types="DYS CVX PertOpt BBOpt"
-grid_sizes="10 20 30 40 50"
+grid_sizes="5 10 15 20 25 30"
 reps="1 2 3"
 data_dir="./src/shortest_path/shortest_path_data"
 weights_dir="./src/shortest_path/saved_weights"
@@ -18,7 +18,7 @@ do
             rep_data_dir="${data_dir}/$rep"
             rep_weights_dir="${weights_dir}/$rep"
             rep_results_dir="${results_dir}/$rep"
-            python -m src.shortest_path.train --model_type $model_type --grid_size $grid_size --num_data 1000 --data_dir $rep_data_dir --results_dir $rep_results_dir
+            python -m src.shortest_path.train --model_type $model_type --grid_size $grid_size --num_data 1000 --data_dir $rep_data_dir --results_dir $rep_results_dir --device cuda:0
             echo "$grid_size"
         done
     done

diff --git a/src/shortest_path/train.py b/src/shortest_path/train.py
@@ -61,6 +61,8 @@ def main(args):
     parser = argparse.ArgumentParser(description="fpo-dys shortest path")
     parser.add_argument('--grid_size', type=int, default=5)
     parser.add_argument('--data_dir', type=str, default='./src/shortest_path/shortest_path_data/')
+    parser.add_argument('--data_deg', type=int, default=4)
+    parser.add_argument('--data_noise_width', type=float, default=0.5)
     parser.add_argument('--weights_dir', type=str, default='./src/shortest_path/saved_weights/')
     parser.add_argument('--results_dir', type=str, default='./src/shortest_path/results/')
     parser.add_argument('--device', type=str, default='mps')