saoyan/dncnn-pytorch

PyTorch implementation of the TIP2017 paper "Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising"

DnCNN-PyTorch

This is a PyTorch implementation of the TIP2017 paper Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising. The author's MATLAB implementation is here.

---

This code was written with PyTorch<0.4, but most people must be using PyTorch>=0.4 today. Migrating the code is easy. Please refer to PyTorch 0.4.0 Migration Guide.

---

How to run

1. Dependences

• PyTorch(<0.4)
• torchvision
• OpenCV for Python
• HDF5 for Python
• tensorboardX (TensorBoard for PyTorch)

2. Train DnCNN-S (DnCNN with known noise level)

python train.py \
  --preprocess True \
  --num_of_layers 17 \
  --mode S \
  --noiseL 25 \
  --val_noiseL 25

NOTE

• If you've already built the training and validation dataset (i.e. train.h5 & val.h5 files), set preprocess to be False.
• According to the paper, DnCNN-S has 17 layers.
• noiseL is used for training and val_noiseL is used for validation. They should be set to the same value for unbiased validation. You can set whatever noise level you need.

3. Train DnCNN-B (DnCNN with blind noise level)

python train.py \
  --preprocess True \
  --num_of_layers 20 \
  --mode B \
  --val_noiseL 25

NOTE

• If you've already built the training and validation dataset (i.e. train.h5 & val.h5 files), set preprocess to be False.
• According to the paper, DnCNN-B has 20 layers.
• noiseL is ingnored when training DnCNN-B. You can set val_noiseL to whatever you need.

4. Test

python test.py \
  --num_of_layers 17 \
  --logdir logs/DnCNN-S-15 \
  --test_data Set12 \
  --test_noiseL 15

NOTE

• Set num_of_layers to be 17 when testing DnCNN-S models. Set num_of_layers to be 20 when testing DnCNN-B model.
• test_data can be Set12 or Set68.
• test_noiseL is used for testing. This should be set according to which model your want to test (i.e. logdir).

Test Results

BSD68 Average RSNR

Noise Level	DnCNN-S	DnCNN-B	DnCNN-S-PyTorch	DnCNN-B-PyTorch
15	31.73	31.61	31.71	31.60
25	29.23	29.16	29.21	29.15
50	26.23	26.23	26.22	26.20

Set12 Average PSNR

Noise Level	DnCNN-S	DnCNN-B	DnCNN-S-PyTorch	DnCNN-B-PyTorch
15	32.859	32.680	32.837	32.725
25	30.436	30.362	30.404	30.344
50	27.178	27.206	27.165	27.138

Tricks useful for boosting performance

• Parameter initialization:
  Use kaiming_normal initialization for Conv; Pay attention to the initialization of BatchNorm

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.kaiming_normal(m.weight.data, a=0, mode='fan_in')
    elif classname.find('Linear') != -1:
        nn.init.kaiming_normal(m.weight.data, a=0, mode='fan_in')
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(mean=0, std=math.sqrt(2./9./64.)).clamp_(-0.025,0.025)
        nn.init.constant(m.bias.data, 0.0)

• The definition of loss function
  Set size_average to be False when defining the loss function. When size_average=True, the pixel-wise average will be computed, but what we need is sample-wise average.

criterion = nn.MSELoss(size_average=False)

The computation of loss will be like:

loss = criterion(out_train, noise) / (imgn_train.size()[0]*2)

where we divide the sum over one batch of samples by 2N, with N being # samples.

dncnn-pytorch's Issues

Can only process grayscale images?

About the residual strategy

