兼容tf2,建议RTX3090使用CUDA11.1版本的ptxas

graph_session.py

@@ -3,6 +3,7 @@
 # Author: kerlomz <[email protected]>
 import os
 import tensorflow as tf
+tf.compat.v1.disable_v2_behavior()
 from tensorflow.python.framework.errors_impl import NotFoundError
 from config import ModelConfig
 
@@ -18,14 +19,14 @@ def __init__(self, model_conf: ModelConfig):
         self.model_name = self.model_conf.model_name
         self.graph_name = self.model_conf.model_name
         self.version = self.model_conf.model_version
-        self.graph = tf.Graph()
-        self.sess = tf.Session(
+        self.graph = tf.compat.v1.Graph()
+        self.sess = tf.compat.v1.Session(
             graph=self.graph,
-            config=tf.ConfigProto(
+            config=tf.compat.v1.ConfigProto(
 
                 # allow_soft_placement=True,
                 # log_device_placement=True,
-                gpu_options=tf.GPUOptions(
+                gpu_options=tf.compat.v1.GPUOptions(
                     # allocator_type='BFC',
                     allow_growth=True,  # it will cause fragmentation.
                     # per_process_gpu_memory_fraction=self.model_conf.device_usage
@@ -49,7 +50,7 @@ def load_model(self):
                 graph_def_file = f.read()
             self.graph_def.ParseFromString(graph_def_file)
             with self.graph.as_default():
-                self.sess.run(tf.global_variables_initializer())
+                self.sess.run(tf.compat.v1.global_variables_initializer())
                 _ = tf.import_graph_def(self.graph_def, name="")
 
             self.logger.info('TensorFlow Session {} Loaded.'.format(self.model_conf.model_name))

package.py

@@ -9,6 +9,7 @@
 import platform
 import distutils
 import tensorflow as tf
+tf.compat.v1.disable_v2_behavior()
 from enum import Enum, unique
 from utils import SystemUtils
 from config import resource_path
@@ -33,7 +34,7 @@ class Version(Enum):
 
 if __name__ == '__main__':
 
-    ver = Version.GPU if tf.test.gpu_device_name() else Version.CPU
+    ver = Version.CPU
 
     upload = False
     server_ip = ""

resource/VERSION

@@ -1 +1 @@
-20201101
+20201109