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

fix: transpose batches for tf msssim

Browse source
This commit is contained in:
jh 2021-03-12 08:26:14 -08:00
commit 0df0d6d5cb
1 changed files with 4 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -325,7 +325,10 @@ def ms_ssim(img1, img2, resolution, kernel_size=11, k1=0.01, k2=0.03, max_value=
if sum(power_factors) < 1.0:
power_factors = [x/sum(power_factors) for x in power_factors]
ms_ssim_val = tf.image.ssim_multiscale(img1, img2, max_val=max_value, power_factors=power_factors,
# Transpose images from NCHW to NHWC
img1_t = tf.transpose(img1, [0, 2, 3, 1])
img2_t = tf.transpose(img2, [0, 2, 3, 1])
ms_ssim_val = tf.image.ssim_multiscale(img1_t, img2_t, max_val=max_value, power_factors=power_factors,
filter_size=kernel_size, k1=k1, k2=k2)
loss = (1.0 - ms_ssim_val) / 2.0