Merge pull request #34 from shubhamagarwal92/shm/speedup

Speedup: torch.cuda.set_device(device) before emptying cache, and less frequent cache refreshes
batra-mlp-lab · May 29, 2020 · 3bdfbf3 · 3bdfbf3
2 parents 556def3 + 0e010a4
commit 3bdfbf3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -37,3 +37,5 @@ logs/
 # virtualenv
 venv/
 ENV/
+
+.idea/
diff --git a/train.py b/train.py
@@ -113,6 +113,7 @@
     if args.gpu_ids[0] >= 0
     else torch.device("cpu")
 )
+torch.cuda.set_device(device)
 
 # Print config and args.
 print(yaml.dump(config, default_flow_style=False))
@@ -279,7 +280,7 @@ def lr_lambda_fun(current_iteration: int) -> float:
 
         scheduler.step(global_iteration_step)
         global_iteration_step += 1
-        torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
 
     # -------------------------------------------------------------------------
     #   ON EPOCH END  (checkpointing and validation)