在我的模型中使用Pytorch SSIM损失函数
我正在尝试这个由this repo实现的SSIM丢失，用于图像恢复。
为了参考作者的GitHub上的原始示例代码，我尝试了：
model.train()
for epo in range(epoch):
for i, data in enumerate(trainloader, 0):
inputs = data
inputs = Variable(inputs)
optimizer.zero_grad()
inputs = inputs.view(bs, 1, 128, 128)
top = model.upward(inputs)
outputs = model.downward(top, shortcut = True)
outputs = outputs.view(bs, 1, 128, 128)
if i % 20 == 0:
out = outputs[0].view(128, 128).detach().numpy() * 255
cv2.imwrite("/home/tk/Documents/recover/SSIM/" + str(epo) + "_" + str(i) + "_re.png", out)
loss = - criterion(inputs, outputs)
ssim_value = - loss.data.item()
print (ssim_value)
loss.backward()
optimizer.step()然而，结果并没有像我预期的那样出来。在前10个时期之后，打印的结果图像都是黑色的。
loss = - criterion(inputs, outputs)是作者提出的，然而，对于经典的Pytorch训练代码，这将是loss = criterion(y_pred, target)，因此这里应该是loss = criterion(inputs, outputs)。
然而，我尝试了loss = criterion(inputs, outputs)，但结果仍然相同。
有人能分享一些关于如何正确利用SSIM损失的想法吗？谢谢。
回答 2
Stack Overflow用户
发布于 2019-02-07 21:58:33
作者正在尝试最大化SSIM值。对pytorch损失函数和优化器工作的自然理解是减少损失。但SSIM值是质量度量，因此越高越好。因此，作者使用
loss = - criterion(inputs, outputs)
您可以尝试使用
loss = 1 - criterion(inputs, outputs)
如本paper所述。
使用this repo测试上述内容的修改代码(max_ssim.py)
import pytorch_ssim
import torch
from torch.autograd import Variable
from torch import optim
import cv2
import numpy as np
npImg1 = cv2.imread("einstein.png")
img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0)/255.0
img2 = torch.rand(img1.size())
if torch.cuda.is_available():
img1 = img1.cuda()
img2 = img2.cuda()
img1 = Variable( img1, requires_grad=False)
img2 = Variable( img2, requires_grad = True)
print(img1.shape)
print(img2.shape)
# Functional: pytorch_ssim.ssim(img1, img2, window_size = 11, size_average = True)
ssim_value = 1-pytorch_ssim.ssim(img1, img2).item()
print("Initial ssim:", ssim_value)
# Module: pytorch_ssim.SSIM(window_size = 11, size_average = True)
ssim_loss = pytorch_ssim.SSIM()
optimizer = optim.Adam([img2], lr=0.01)
while ssim_value > 0.05:
optimizer.zero_grad()
ssim_out = 1-ssim_loss(img1, img2)
ssim_value = ssim_out.item()
print(ssim_value)
ssim_out.backward()
optimizer.step()
cv2.imshow('op',np.transpose(img2.cpu().detach().numpy()[0],(1,2,0)))
cv2.waitKey()Stack Overflow用户
发布于 2021-01-11 06:21:33
将相似度(越高越好)转换为损失的常用方法是计算1 - similarity(x, y)。
要创建此损失，您可以创建一个新的“函数”。
def ssim_loss(x, y):
return 1. - ssim(x, y)或者，如果相似性是一个类(nn.Module)，您可以重载它以创建一个新的类。
class SSIMLoss(SSIM):
def forward(self, x, y):
return 1. - super().forward(x, y)此外，还有比这个repo更好的SSIM实现。例如，piqa Python包的速度更快。该程序包可与一起安装
pip install piqa解决你的问题
from piqa import SSIM
class SSIMLoss(SSIM):
def forward(self, x, y):
return 1. - super().forward(x, y)
criterion = SSIMLoss() # .cuda() if you need GPU support
...
loss = criterion(x, y)
...应该工作得很好。
https://stackoverflow.com/questions/53956932
复制相似问题
腾讯云开发者
