github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-05 20:42:11 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

AMP fix

Browse source
This commit is contained in:
iperov 2021-05-30 09:24:23 +04:00
parent 11a7993238
commit e52b53f87c
2 changed files with 4 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

9
core/leras/ops/__init__.py
View file

@ -333,7 +333,9 @@ def depth_to_space(x, size):
x = tf.reshape(x, (-1, oh, ow, oc, )) x = tf.reshape(x, (-1, oh, ow, oc, ))
return x return x
else: 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() b,c,h,w = x.shape.as_list()
oh, ow = h * size, w * size oh, ow = h * size, w * size
oc = c // (size * size) oc = c // (size * size)
@ -344,11 +346,6 @@ def depth_to_space(x, size):
return x return x
nn.depth_to_space = depth_to_space 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): def rgb_to_lab(srgb):
srgb_pixels = tf.reshape(srgb, [-1, 3]) srgb_pixels = tf.reshape(srgb, [-1, 3])
linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32) linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)

5
models/Model_AMP/Model.py
View file

@ -121,7 +121,7 @@ class AMPModel(ModelBase):
def on_initialize(self): def on_initialize(self):
device_config = nn.getCurrentDeviceConfig() device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices 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) nn.initialize(data_format=self.model_data_format)
tf = nn.tf tf = nn.tf
@ -262,8 +262,6 @@ class AMPModel(ModelBase):
m = tf.nn.sigmoid(self.out_convm(m)) m = tf.nn.sigmoid(self.out_convm(m))
return x, m return x, m
self.face_type = {'wf' : FaceType.WHOLE_FACE, self.face_type = {'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ] 'head' : FaceType.HEAD}[ self.options['face_type'] ]
@ -287,7 +285,6 @@ class AMPModel(ModelBase):
if ct_mode == 'none': if ct_mode == 'none':
ct_mode = None ct_mode = None
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu'] 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' 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' optimizer_vars_on_cpu = models_opt_device=='/CPU:0'