soft placement

core/leras/layers/MsSsim.py

@@ -1,6 +1,5 @@
 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)
@@ -21,7 +20,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])
 
-        loss = tf2.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
+        loss = tf.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,7 +31,6 @@ class nn():
     current_DeviceConfig = None
 
     tf = None
-    tf2 = None
     tf_sess = None
     tf_sess_config = None
     tf_default_device = None
@@ -87,8 +86,6 @@ class nn():
 
             if tf_version[0] == '2':
                 tf = tensorflow.compat.v1
-                tf2 = tensorflow
-                nn.tf2 = tf2
             else:
                 tf = tensorflow
 
@@ -115,7 +112,7 @@ class nn():
                 config = tf.ConfigProto(device_count={'GPU': 0})
             else:
                 nn.tf_default_device = "/GPU:0"
-                config = tf.ConfigProto()
+                config = tf.ConfigProto(allow_soft_placement=True)
                 config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices])
 
             config.gpu_options.force_gpu_compatible = True