Используйте функцию потери Pytorch SSIM в моей модели

Question

Я испытываю эту потерю 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? Благодарю.

Answer 1

Автор пытается максимизировать значение SSIM. Естественное понимание функции потери мощности и оптимизатора заключается в уменьшении потерь. Но значение SSIM является мерой качества и, следовательно, чем выше, тем лучше. Следовательно, автор использует
loss = - criterion(inputs, outputs)

Вместо этого вы можете попробовать использовать
loss = 1 - criterion(inputs, outputs)
как описано в этой статье.


Модифицированный код (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()

Answer 2

Обычный способ превратить сходство (чем выше, тем лучше) в потерю - это вычислить 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)

Кроме того, существуют лучшие реализации 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)
...

должно работать хорошо.