Multi-scale discriminator feature-wise loss function

Graphics and Displays group - University of Cambridge

Last update: Dec 12, 2022

Related tags

Deep Learning mdf

Overview

Multi-Scale Discriminative Feature Loss

This repository provides code for Multi-Scale Discriminative Feature (MDF) loss for image reconstruction algorithms.

Description

Central to the application of neural networks in image restoration problems, such as single image super resolution, is the choice of a loss function that encourages natural and perceptually pleasing results. We provide a lightweight feature extractor that outperforms state-of-the-art loss functions in single image super resolution, denoising, and JPEG artefact removal. We propose a novel Multi-Scale Discriminative Feature (MDF) loss comprising a series of discriminators, trained to penalize errors introduced by a generator. For further information please refer to the project webpage.

Usage

The code runs in Python3 and Pytorch.

First install the dependencies by running:

pip3 install -r requirements.txt

To run a simple example, optimizing image pixels:

import torch as pt
import torch.optim as optim
import imageio
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable

from mdfloss import MDFLoss


# Set parameters
cuda_available = False
epochs = 25
application = 'Denoising'
image_path = './misc/i10.png'

if application =='SISR':
    path_disc = "./weights/Ds_SISR.pth"
elif application == 'Denoising':
    path_disc = "./weights/Ds_Denoising.pth"
elif application == 'JPEG':
    path_disc = "./weights/Ds_JPEG.pth"

# Read reference images
imgr = imageio.imread(image_path)
imgr = pt.from_numpy(imageio.core.asarray(imgr/255.0))
imgr = imgr.type(dtype=pt.float64)
imgr = imgr.permute(2,0,1)
imgr = imgr.unsqueeze(0).type(pt.FloatTensor)

# Create a noisy image 
imgd = pt.rand(imgr.size())

if cuda_available:
    imgr = imgr.cuda()
    imgd = imgd.cuda()

# Convert images to variables to support gradients
imgrb = Variable( imgr, requires_grad = False)
imgdb = Variable( imgd, requires_grad = True)

optimizer = optim.Adam([imgdb], lr=0.1)

# Initialise the loss
criterion = MDFLoss(path_disc, cuda_available=cuda_available)

# Iterate over the epochs optimizing for the noisy image
for ii in range(0,epochs):
    
    optimizer.zero_grad()
    loss = criterion(imgrb,imgdb) 
    print("Epoch: ",ii," loss: ", loss.item())
    loss.backward()
    optimizer.step()

Citing

If using, please cite:

@article{mustafa2021training,
  title={Training a Better Loss Function for Image Restoration},
  author={Mustafa, Aamir and Mikhailiuk, Aliaksei and Iliescu, Dan Andrei and Babbar, Varun and Mantiuk, Rafal K},
  journal={arXiv preprint arXiv:2103.14616},
  year={2021}
}

Acknowledgement

This project has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement N◦ 725253–EyeCode).

Comments

Loading SinGAN

Hey, I'm trying to load the SISR weights after cloning the project and I got the following error:

ModuleNotFoundError: No module named 'SinGAN'

May the .pth files are courrupted?

Thanks

opened by korenleven 6
how to define the generator

Dear sir, after I read the code, there are a few questions. How to train the generator mentioned in the paper? And how to define the generator mentioned in this paper?

opened by KKKloveQbh 4
Load the weights file for multi-gpu training

Hi! I liked your approach and wanted to try it for SISR task, however, I am facing issues while loading the weights file when I want to train it on multiple GPUs. I guess the torch.load() directly is causing the problem here https://github.com/gfxdisp/mdf/blob/85c9ad36734f50e30d162e327194e8dbfa52f0cf/mdfloss.py#L10

Can you please pass it as a model instance where we can load it as model.load_state_dict(torch.load("...")) ?? That might be helpful if anyone wants to train using multiple GPUs. Also can you please share the opt argparse arguments mentioned in the SinGAN/models.py file?? https://github.com/gfxdisp/mdf/blob/85c9ad36734f50e30d162e327194e8dbfa52f0cf/SinGAN/models.py#L19

Please let me know if you have a workaround for this (loading the weights file on multiple GPUs) and it would also be helpful if you could release the training code.

Thanks!

opened by rushi-the-neural-arch 2
ModuleNotFoundError: No module named 'SinGAN'

Hello, I was pretty thrilled to attempt your loss after reading your paper, however, when I tried to adopt the MDF loss by using Denoising.pth it didn't work, it keeps giving me this issue, could you please tell me how to use your loss in another model?

opened by Tihetna1891 1
How to train the discriminators?

Hi, I'm interested in your work "Training a Task-Specific Image Reconstruction Loss". I'm wondering how to train the discriminators used for calculate the MDF loss (the phase 1 mentioned in the paper) ? It seems that the training code is not provided in this repo.

opened by prokaryote-hgy 0

Multi-scale discriminator feature-wise loss function

Related tags

Overview

Multi-Scale Discriminative Feature Loss

Description

Usage

Citing

Acknowledgement

Comments

Loading SinGAN

how to define the generator

Load the weights file for multi-gpu training

ModuleNotFoundError: No module named 'SinGAN'

How to train the discriminators?

Owner

Graphics and Displays group - University of Cambridge

An implementation for the loss function proposed in Decoupled Contrastive Loss paper.

HistoSeg : Quick attention with multi-loss function for multi-structure segmentation in digital histology images

Code for the paper "Training GANs with Stronger Augmentations via Contrastive Discriminator" (ICLR 2021)

Official PyTorch implementation of the paper "Recycling Discriminator: Towards Opinion-Unaware Image Quality Assessment Using Wasserstein GAN", accepted to ACM MM 2021 BNI Track.

StyleGAN2 with adaptive discriminator augmentation (ADA) - Official TensorFlow implementation

Vertical Federated Principal Component Analysis and Its Kernel Extension on Feature-wise Distributed Data based on Pytorch Framework

Recall Loss for Semantic Segmentation (This repo implements the paper: Recall Loss for Semantic Segmentation)

Implement of "Training deep neural networks via direct loss minimization" in PyTorch for 0-1 loss

Seach Losses of our paper 'Loss Function Discovery for Object Detection via Convergence-Simulation Driven Search', accepted by ICLR 2021.

PyTorch implementation of Soft-DTW: a Differentiable Loss Function for Time-Series in CUDA

clDice - a Novel Topology-Preserving Loss Function for Tubular Structure Segmentation

Implementation of "A Deep Learning Loss Function based on Auditory Power Compression for Speech Enhancement" by pytorch

Supervised Sliding Window Smoothing Loss Function Based on MS-TCN for Video Segmentation

"MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction" (CVPRW 2022) & (Winner of NTIRE 2022 Challenge on Spectral Reconstruction from RGB)

Official pytorch implementation of "Feature Stylization and Domain-aware Contrastive Loss for Domain Generalization" ACMMM 2021 (Oral)

PyTorch reimplementation of the Smooth ReLU activation function proposed in the paper "Real World Large Scale Recommendation Systems Reproducibility and Smooth Activations" [arXiv 2022].

Multi-Scale Geometric Consistency Guided Multi-View Stereo

The code repository for "RCNet: Reverse Feature Pyramid and Cross-scale Shift Network for Object Detection" (ACM MM'21)

Official implementation for the paper: "Multi-label Classification with Partial Annotations using Class-aware Selective Loss"