Update loss_utils.py

Updated the function ssim(), the optimized version of ssim() is ssim_optimized(), which reduce the computational complexity. After modified, the function create_window() just need to be called once in the main function in train.py, no need to be called in ssim() in every iteration.
This commit is contained in:
Lixing Xiao 2024-07-12 16:56:48 +08:00 committed by GitHub
parent c7ea596c57
commit b66e1ad13e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -26,7 +26,7 @@ def gaussian(window_size, sigma):
def create_window(window_size, channel):
"""Create a 2D Gaussian window."""
_1D_window = gaussian(window_size, 1.5).unsqueeze(1) #此处高斯窗口的sigma固定为1.5
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
return window
@ -63,7 +63,7 @@ def _ssim(img1, img2, window, window_size, channel, size_average=True):
else:
return ssim_map.mean(1).mean(1).mean(1)
#-----------------modify-----------
#-----------------modify------------------------------------
def ssim_optimized(img1, img2, window=None, window_size=11, size_average=True):
channel = img1.size(-3)
if window is None: