DmitryUlyanov / deep-image-prior

Image restoration with neural networks but without learning.
https://dmitryulyanov.github.io/deep_image_prior
Other
7.84k stars 1.43k forks source link

TVLoss seems to be too sensitive #8

Closed babaozhouy5 closed 6 years ago

babaozhouy5 commented 6 years ago

In the Super-resolution section, there is a tv_weight&tv_loss, but when i change tv_weight to a small value (default is 0.0), like 0.1/0.01/0.001, i got a -inf PSNR value for LR&HR images. It seems that only tv_weight < 1e-3, eg. tv_weight=1e-4, can get a proper result.

So I change the tv_loss function like this:

def tv_loss(x):
    def _tensor_size(t):
        return t.size()[1]*t.size()[2]*t.size()[3]
    batch_size = x.size()[0]
    h_x = x.size()[2]
    w_x = x.size()[3]
    count_h = _tensor_size(x[:,:,1:,:])
    count_w = _tensor_size(x[:,:,:,1:])
    h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]), 2).sum()
    w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]), 2).sum()
    return 2*(h_tv/count_h+w_tv/count_w)/batch_size

(Thanks for TV Loss) It seems works well, not sensitive. Any suggestions?

DmitryUlyanov commented 6 years ago

Yes, I also experience that a nice values for tv_weigth are about 1e-6 or 1e-7, almost at the point where you lose accuracy due to numerical instabilities. So one way could be to multiply other than TV loss terms by a factor of, say, 100 while dividing learning rate by 100 at the same time.

Notice, that you do not apply square root in your code, to get a vanilla TV loss you should. You TV loss will behave differently to the one with square root. See figure 2 of [1] for example.

[1] Understanding Deep Image Representations by Inverting Them, Aravindh Mahendran, Andrea Vedaldi

babaozhouy5 commented 6 years ago

Thanks for apply :). Yes, i will try this method. The paper looks very good 👍 , I should carefully read it.