From 288e51b62333d8b81de0e7ce1ba00013dbfcbd09 Mon Sep 17 00:00:00 2001 From: chen-yiqiang Date: Mon, 31 May 2021 22:25:03 +0800 Subject: [PATCH] add dncc to model zoo --- model_zoo/official/cv/dncnn/README.md | 139 +++++++++++++ model_zoo/official/cv/dncnn/eval.py | 146 ++++++++++++++ model_zoo/official/cv/dncnn/export.py | 52 +++++ model_zoo/official/cv/dncnn/requirements.txt | 5 + .../official/cv/dncnn/scripts/run_eval_gpu.sh | 16 ++ .../cv/dncnn/scripts/run_train_gpu.sh | 21 ++ model_zoo/official/cv/dncnn/src/dataset.py | 182 ++++++++++++++++++ model_zoo/official/cv/dncnn/src/model.py | 59 ++++++ model_zoo/official/cv/dncnn/train.py | 92 +++++++++ 9 files changed, 712 insertions(+) create mode 100755 model_zoo/official/cv/dncnn/README.md create mode 100755 model_zoo/official/cv/dncnn/eval.py create mode 100755 model_zoo/official/cv/dncnn/export.py create mode 100644 model_zoo/official/cv/dncnn/requirements.txt create mode 100644 model_zoo/official/cv/dncnn/scripts/run_eval_gpu.sh create mode 100644 model_zoo/official/cv/dncnn/scripts/run_train_gpu.sh create mode 100755 model_zoo/official/cv/dncnn/src/dataset.py create mode 100755 model_zoo/official/cv/dncnn/src/model.py create mode 100755 model_zoo/official/cv/dncnn/train.py diff --git a/model_zoo/official/cv/dncnn/README.md b/model_zoo/official/cv/dncnn/README.md new file mode 100755 index 00000000000..5a3e67edf72 --- /dev/null +++ b/model_zoo/official/cv/dncnn/README.md @@ -0,0 +1,139 @@ + +### 目录 + +[TOC] + +### DnCNN描述 + +DnCNN是一个使用FCN处理图像降噪任务的模型， 本项目是图像去躁模型DnCNN在mindspore上的复现。 +论文\: Zhang, K., Zuo, W., Chen, Y., Meng, D., & Zhang, L. (2017). Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26(7), 3142-3155. + +### 模型结构 + +网络由N层convolution block组成。其中第一层是conv加reLU，中间n-2层是conv+BN+ReLU，最后一层是单独的conv + +### 数据集 + +训练集DnCNN-S、DnCNN-B使用BSD500中的400张图片 +DnCNN-3使用BSD500中的200张图片和T91中的91张图片 + +测试集包括BDS68，Set5， Set14， clasic5， live1等 + +### 环境要求 + +mindspore=1.1 +skimage=0.18.1 +numpy +PIL +opencv +argparse + +### 快速入门 + +通过官方网站安装MindSpore后，您可以按照如下步骤进行训练和评估： + +```bash +# 训练示例 +python train.py --dataset_path=/path/to/training/data --model_type DnCNN-S --ckpt-prefix=DnCNN-S_25noise --noise_level=25 +# 或者 +sh ./scripts/run_train_gpu.sh /path/to/training/data DnCNN-S DnCNN-S_25noise 25 + +# 评估示例 +python eval.py --dataset_path=/path/to/test/data --ckpt_path=./ckpt/DnCNN-S-50_1800.ckpt --model_type=DnCNN-S --noise_level=25 --noise_type=denoise +# 或者 +sh ./scripts/run_eval_gpu.sh /path/to/test/data ./ckpt/DnCNN-S-50_1800.ckpt DnCNN-S 25 denoise +``` + +### 脚本说明 + +├── readme.md +├── scripts +│ ├── run_eval_gpu.sh //训练shell脚本 +│ └── run_train_gpu.sh //评估shell脚本 +├── src +│ ├── dataset.py //数据读取 +│ └──model.py //模型定义 +├── eval.py //评估脚本 +├── export.py //导出模型 +└── train.py //训练脚本 + +### 训练过程 + +可通过`train.py`脚本中的参数修改训练行为。`train.py`脚本中的参数如下： + +#### 训练参数 + +--dataset_path 训练数据路径 +--model_type 模型类型 = ['DnCNN-S', 'DnCNN-B', 'DnCNN-3'] +--ckpt-prefix 检查点前缀 +--noise_level 噪音等级 +--batch_size 批次大小 +--lr 学习率 +--epoch_num 轮次数 + +#### 默认训练参数 + +optimizer=adam +learning rate=0.001 +batch_size=128 +weight_decay=0.0001 +epoch=50 + +#### 训练 + +只有DnCNN-S 需要指定noise_level + +```python +python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-S --ckpt-prefix=DnCNN-S_25noise --noise_level=25 +python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-B --ckpt-prefix=DnCNN-B +python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-3 --ckpt-prefix=DnCNN-3 +``` + +在ckpt文件夹下保存检查点 + +### 评估过程 + +评估需要通过命令行提供以下参数： +--dataset_path 数据路径 +--ckpt_path 检查点路径 +--model_type 模型类型 +--noise_type 噪音类型， 通过noise_type选择图像测试噪音的类型：["denoise", "super-resolution","jpeg-deblock"] +--noise_level 噪音等级：对应三种noise type的强度，噪音sigma/下采样上采样scale/jpeg压缩quality + +ex: + +```python +python eval.py --dataset_path=/path/to/test/data --ckpt_path=./ckpt/DnCNN-B-50_3000.ckpt --model_type=DnCNN-B --noise_level=50 --noise_type=denoise +``` + +### 模型描述 + +#### 训练准确率结果 + +| 参数 | GPU | +| ------------- | -------------------------------------------------------- | +| 模型版本 | DnCNN-S | +| 资源 | Nvidia V100 | +| mindspore版本 | mindspore 1.1 | +| 数据集 | Berkeley Segmentation Datase | +| 轮次 | 50 | +| 输出 | noise残差 | +| 性能 | 在BSD68测试，PSNR=32.92(σ=15)， 31.73(σ=25)，30.59(σ=50) | + +#### 训练性能结果 + +| 参数 | GPU | +| ------------- | ---------------------------------- | +| 模型版本 | DnCNN-S | +| 资源 | Nvidia V100 | +| mindspore版本 | mindspore 1.1 | +| 训练参数 | lr 0.001, batch_size 128, epoch 50 | +| 优化器 | adam | +| 损失函数 | MSE | +| 输出 | noise残差 | +| 速度 | 320ms/batch | +| 总时长 | 7h9min | +| 检查点 | 6.38M | + + + diff --git a/model_zoo/official/cv/dncnn/eval.py b/model_zoo/official/cv/dncnn/eval.py new file mode 100755 index 00000000000..14683113fbf --- /dev/null +++ b/model_zoo/official/cv/dncnn/eval.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import os +import io +import argparse +import glob + +import PIL +import numpy as np +import cv2 +import skimage.metrics +import mindspore +import mindspore.dataset as ds +from mindspore import context +from mindspore.train.serialization import load_checkpoint, load_param_into_net +import mindspore.dataset.transforms.c_transforms as C +from src.model import DnCNN + +class DnCNN_eval_Dataset(): + def __init__(self, dataset_path, task_type, noise_level): + self.im_list = [] + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.png"))) + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.bmp"))) + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.jpg"))) + self.task_type = task_type + self.noise_level = noise_level + + def __getitem__(self, i): + img = cv2.imread(self.im_list[i], 0) + + if self.task_type == "denoise": + noisy = self.add_noise(img, self.noise_level) + elif self.task_type == "super-resolution": + h, w = img.shape + noisy = cv2.resize(img, (int(w/self.noise_level), int(h/self.noise_level))) + noisy = cv2.resize(noisy, (w, h)) + elif self.task_type == "jpeg-deblock": + noisy = self.jpeg_compression(img, self.noise_level) + + #add channel dimension + noisy = noisy[np.newaxis, :, :] + noisy = noisy / 255.0 + return noisy, img + + def __len__(self): + return len(self.im_list) + + def add_noise(self, im, sigma): + gauss = np.random.normal(0, sigma, im.shape) + noisy = im + gauss + noisy = np.clip(noisy, 0, 255) + noisy = noisy.astype('float32') + return noisy + + def jpeg_compression(self, img, severity): + im_pil = PIL.Image.fromarray(img) + output = io.BytesIO() + im_pil.save(output, 'JPEG', quality=severity) + im_pil = PIL.Image.open(output) + img_np = np.asarray(im_pil) + return img_np + + +def create_eval_dataset(data_path, task_type, noise_level, batch_size=1): + # define dataset + dataset = DnCNN_eval_Dataset(data_path, task_type, noise_level) + dataloader = ds.GeneratorDataset(dataset, ["noisy", "clear"]) + # apply map operations on images + dataloader = dataloader.map(input_columns="noisy", operations=C.TypeCast(mindspore.float32)) + dataloader = dataloader.map(input_columns="clear", operations=C.TypeCast(mindspore.uint8)) + dataloader = dataloader.batch(batch_size, drop_remainder=False) + return dataloader + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="DnCNN") + parser.add_argument("--dataset_path", type=str, default="/code/12imgs-TestingSet/", help='eval image path') + parser.add_argument('--ckpt_path', type=str, default=False, help='trained ckpt') + parser.add_argument('--model_type', type=str, default='DnCNN-S', \ + choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN') + parser.add_argument('--noise_type', type=str, default=False, \ + choices=["denoise", "super-resolution", "jpeg-deblock"], help='trained ckpt') + parser.add_argument('--noise_level', type=int, default=False, help='trained ckpt') + args = parser.parse_args() + + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + ds_eval = create_eval_dataset(args.dataset_path, args.noise_type, args.noise_level, batch_size=1) + print("evaluation image number:", ds_eval.get_dataset_size()) + + if args.model_type == 'DnCNN-S': + network = DnCNN(1, num_of_layers=17) + elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B': + network = DnCNN(1, num_of_layers=20) + else: + print("wrong model type") + exit() + + # load parameter to the network + param_dict = load_checkpoint(args.ckpt_path) + load_param_into_net(network, param_dict) + + mean_psnr = 0 + mean_ssim = 0 + count = 0 + for data in ds_eval.create_dict_iterator(): + clear = data["clear"].asnumpy() + #get denoised image + residual = network(data["noisy"]).asnumpy() * 255 + noisy_img = data["noisy"].asnumpy() * 255 + denoised = np.clip(noisy - residual, 0, 255).astype("uint8") + denoised = np.squeeze(denoised) + clear = np.squeeze(clear) + noisy_img = np.squeeze(noisy_img) + if count == 0: #save example result + cv2.imwrite("noisy.jpg", noisy_img.astype("uint8")) + cv2.imwrite("denoised.jpg", denoised) + cv2.imwrite("original.jpg", clear) + + #calculate psnr + mse = np.mean((clear - denoised)**2) + psnr = 10*np.log10(255*255/mse) + #calculate ssim + ssim = skimage.metrics.structural_similarity(clear, denoised, data_range=255) #skimage 0.18 + + mean_psnr += psnr + mean_ssim += ssim + count += 1 + + mean_psnr = mean_psnr / count + mean_ssim = mean_ssim / count + print("mean psnr", mean_psnr) + print("mean_ssim", mean_ssim) diff --git a/model_zoo/official/cv/dncnn/export.py b/model_zoo/official/cv/dncnn/export.py new file mode 100755 index 00000000000..fbfb9b435d0 --- /dev/null +++ b/model_zoo/official/cv/dncnn/export.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import argparse +import numpy as np + +import mindspore as ms +from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context + +from src.model import DnCNN + +parser = argparse.ArgumentParser(description='DnCNN') +parser.add_argument("--batch_size", type=int, default=1, help="batch size") +parser.add_argument("--image_height", type=int, default=256, help="image_height") +parser.add_argument("--image_width", type=int, default=256, help="image_width") +parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.") +parser.add_argument("--file_name", type=str, default="DnCNN", help="output file name.") +parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='MINDIR', help='file format') +parser.add_argument('--model_type', type=str, default='DnCNN-S', \ + choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN') + +args = parser.parse_args() + + +if __name__ == '__main__': + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + if args.model_type == 'DnCNN-S': + network = DnCNN(1, num_of_layers=17) + elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B': + network = DnCNN(1, num_of_layers=20) + else: + print("wrong model type") + exit() + + param_dict = load_checkpoint(args.ckpt_file) + load_param_into_net(network, param_dict) + + input_arr = Tensor(np.ones([args.batch_size, 1, args.image_height, args.image_width]), ms.float32) + export(network, input_arr, file_name=args.file_name, file_format=args.file_format) diff --git a/model_zoo/official/cv/dncnn/requirements.txt b/model_zoo/official/cv/dncnn/requirements.txt new file mode 100644 index 00000000000..5020bd1705d --- /dev/null +++ b/model_zoo/official/cv/dncnn/requirements.txt @@ -0,0 +1,5 @@ +mindspore-gpu==1.1.1 +numpy==1.17.0 +Pillow=6.2.2 +opencv-python==4.2.0.34 +scikit-image==0.18.1 diff --git a/model_zoo/official/cv/dncnn/scripts/run_eval_gpu.sh b/model_zoo/official/cv/dncnn/scripts/run_eval_gpu.sh new file mode 100644 index 00000000000..e722d2dc43c --- /dev/null +++ b/model_zoo/official/cv/dncnn/scripts/run_eval_gpu.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +python eval.py --dataset_path=$1 --ckpt_path=$2 --model_type=$3 --noise_level=$4 --noise_type=$5 diff --git a/model_zoo/official/cv/dncnn/scripts/run_train_gpu.sh b/model_zoo/official/cv/dncnn/scripts/run_train_gpu.sh new file mode 100644 index 00000000000..bf12092d62e --- /dev/null +++ b/model_zoo/official/cv/dncnn/scripts/run_train_gpu.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +if [ "$2" == "DnCNN-S" ] +then + python train.py --dataset_path=$1 --model_type=$2 --ckpt_prefix=$3 --noise_level=$4 +else + python train.py --dataset_path=$1 --model_type=$2 --ckpt_prefix=$3 +fi diff --git a/model_zoo/official/cv/dncnn/src/dataset.py b/model_zoo/official/cv/dncnn/src/dataset.py new file mode 100755 index 00000000000..738974cb287 --- /dev/null +++ b/model_zoo/official/cv/dncnn/src/dataset.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import os +import random +import glob +import io + +import numpy as np +import cv2 +import PIL +import mindspore +import mindspore.dataset as ds +import mindspore.dataset.transforms.c_transforms as C + +def create_train_dataset(data_path, model_type, noise_level=25, batch_size=128): + # define dataset + if model_type == "DnCNN-S": + dataset = DnCNN_train_Dataset(data_path, model_type, batch_size, noise_level, patch_shape=(40, 40)) + if model_type in ["DnCNN-B", "DnCNN-3"]: + dataset = DnCNN_train_Dataset(data_path, model_type, batch_size, patch_shape=(50, 50)) + + print("total training patch numbers per epoch", len(dataset)) + dataloader = ds.GeneratorDataset(dataset, ["noisy", "gt"]) + # apply map operations on images + dataloader = dataloader.map(input_columns="noisy", operations=C.TypeCast(mindspore.float32)) + dataloader = dataloader.map(input_columns="gt", operations=C.TypeCast(mindspore.float32)) + + # apply DatasetOps + dataloader = dataloader.shuffle(buffer_size=10000) + dataloader = dataloader.batch(batch_size, drop_remainder=True) + dataloader = dataloader.repeat(1) #here 400 images as an epoch , on the paper 128x1600 patches as a epoch + return dataloader + +class DnCNN_train_Dataset(): + def __init__(self, dataset_path, model_type, batch_size=128, noise_level=25, \ + image_shape=(180, 180), patch_shape=(50, 50)): + #DnCNN-S/B uses 200 training and 200 test images in BDS500 as training set + #DnCNN-3 uses 200 training images in BDS500 and T91 images as training set + self.im_list = [] + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.jpg"))) + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.bmp"))) + self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.png"))) + self.patch_shape = patch_shape + self.image_shape = image_shape + self.batch_size = batch_size + self.model_type = model_type + self.noise_level = noise_level + self.scales = [0.8, 1, 1.25] + self.sr_scales = [2, 3, 4] + self.compression_range = (5, 100) + self.sigma_range = (0, 55) + + def __getitem__(self, i): + i = i % len(self.im_list) + img = cv2.imread(self.im_list[i], 0) + #random scale + scale = random.choice(self.scales) + img_h = int(self.image_shape[1]*scale) + img_w = int(self.image_shape[0]*scale) + img = cv2.resize(img, (img_w, img_h)) + + #crop random patch + start_w = random.randint(0, img_w - self.patch_shape[0]) + start_h = random.randint(0, img_h - self.patch_shape[1]) + patch = img[start_h:start_h+self.patch_shape[1], start_w:start_w+self.patch_shape[0]] + + #random flip & rotation + patch = self.data_augment(patch) + + #add noise + if self.model_type == "DnCNN-S": + #add specific level of noise + noisy = self.add_noise(patch, self.noise_level) + elif self.model_type == "DnCNN-B": + #add random level of noise + sigma = random.uniform(*self.sigma_range) + noisy = self.add_noise(patch, sigma) + elif self.model_type == "DnCNN-3": + #randomly add noise, down-up sampling blur and jpeg bloc + noisy = self.add_random_type_noise(patch) + else: + print("wrong type") + exit() + noise = np.float32(noisy) - np.float32(patch) + + #add channel dimension + noisy = noisy[np.newaxis, :, :] + noise = noise[np.newaxis, :, :] + + #normalize + noisy = noisy / 255.0 + noise = noise / 255.0 + + return noisy, noise + + def __len__(self): + #To build the same epoch size as original paper + #on the paper, DnCNN-S has 1600 iteratisons per epoch + #DnCNN-B has 3000 and DnCNN-3 has 8000 + if self.model_type == "DnCNN-S": + epoch_size = self.batch_size * 1600 + elif self.model_type == "DnCNN-B": + epoch_size = self.batch_size * 3000 + elif self.model_type == "DnCNN-3": + epoch_size = self.batch_size * 8000 + else: + print("wrong model type") + exit() + return epoch_size + + def data_augment(self, patch): + #random flip & rotation + if random.random() < 0.5: + patch = np.fliplr(patch) + if random.random() < 0.5: + patch = np.flipud(patch) + alea = random.random() + if alea > 0.25: + patch = np.rot90(patch) + elif alea > 0.5: + patch = np.rot90(patch, k=2) + elif alea > 0.75: + patch = np.rot90(patch, k=3) + return patch + + def add_random_type_noise(self, patch): + #random noise/down-up resampling/JPEG compression + alea = random.random() + if alea < 0.33: #generate noisy image + sigma = random.uniform(*self.sigma_range) + noisy = self.add_noise(patch, sigma) + elif alea < 0.66: #generate blur image + sr_scale = random.choice(self.sr_scales) + noisy = cv2.resize(patch, (int(self.patch_shape[0]/sr_scale), int(self.patch_shape[1]/sr_scale))) + noisy = cv2.resize(noisy, (int(self.patch_shape[0]), int(self.patch_shape[1]))) + else: #generate JPEG blocking image + compression_level = random.randint(*self.compression_range) + noisy = self.jpeg_compression(patch, compression_level) + return noisy + + def add_noise(self, im, sigma): + gauss = np.random.normal(0, sigma, self.patch_shape) + noisy = im + gauss + noisy = np.clip(noisy, 0, 255) + noisy = noisy.astype('uint8') + return noisy + + def jpeg_compression(self, img, quality): + im_pil = PIL.Image.fromarray(img) + output = io.BytesIO() + im_pil.save(output, 'JPEG', quality=quality) + im_pil = PIL.Image.open(output) + img_np = np.asarray(im_pil) + return img_np + + +if __name__ == "__main__": + #only for test + test_dataset_path = "/code/BSR_bsds500/BSR/BSDS500/data/images/" + ds_train = create_train_dataset(test_dataset_path, "DnCNN-S", batch_size=128) + print("batch number:", ds_train.get_dataset_size()) + + for data in ds_train.create_dict_iterator(): + print(type(data)) + print(data["noisy"].shape) + print(data["gt"].shape) + print(type(data["noisy"])) + break diff --git a/model_zoo/official/cv/dncnn/src/model.py b/model_zoo/official/cv/dncnn/src/model.py new file mode 100755 index 00000000000..478c88238c6 --- /dev/null +++ b/model_zoo/official/cv/dncnn/src/model.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import numpy as np +import mindspore +import mindspore.nn as nn +from mindspore import context +from mindspore import Tensor +from mindspore.common.initializer import HeUniform + + +class DnCNN(nn.Cell): + def __init__(self, channels, num_of_layers=20): + super(DnCNN, self).__init__() + kernel_size = 3 + padding = 1 + features = 64 + layers = [] + layers.append(nn.Conv2d(channels, out_channels=features, kernel_size=kernel_size, \ + pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform())) + layers.append(nn.ReLU()) + + for _ in range(num_of_layers-2): + layers.append(nn.Conv2d(features, out_channels=features, kernel_size=kernel_size, \ + pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform())) + layers.append(nn.BatchNorm2d(features)) + layers.append(nn.ReLU()) + + layers.append(nn.Conv2d(features, out_channels=channels, kernel_size=kernel_size, \ + pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform())) + self.dncnn = nn.SequentialCell(layers) + + def construct(self, x): + out = self.dncnn(x) + return out + +if __name__ == "__main__": + #for test + context.set_context(mode=context.GRAPH_MODE, device_target='GPU') + net = DnCNN(1, num_of_layers=17) + a = Tensor(np.ones((2, 1, 40, 40)), mindspore.float32) + output = net(a) + print(output) + print(type(output)) + np_out = output.asnumpy() + print(type(np_out)) diff --git a/model_zoo/official/cv/dncnn/train.py b/model_zoo/official/cv/dncnn/train.py new file mode 100755 index 00000000000..1584ac4cabb --- /dev/null +++ b/model_zoo/official/cv/dncnn/train.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +# Copyright 2021 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +import argparse +import datetime +import mindspore.nn as nn +from mindspore import context +from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, LearningRateScheduler +from mindspore.train import Model +from mindspore.train.callback import Callback + +from src.dataset import create_train_dataset +from src.model import DnCNN + + +class BatchAverageMSELoss(nn.Cell): + def __init__(self, batch_size): + super(BatchAverageMSELoss, self).__init__() + self.batch_size = batch_size + self.sumMSELoss = nn.MSELoss(reduction='sum') + + def construct(self, logits, labels): + #equation 1 on the paper + loss = self.sumMSELoss(logits, labels) / self.batch_size / 2 + return loss + +class Print_info(Callback): + def epoch_end(self, run_context): + cb_params = run_context.original_args() + print(datetime.datetime.now(), "end epoch", cb_params.cur_epoch_num) + +def learning_rate_function(lr, cur_step_num): + if cur_step_num % 40000 == 0: + lr = lr*0.8 + print("current lr: ", str(lr)) + return lr + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="DnCNN") + parser.add_argument("--dataset_path", type=str, default="/code/BSR_bsds500/BSR/BSDS500/data/images/", \ + help='training image path') + parser.add_argument("--batch_size", type=int, default=128, help='training batch size') + parser.add_argument('--lr', type=float, default=0.001, help='learning rate') + parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight_decay') + parser.add_argument('--model_type', type=str, default='DnCNN-S', \ + choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN') + parser.add_argument('--noise_level', type=int, default=25, help="noise level only for DnCNN-S") + parser.add_argument('--ckpt_prefix', type=str, default="dncnn_mindspore", help='ckpt name prefix') + parser.add_argument('--epoch_num', type=int, default=50, help='epoch number') + args = parser.parse_args() + + context.set_context(mode=context.GRAPH_MODE, device_target="GPU") + + if args.model_type == 'DnCNN-S': + network = DnCNN(1, num_of_layers=17) + elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B': + network = DnCNN(1, num_of_layers=20) + else: + print("wrong model type") + exit() + + ds_train = create_train_dataset(args.dataset_path, args.model_type, noise_level=args.noise_level, \ + batch_size=args.batch_size) + opt = nn.AdamWeightDecay(network.trainable_params(), args.lr, weight_decay=args.weight_decay) + loss_fun = BatchAverageMSELoss(args.batch_size) + + model = Model(network, loss_fun, opt) + + #training callbacks + checkpoint_config = CheckpointConfig(save_checkpoint_steps=1000, keep_checkpoint_max=3) + ckpoint_cb = ModelCheckpoint(prefix=args.ckpt_prefix, directory='./ckpt/', config=checkpoint_config) + print_cb = Print_info() + lr_cb = LearningRateScheduler(learning_rate_function) + loss_monitor_cb = LossMonitor(per_print_times=100) + + print(datetime.datetime.now(), " training starts") + model.train(args.epoch_num, ds_train, callbacks=[lr_cb, ckpoint_cb, print_cb, loss_monitor_cb], \ + dataset_sink_mode=False)