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

assign device

Browse source
This commit is contained in:
jh 2021-03-12 18:27:00 -08:00
commit beb8c30970
1 changed files with 11 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

13
core/leras/layers/MsSsim.py
View file

@ -19,7 +19,16 @@ class MsSsim(nn.LayerBase):
# Transpose images from NCHW to NHWC # Transpose images from NCHW to NHWC
y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) 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]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1])
loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors)
return (1.0 - loss) / 2.0
def assign_device(op):
if op.type != 'Assert':
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)
return (1.0 - loss) / 2.0
nn.MsSsim = MsSsim nn.MsSsim = MsSsim