Batch resizing (#286)

* Add utility to help manage batch resizer * Add scaling config * cleanrljax pins * Update scaling files * Update batch sizer model * Latest batch size model * Update pins --------- Co-authored-by: pierre.delaunay <[email protected]>
mila-iqia · Sep 19, 2024 · dbebc7c · dbebc7c
1 parent cf0df51
commit dbebc7c
Show file tree

Hide file tree

Showing 40 changed files with 626 additions and 259 deletions.
diff --git a/.pin/constraints-cuda-torch.txt b/.pin/constraints-cuda-torch.txt
diff --git a/benchmarks/brax/main.py b/benchmarks/brax/main.py
@@ -85,6 +85,9 @@ def run():
 
     args = parser.parse_args()
 
+    # args.num_envs = (args.batch_size * args.num_minibatches)  
+
+
     train(
         environment=envs.get_environment(env_name=args.env),
         num_timesteps=args.num_timesteps,

diff --git a/benchmarks/brax/requirements.cuda.txt b/benchmarks/brax/requirements.cuda.txt
diff --git a/benchmarks/brax/voirfile.py b/benchmarks/brax/voirfile.py
@@ -20,10 +20,10 @@ class Config:
     skip: int = 5
 
     # Number of rates to log before stopping
-    stop: int = 20
+    stop: int = 60
 
     # Number of seconds between each gpu poll
-    gpu_poll: int = 3
+    gpu_poll: int = 1
 
 
 @configurable