Merge pull request #176 from PhanatosZou/master

Compatible with TF2.0
nl8590687 · Mar 28, 2020 · 4e4aad0 · 4e4aad0
2 parents a6a8af3 + 7b9630a
commit 4e4aad0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/test_mspeech.py b/test_mspeech.py
@@ -17,10 +17,12 @@
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 #进行配置，使用90%的GPU
-config = tf.ConfigProto()
+config = tf.compat.v1.ConfigProto()
 config.gpu_options.per_process_gpu_memory_fraction = 0.9
 #config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
-set_session(tf.Session(config=config))
+sess = tf.compat.v1.Session(config=config)
+tf.compat.v1.keras.backend.set_session(sess)
+
 
 
 datapath = ''

diff --git a/train_mspeech.py b/train_mspeech.py
@@ -16,10 +16,11 @@
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 #进行配置，使用95%的GPU
-config = tf.ConfigProto()
+config = tf.compat.v1.ConfigProto()
 config.gpu_options.per_process_gpu_memory_fraction = 0.95
 #config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
-set_session(tf.Session(config=config))
+sess = tf.compat.v1.Session(config=config)
+tf.compat.v1.keras.backend.set_session(sess)
 
 
 datapath = ''