tf2

core/leras/layers/MsSsim.py

@@ -1,5 +1,6 @@
 from core.leras import nn
 tf = nn.tf
+tf2 = nn.tf2
 
 class MsSsim(nn.LayerBase):
     default_power_factors = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
@@ -20,15 +21,7 @@ class MsSsim(nn.LayerBase):
         y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1])
         y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1])
 
-
-        def assign_device(op):
-            if op.type != 'Assert' or op.type != 'ListDiff':
-                return '/gpu:0'
-            else:
-                return '/cpu:0'
-
-        with tf.device(assign_device):
-            loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
+        loss = tf2.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
         return (1.0 - loss) / 2.0
 
 nn.MsSsim = MsSsim

core/leras/nn.py

@@ -31,6 +31,7 @@ class nn():
     current_DeviceConfig = None
 
     tf = None
+    tf2 = None
     tf_sess = None
     tf_sess_config = None
     tf_default_device = None
@@ -86,6 +87,8 @@ class nn():
 
             if tf_version[0] == '2':
                 tf = tensorflow.compat.v1
+                tf2 = tensorflow
+                nn.tf2 = tf2
             else:
                 tf = tensorflow