Dear author,
I `ve found a difference between your train code and CSZN DNCNN pytorch train code.

1. In his train code, the loss is calculated like this:

loss = criterion(model(batch_y), batch_x)_
Since the output of his model was clean image. He define the loss by output and clean-groundtruth.

2. But in ur train code, the loss is calculated like this:
   loss = criterion(out_train, noise) / (imgn_train.size()[0]*2)
   Since the output of ur model was the residual noise map. You define the loss by output and the noise-groundtruth

So the question is:

• Would this difference influence the training result?
• For denoising, is it maybe a better way to take clean image as ground-truth instead of noise residual?

Implementation may be wrong

Hello and thanks for your reproduction!
When I went through model.py, I noticed that in the forward method, you returned out instead of x - out, which is corresponding to residual learning in the author's article, I wonder if it's a mistake or I missed something. Thanks!

Testing DnCNN on CPU

Hi SaoYan,

Thank you for sharing your codes at GitHub!
I was trying to test the model on my Mac machine -- but I am having issue with it, can you please advise as how to proceed with that. The models have been trained using GPU and I am trying on CPU -- so it is causing problem when I try to load the model. I am a newbie to PyTorch, any help in this regards will be great.

Thanks,
Touqeer

FFDNet in PyTorch

Hi SaoYan,

I was wondering if by any chance you are working or already have PyTorch version of FFDNet. It is an advanced version of DnCNN -- but only MatConvNet version is available from author's GitHub.

https://github.com/cszn/FFDNet

Thanks,
Touqeer

make my own dataset

i have some images with fuzzy&small target,i wanna to get a higher signal-to-miscellaneous ratio and higher quality image,how can i make noise images and .h5 train &test file??

AttributeError: 'Variable' object has no attribute 'item'

Hi, thank you for your share. When I train DnCNN-S, I get this error after I get the train.h5 and the val.h5.
Could you please help me solve it?
......
file: data/Set12/09.png
file: data/Set12/10.png
file: data/Set12/11.png
file: data/Set12/12.png
training set, # samples 238400

val set, # samples 12

Loading dataset ...

of training samples: 238400

learning rate 0.001000
Traceback (most recent call last):
File "train.py", line 126, in
main()
File "train.py", line 88, in main
(epoch+1, i+1, len(loader_train), loss.item(), psnr_train))
File "/usr/local/lib/python2.7/dist-packages/torch/autograd/variable.py", line 67, in getattr
return object.getattribute(self, name)
AttributeError: 'Variable' object has no attribute 'item'

color images

What changes required for color images?

Torch not compiled with CUDA enabled

how to solve

h5py wrong

File "h5py_objects.pyx", line 54, in h5py._objects.with_phil.wrapper
File "h5py_objects.pyx", line 55, in h5py._objects.with_phil.wrapper
File "h5py\h5f.pyx", line 106, in h5py.h5f.open
OSError: Unable to open file (file signature not found)
下面这段代码出错（wrong)
class Dataset(udata.Dataset):
def init(self, train=True):
super(Dataset, self).init()
self.train = train
if self.train:
h5f = h5py.File('train.h5', 'r')
# h5f = h5py.File('C:/pytorch/DnCNN-PyTorch-master/train.py', 'r')
else:
h5f = h5py.File('val.h5', 'r')
self.keys = list(h5f.keys())
random.shuffle(self.keys)
h5f.close()
def len(self):
return len(self.keys)
def getitem(self, index):
if self.train:
h5f = h5py.File('train.h5', 'r')
else:
h5f = h5py.File('val.h5', 'r')
key = self.keys[index]
data = np.array(h5f[key])
h5f.close()
return torch.Tensor(data)

Why do you set bias = False?

Why is the bias term in all the convolution layers is false?

Thanks!

关于残差与loss function问题

您好！
在您写的代码中，有一段没有看懂，望指点：
在train.py中，
out_train = model(imgn_train)
loss = criterion(out_train, noise) / (imgn_train.size()[0]*2)
...

results

model.eval()
out_train = torch.clamp(imgn_train-model(imgn_train), 0., 1.)
psnr_train = batch_PSNR(out_train, img_train, 1.)

1.为什么已经定义了out_train，后面还要重新定义out_train呢
2.为什么要除imgn_train.size()[0]*2。0.,1. 是什么意思啊
3.残差是指的代码中的哪个呢，文章中是 l = 1/2N ||R(y) - (y -x) ||2，文章中的残差是R(y)吧

新手，谢谢指点