more general atari/common

examples/Atari2600/DQN.py

@@ -171,7 +171,7 @@ def get_config():
             HumanHyperParamSetter(ObjAttrParam(dataset_train, 'exploration'), 'hyper.txt'),
             RunOp(lambda: M.update_target_param()),
             dataset_train,
-            PeriodicCallback(Evaluator(EVAL_EPISODE), 2),
+            PeriodicCallback(Evaluator(EVAL_EPISODE, 'fct/output:0'), 2),
         ]),
         # save memory for multiprocess evaluator
         session_config=get_default_sess_config(0.3),
@@ -194,10 +194,15 @@ if __name__ == '__main__':
         assert args.load is not None
     ROM_FILE = args.rom
 
+    if args.task != 'train':
+        cfg = PredictConfig(
+                model=Model(),
+                session_init=SaverRestore(args.load),
+                output_var_names=['fct/output:0'])
         if args.task == 'play':
-        play_model(Model(), args.load)
+            play_model(cfg)
         elif args.task == 'eval':
-        eval_model_multithread(Model(), args.load, EVAL_EPISODE)
+            eval_model_multithread(cfg, EVAL_EPISODE)
     else:
         config = get_config()
         if args.load:

examples/Atari2600/common.py

@@ -28,13 +28,8 @@ def play_one_episode(player, func, verbose=False):
         return act
     return np.mean(player.play_one_episode(f))
 
-def play_model(M, model_path):
+def play_model(cfg):
     player = get_player(viz=0.01)
-    cfg = PredictConfig(
-            model=M,
-            input_data_mapping=[0],
-            session_init=SaverRestore(model_path),
-            output_var_names=['fct/output:0'])
     predfunc = get_predict_func(cfg)
     while True:
         score = play_one_episode(player, predfunc)
@@ -73,25 +68,21 @@ def eval_with_funcs(predict_funcs, nr_eval):
             return (stat.average, stat.max)
         return (0, 0)
 
-def eval_model_multithread(M, model_path, nr_eval):
-    cfg = PredictConfig(
-            model=M,
-            input_data_mapping=[0],
-            session_init=SaverRestore(model_path),
-            output_var_names=['fct/output:0'])
+def eval_model_multithread(cfg, nr_eval):
     func = get_predict_func(cfg)
     NR_PROC = min(multiprocessing.cpu_count() // 2, 8)
     mean, max = eval_with_funcs([func] * NR_PROC, nr_eval)
     logger.info("Average Score: {}; Max Score: {}".format(mean, max))
 
 class Evaluator(Callback):
-    def __init__(self, nr_eval):
+    def __init__(self, nr_eval, output_name):
         self.eval_episode = nr_eval
+        self.output_name = output_name
 
     def _before_train(self):
         NR_PROC = min(multiprocessing.cpu_count() // 2, 8)
         self.pred_funcs = [self.trainer.get_predict_func(
-           ['state'], ['fct/output'])] * NR_PROC
+           ['state'], [self.output_name])] * NR_PROC
 
     def _trigger_epoch(self):
         t = time.time()

tensorpack/predict/common.py

@@ -84,9 +84,6 @@ def get_predict_func(config):
     config.session_init.init(sess)
 
     def run_input(dp):
-        assert len(input_map) == len(dp), \
-            "Graph has {} inputs but dataset only gives {} components!".format(
-                    len(input_map), len(dp))
         feed = dict(zip(input_map, dp))
         return sess.run(output_vars, feed_dict=feed)
     return run_input