Skip to content

Commit 4a4b5b8

Browse files
cszncszn
authored
Add training code with DagNN
1 parent 52faf88 commit 4a4b5b8

26 files changed

+1565
-0
lines changed

‎TrainingCodes/DnCNN_TrainingCodes_DagNN_v1.1/Demo_Test_DnCNN_DAG.m

+120
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
2+
% @article{zhang2017beyond,
3+
% title={Beyond a {Gaussian} denoiser: Residual learning of deep {CNN} for image denoising},
4+
% author={Zhang, Kai and Zuo, Wangmeng and Chen, Yunjin and Meng, Deyu and Zhang, Lei},
5+
% journal={IEEE Transactions on Image Processing},
6+
% year={2017},
7+
% volume={26},
8+
% number={7},
9+
% pages={3142-3155},
10+
% }
11+
12+
% by Kai Zhang (1/2018)
13+
% [email protected]
14+
% https://github.com/cszn
15+
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
16+
17+
% clear; clc;
18+
19+
%% testing set
20+
addpath(fullfile('utilities'));
21+
22+
folderModel = 'model';
23+
folderTest = 'testsets';
24+
folderResult= 'results';
25+
imageSets = {'BSD68','Set12'}; % testing datasets
26+
setTestCur = imageSets{2}; % current testing dataset
27+
28+
29+
showresult = 1;
30+
gpu = 1;
31+
32+
33+
noiseSigma = 25;
34+
35+
% load model
36+
epoch = 50;
37+
38+
modelName = 'DnCNN';
39+
40+
% case one: for the model in 'data/model'
41+
%load(fullfile('data',folderModel,[modelName,'-epoch-',num2str(epoch),'.mat']));
42+
43+
% case two: for the model in 'utilities'
44+
load(fullfile('utilities',[modelName,'-epoch-',num2str(epoch),'.mat']));
45+
46+
47+
48+
49+
net = dagnn.DagNN.loadobj(net) ;
50+
51+
net.removeLayer('loss') ;
52+
out1 = net.getVarIndex('prediction') ;
53+
net.vars(net.getVarIndex('prediction')).precious = 1 ;
54+
55+
net.mode = 'test';
56+
57+
if gpu
58+
net.move('gpu');
59+
end
60+
61+
% read images
62+
ext = {'*.jpg','*.png','*.bmp'};
63+
filePaths = [];
64+
for i = 1 : length(ext)
65+
filePaths = cat(1,filePaths, dir(fullfile(folderTest,setTestCur,ext{i})));
66+
end
67+
68+
folderResultCur = fullfile(folderResult, [setTestCur,'_',int2str(noiseSigma)]);
69+
if ~isdir(folderResultCur)
70+
mkdir(folderResultCur)
71+
end
72+
73+
74+
% PSNR and SSIM
75+
PSNRs = zeros(1,length(filePaths));
76+
SSIMs = zeros(1,length(filePaths));
77+
78+
79+
for i = 1 : length(filePaths)
80+
81+
% read image
82+
label = imread(fullfile(folderTest,setTestCur,filePaths(i).name));
83+
[~,nameCur,extCur] = fileparts(filePaths(i).name);
84+
[w,h,c]=size(label);
85+
if c==3
86+
label = rgb2gray(label);
87+
end
88+
89+
% add additive Gaussian noise
90+
randn('seed',0);
91+
noise = noiseSigma/255.*randn(size(label));
92+
input = im2single(label) + single(noise);
93+
94+
if gpu
95+
input = gpuArray(input);
96+
end
97+
net.eval({'input', input}) ;
98+
% output (single)
99+
output = gather(squeeze(gather(net.vars(out1).value)));
100+
101+
102+
% calculate PSNR and SSIM
103+
[PSNRCur, SSIMCur] = Cal_PSNRSSIM(label,im2uint8(output),0,0);
104+
if showresult
105+
imshow(cat(2,im2uint8(input),im2uint8(label),im2uint8(output)));
106+
title([filePaths(i).name,' ',num2str(PSNRCur,'%2.2f'),'dB',' ',num2str(SSIMCur,'%2.4f')])
107+
imwrite(im2uint8(output), fullfile(folderResultCur, [nameCur, '_' int2str(noiseSigma),'_PSNR_',num2str(PSNRCur*100,'%4.0f'), extCur] ));
108+
drawnow;
109+
% pause()
110+
end
111+
PSNRs(i) = PSNRCur;
112+
SSIMs(i) = SSIMCur;
113+
end
114+
115+
116+
disp([mean(PSNRs),mean(SSIMs)]);
117+
118+
119+
120+

