enable graph kernel when training dqn on GPU back-end

2021-05-31 15:29:09 +08:00 · 2021-05-31 15:29:09 +08:00 · 65f9fc3f48
parent 209a6e1df6
commit 65f9fc3f48
1 changed files with 3 additions and 0 deletions
--- a/model_zoo/official/rl/dqn/train.py
+++ b/model_zoo/official/rl/dqn/train.py
@ -34,6 +34,9 @@ set_seed(1)

 if __name__ == "__main__":
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
+    if args.device_target == 'GPU':
+        # Enable graph kernel
+        context.set_context(enable_graph_kernel=True, graph_kernel_flags="--enable_parallel_fusion")
    env = gym.make('CartPole-v1')
    cfg.state_space_dim = env.observation_space.shape[0]
    cfg.action_space_dim = env.action_space.n