在ubuntu 18.04下,使用PyCharm作为IDE,运行相同的代码,只是一个3线性层网络,以训练RL代理在OpenAI Gym上玩CartPole-v0
,进行3000步。
使用tensorflow 2.0beta1,仅花费14秒:
TensorFlow Version: 2.0.0-beta1
num_actions: 2
obs_shape: (4,)
2019-10-08 09:12:39.648644: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-10-08 09:12:39.669515: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2304000000 Hz
2019-10-08 09:12:39.669920: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x557b0e7eaaa0 executing computations on platform Host. Devices:
2019-10-08 09:12:39.669956: I tensorflow/compiler/xla/service/service.cc:175] StreamExecutor device (0): <undefined>, <undefined>
2019-10-08 09:12:39.747755: W tensorflow/compiler/jit/mark_for_compilation_pass.cc:1483] (One-time warning): Not using XLA:CPU for cluster because envvar TF_XLA_FLAGS=--tf_xla_cpu_global_jit was not set. If you want XLA:CPU, either set that envvar, or use experimental_jit_scope to enable XLA:CPU. To confirm that XLA is active, pass --vmodule=xla_compilation_cache=1 (as a proper command-line flag, not via TF_XLA_FLAGS) or set the envvar XLA_FLAGS=--xla_hlo_profile.
res of test model: 0 [ 0.06576662 -0.00780154]
Before Training: 9 out of 200
losses each 1000 steps: 0.004245305
losses each 1000 steps: 0.002508364
losses each 1000 steps: 0.009408922
The total train time: 14.006552696228027
After Training: 200 out of 200
但是当使用tensorflow 2.0official时,代码完全相同,但是时间成本却高出18倍,这也确实消耗了内存,几乎杀死了我的笔记本电脑:
TensorFlow Version: 2.0.0
num_actions: 2
obs_shape: (4,)
2019-10-08 09:24:19.170213: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-10-08 09:24:19.193118: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2304000000 Hz
2019-10-08 09:24:19.193366: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x564a84d34fb0 executing computations on platform Host. Devices:
2019-10-08 09:24:19.193384: I tensorflow/compiler/xla/service/service.cc:175] StreamExecutor device (0): Host, Default Version
res of test model: 0 [0. 0.]
Before Training: 8 out of 200
losses each 1000 steps: 0.037620276
losses each 1000 steps: 0.0022274784
losses each 1000 steps: 0.002310919
The total train time: 261.44773054122925
After Training: 123 out of 200
我不知道原因是什么,PyCharm或TensorFlow 2.0。
请注意,两者都是tensorflow cpu版本和eager模式。