‎TrainingCodes/DnCNN_TrainingCodes_DagNN_v1.1/Demo_Train_DnCNN_DAG.m

+61
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
2+
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
3+
% @article{zhang2017beyond,
4+
% title={Beyond a {Gaussian} denoiser: Residual learning of deep {CNN} for image denoising},
5+
% author={Zhang, Kai and Zuo, Wangmeng and Chen, Yunjin and Meng, Deyu and Zhang, Lei},
6+
% journal={IEEE Transactions on Image Processing},
7+
% year={2017},
8+
% volume={26},
9+
% number={7},
10+
% pages={3142-3155},
11+
% }
12+
13+
% by Kai Zhang (1/2018)
14+
% [email protected]
15+
% https://github.com/cszn
16+
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
17+
18+
% The training data is generated by '[imdb] = generatepatches;' in line 126 of 'DnCNN_train_dag.m'.
19+
20+
rng('default')
21+
22+
addpath('utilities');
23+
%-------------------------------------------------------------------------
24+
% Configuration
25+
%-------------------------------------------------------------------------
26+
opts.modelName = 'DnCNN'; % model name
27+
opts.learningRate = [logspace(-3,-3,22) logspace(-4,-4,105)];% you can change the learning rate
28+
opts.batchSize = 128; %
29+
opts.gpus = [1];
30+
opts.numSubBatches = 2;
31+
32+
% solver
33+
opts.solver = 'Adam'; % global
34+
opts.derOutputs = {'objective',1} ;
35+
36+
opts.backPropDepth = Inf;
37+
%-------------------------------------------------------------------------
38+
% Initialize model
39+
%-------------------------------------------------------------------------
40+
41+
net = feval([opts.modelName,'_Init']);
42+
43+
%-------------------------------------------------------------------------
44+
% Train
45+
%-------------------------------------------------------------------------
46+
47+
[net, info] = DnCNN_train_dag(net, ...
48+
'learningRate',opts.learningRate, ...
49+
'derOutputs',opts.derOutputs, ...
50+
'numSubBatches',opts.numSubBatches, ...
51+
'backPropDepth',opts.backPropDepth, ...
52+
'solver',opts.solver, ...
53+
'batchSize', opts.batchSize, ...
54+
'modelname', opts.modelName, ...
55+
'gpus',opts.gpus) ;
56+
57+
58+
59+
60+
61+

‎TrainingCodes/DnCNN_TrainingCodes_DagNN_v1.1/DnCNN_Init.m

