AMP fix

core/leras/ops/__init__.py

@@ -333,7 +333,9 @@ def depth_to_space(x, size):
         x = tf.reshape(x, (-1, oh, ow, oc, ))
         return x
     else:
-        return tf.depth_to_space(x, size, data_format=nn.data_format)
+        cfg = nn.getCurrentDeviceConfig()
+        if not cfg.cpu_only:
+            return tf.depth_to_space(x, size, data_format=nn.data_format)
         b,c,h,w = x.shape.as_list()
         oh, ow = h * size, w * size
         oc = c // (size * size)
@@ -344,11 +346,6 @@ def depth_to_space(x, size):
         return x
 nn.depth_to_space = depth_to_space
 
-def pixel_norm(x, power = 1.0):
-    return x * power * tf.rsqrt(tf.reduce_mean(tf.square(x), axis=nn.conv2d_spatial_axes, keepdims=True) + 1e-06)
-nn.pixel_norm = pixel_norm
-
-
 def rgb_to_lab(srgb):
     srgb_pixels = tf.reshape(srgb, [-1, 3])
     linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)

models/Model_AMP/Model.py

@@ -121,7 +121,7 @@ class AMPModel(ModelBase):
     def on_initialize(self):
         device_config = nn.getCurrentDeviceConfig()
         devices = device_config.devices
-        self.model_data_format = "NCHW"# if len(devices) != 0 and not self.is_debug() else "NHWC"
+        self.model_data_format = "NCHW"
         nn.initialize(data_format=self.model_data_format)
         tf = nn.tf
 
@@ -262,8 +262,6 @@ class AMPModel(ModelBase):
                 m = tf.nn.sigmoid(self.out_convm(m))
                 return x, m
 
-
-
         self.face_type = {'wf' : FaceType.WHOLE_FACE,
                           'head' : FaceType.HEAD}[ self.options['face_type'] ]
 
@@ -287,7 +285,6 @@ class AMPModel(ModelBase):
         if ct_mode == 'none':
             ct_mode = None
 
-
         models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
         models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
         optimizer_vars_on_cpu = models_opt_device=='/CPU:0'