Create model if not exists in the trainning scripts

BabyBlue26 · May 31, 2018 · 50bac1a · 50bac1a
1 parent b740d5a
commit 50bac1a
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 5 deletions.
diff --git a/rl/scripts/train_il.py b/rl/scripts/train_il.py
@@ -81,7 +81,8 @@
 # Define actor-critic model
 
 acmodel = utils.load_model(obss_preprocessor.obs_space, env.action_space, model_name,
-                           not(args.no_instr), not(args.no_mem), args.arch)
+                           not(args.no_instr), not(args.no_mem), args.arch,
+                           create_if_not_exists=True)
 if torch.cuda.is_available():
     acmodel.cuda()
 

diff --git a/rl/scripts/train_rl.py b/rl/scripts/train_rl.py
@@ -94,7 +94,8 @@
 # Define actor-critic model
 
 acmodel = utils.load_model(obss_preprocessor.obs_space, envs[0].action_space, model_name,
-                           not(args.no_instr), not(args.no_mem), args.arch)
+                           not(args.no_instr), not(args.no_mem), args.arch,
+                           create_if_not_exists=True)
 if torch.cuda.is_available():
     acmodel.cuda()
 

diff --git a/rl/scripts/train_wd.py b/rl/scripts/train_wd.py
@@ -89,7 +89,8 @@
 
 # Define actor-critic model
 acmodel = utils.load_model(obss_preprocessor.obs_space, env.action_space, model_name,
-                           not(args.no_instr), not(args.no_mem), args.arch)
+                           not(args.no_instr), not(args.no_mem), args.arch,
+                           create_if_not_exists=True)
 
 acmodel.train()
 

diff --git a/rl/utils/model.py b/rl/utils/model.py
@@ -11,13 +11,16 @@ def get_model_path(model_name):
     return os.path.join(get_model_dir(model_name), "model.pt")
 
 def load_model(observation_space, action_space, model_name,
-               use_instr=False, use_memory=False, arch="cnn1"):
+               use_instr=False, use_memory=False, arch="cnn1",
+               create_if_not_exists=False):
     path = get_model_path(model_name)
     if os.path.exists(path):
         acmodel = torch.load(path)
-    else:
+    elif create_if_not_exists:
         acmodel = ACModel(observation_space, action_space,
                           use_instr, use_memory, arch)
+    else:
+        raise ValueError("No model at `{}`".format(path))
     acmodel.eval()
     return acmodel