+213
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
function net = DnCNN_Init()
2+
3+
% by Kai Zhang (1/2018)
4+
% [email protected]
5+
% https://github.com/cszn
6+
7+
% Create DAGNN object
8+
net = dagnn.DagNN();
9+
10+
% conv + relu
11+
blockNum = 1;
12+
inVar = 'input';
13+
channel= 1; % grayscale image
14+
dims = [3,3,channel,64];
15+
pad = [1,1];
16+
stride = [1,1];
17+
lr = [1,1];
18+
[net, inVar, blockNum] = addConv(net, blockNum, inVar, dims, pad, stride, lr);
19+
[net, inVar, blockNum] = addReLU(net, blockNum, inVar);
20+
21+
for i = 1:15
22+
% conv + bn + relu
23+
dims = [3,3,64,64];
24+
pad = [1,1];
25+
stride = [1,1];
26+
lr = [1,0];
27+
[net, inVar, blockNum] = addConv(net, blockNum, inVar, dims, pad, stride, lr);
28+
n_ch = dims(4);
29+
[net, inVar, blockNum] = addBnorm(net, blockNum, inVar, n_ch);
30+
[net, inVar, blockNum] = addReLU(net, blockNum, inVar);
31+
end
32+
33+
% conv
34+
dims = [3,3,64,channel];
35+
pad = [1,1];
36+
stride = [1,1];
37+
lr = [1,0]; % or [1,1], it does not influence the results
38+
[net, inVar, blockNum] = addConv(net, blockNum, inVar, dims, pad, stride, lr);
39+
40+
% sum
41+
inVar = {inVar,'input'};
42+
[net, inVar, blockNum] = addSum(net, blockNum, inVar);
43+
44+
outputName = 'prediction';
45+
net.renameVar(inVar,outputName)
46+
47+
% loss
48+
net.addLayer('loss', dagnn.Loss('loss','L2'), {'prediction','label'}, {'objective'},{});
49+
net.vars(net.getVarIndex('prediction')).precious = 1;
50+
51+
52+
end
53+
54+
55+
56+
57+
% Add a Concat layer
58+
function [net, inVar, blockNum] = addConcat(net, blockNum, inVar)
59+
60+
outVar = sprintf('concat%d', blockNum);
61+
layerCur = sprintf('concat%d', blockNum);
62+
63+
block = dagnn.Concat('dim',3);
64+
net.addLayer(layerCur, block, inVar, {outVar},{});
65+
66+
inVar = outVar;
67+
blockNum = blockNum + 1;
68+
end
69+
70+
71+
% Add a loss layer
72+
function [net, inVar, blockNum] = addLoss(net, blockNum, inVar)
73+
74+
outVar = 'objective';
75+
layerCur = sprintf('loss%d', blockNum);
76+
77+
block = dagnn.Loss('loss','L2');
78+
net.addLayer(layerCur, block, inVar, {outVar},{});
79+
80+
inVar = outVar;
81+
blockNum = blockNum + 1;
82+
end
83+
84+
85+
% Add a sum layer
86+
function [net, inVar, blockNum] = addSum(net, blockNum, inVar)
87+
88+
outVar = sprintf('sum%d', blockNum);
89+
layerCur = sprintf('sum%d', blockNum);
90+
91+
block = dagnn.Sum();
92+
net.addLayer(layerCur, block, inVar, {outVar},{});
93+
94+
inVar = outVar;
95+
blockNum = blockNum + 1;
96+
end
97+
98+
99+
% Add a relu layer
100+
function [net, inVar, blockNum] = addReLU(net, blockNum, inVar)
101+
102+
outVar = sprintf('relu%d', blockNum);
103+
layerCur = sprintf('relu%d', blockNum);
104+
105+
block = dagnn.ReLU('leak',0);
106+
net.addLayer(layerCur, block, {inVar}, {outVar},{});
107+
108+
inVar = outVar;
109+
blockNum = blockNum + 1;
110+
end
111+
112+
113+
% Add a bnorm layer
114+
function [net, inVar, blockNum] = addBnorm(net, blockNum, inVar, n_ch)
115+
116+
trainMethod = 'adam';
117+
outVar = sprintf('bnorm%d', blockNum);
118+
layerCur = sprintf('bnorm%d', blockNum);
119+
120+
params={[layerCur '_g'], [layerCur '_b'], [layerCur '_m']};
121+
net.addLayer(layerCur, dagnn.BatchNorm('numChannels', n_ch), {inVar}, {outVar},params) ;
122+
123+
pidx = net.getParamIndex({[layerCur '_g'], [layerCur '_b'], [layerCur '_m']});
124+
b_min = 0.025;
125+
net.params(pidx(1)).value = clipping(sqrt(2/(9*n_ch))*randn(n_ch,1,'single'),b_min);
126+
net.params(pidx(1)).learningRate= 1;
127+
net.params(pidx(1)).weightDecay = 0;
128+
net.params(pidx(1)).trainMethod = trainMethod;
129+
130+
net.params(pidx(2)).value = zeros(n_ch, 1, 'single');
131+
net.params(pidx(2)).learningRate= 1;
132+
net.params(pidx(2)).weightDecay = 0;
133+
net.params(pidx(2)).trainMethod = trainMethod;
134+
135+
net.params(pidx(3)).value = [zeros(n_ch,1,'single'), 0.01*ones(n_ch,1,'single')];
136+
net.params(pidx(3)).learningRate= 1;
137+
net.params(pidx(3)).weightDecay = 0;
138+
net.params(pidx(3)).trainMethod = 'average';
139+
140+
inVar = outVar;
141+
blockNum = blockNum + 1;
142+
end
143+
144+
145+
% add a ConvTranspose layer
146+
function [net, inVar, blockNum] = addConvt(net, blockNum, inVar, dims, crop, upsample, lr)
147+
opts.cudnnWorkspaceLimit = 1024*1024*1024*2; % 2GB
148+
convOpts = {'CudnnWorkspaceLimit', opts.cudnnWorkspaceLimit} ;
149+
trainMethod = 'adam';
150+
151+
outVar = sprintf('convt%d', blockNum);
152+
153+
layerCur = sprintf('convt%d', blockNum);
154+
155+
convBlock = dagnn.ConvTranspose('size', dims, 'crop', crop,'upsample', upsample, ...
156+
'hasBias', true, 'opts', convOpts);
157+
158+
net.addLayer(layerCur, convBlock, {inVar}, {outVar},{[layerCur '_f'], [layerCur '_b']});
159+
160+
f = net.getParamIndex([layerCur '_f']) ;
161+
sc = sqrt(2/(dims(1)*dims(2)*dims(4))) ; %improved Xavier
162+
net.params(f).value = sc*randn(dims, 'single');
163+
net.params(f).learningRate = lr(1);
164+
net.params(f).weightDecay = 1;
165+
net.params(f).trainMethod = trainMethod;
166+
167+
f = net.getParamIndex([layerCur '_b']) ;
168+
net.params(f).value = zeros(dims(3), 1, 'single');
169+
net.params(f).learningRate = lr(2);
170+
net.params(f).weightDecay = 1;
171+
net.params(f).trainMethod = trainMethod;
172+
173+
inVar = outVar;
174+
blockNum = blockNum + 1;
175+
end
176+
177+
178+
% add a Conv layer
179+
function [net, inVar, blockNum] = addConv(net, blockNum, inVar, dims, pad, stride, lr)
180+
opts.cudnnWorkspaceLimit = 1024*1024*1024*2; % 2GB
181+
convOpts = {'CudnnWorkspaceLimit', opts.cudnnWorkspaceLimit} ;
182+
trainMethod = 'adam';
183+
184+
outVar = sprintf('conv%d', blockNum);
185+
layerCur = sprintf('conv%d', blockNum);
186+
187+
convBlock = dagnn.Conv('size', dims, 'pad', pad,'stride', stride, ...
188+
'hasBias', true, 'opts', convOpts);
189+
190+
net.addLayer(layerCur, convBlock, {inVar}, {outVar},{[layerCur '_f'], [layerCur '_b']});
191+
192+
f = net.getParamIndex([layerCur '_f']) ;
193+
sc = sqrt(2/(dims(1)*dims(2)*max(dims(3), dims(4)))) ; %improved Xavier
194+
net.params(f).value = sc*randn(dims, 'single') ;
195+
net.params(f).learningRate = lr(1);
196+
net.params(f).weightDecay = 1;
197+
net.params(f).trainMethod = trainMethod;
198+
199+
f = net.getParamIndex([layerCur '_b']) ;
200+
net.params(f).value = zeros(dims(4), 1, 'single');
201+
net.params(f).learningRate = lr(2);
202+
net.params(f).weightDecay = 1;
203+
net.params(f).trainMethod = trainMethod;
204+
205+
inVar = outVar;
206+
blockNum = blockNum + 1;
207+
end
208+
209+
210+
function A = clipping(A,b)
211+
A(A>=0&A<b) = b;
212+
A(A<0&A>-b) = -b;
213+
end

0 commit comments

Comments
 (0)