From e52b53f87c38b0dad3fca8dd9a76c15fcda63e46 Mon Sep 17 00:00:00 2001 From: iperov Date: Sun, 30 May 2021 09:24:23 +0400 Subject: [PATCH] AMP fix --- core/leras/ops/__init__.py | 9 +++------ models/Model_AMP/Model.py | 5 +---- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py index 09e8e7a..0cf23a8 100644 --- a/core/leras/ops/__init__.py +++ b/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) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 6b575f3..ab5b8e1 100644 --- a/models/Model_AMP/Model.py +++ b/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'