diff --git a/model_zoo/research/cv/Pix2Pix/README.md b/model_zoo/research/cv/Pix2Pix/README.md new file mode 100644 index 00000000000..d6758e026db --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/README.md @@ -0,0 +1,219 @@ +# Contents + +- [Pix2Pix Description](#Pix2Pix-description) +- [Model Architecture](#model-architecture) +- [Dataset](#dataset) +- [Environment Requirements](#environment-requirements) +- [Script Description](#script-description) + - [Script and Sample Code](#script-and-sample-code) + - [Script Parameters](#script-parameters) + - [Training](#training-process) + - [Evaluation](#evaluation-process) + - [Prediction Process](#prediction-process) +- [Model Description](#model-description) + - [Performance](#performance) + - [Training Performance](#evaluation-performance) + - [Evaluation Performance](#evaluation-performance) +- [ModelZoo Homepage](#modelzoo-homepage) + +# [Pix2Pix Description](#contents) + +Many problems in image processing, computer graphics, and computer vision can be posed as “translating” an input image into a corresponding output image, each of these tasks has been tackled with separate, special-purpose machinery, despite the fact that the setting is always the same: predict pixels from pixels. +Our goal in this paper is to develop a common framework for all these problems. Pix2pix model is a conditional GAN, which includes two modules--generator and discriminator. This model transforms an input image into a corresponding output image. The essence of the model is the mapping from pixel to pixel. + +[Paper](https://arxiv.org/abs/1611.07004): Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros. "Image-to-Image Translation with Conditional Adversarial Networks", in CVPR 2017. + +![Pix2Pix Imgs](imgs/Pix2Pix-examples.jpg) + +# [Model Architecture](#contents) + +The Pix2Pix contains a generation network and a discriminant networks.In the generator part, the model can be any pixel to pixel mapping network (in the raw paper, the author proposed to use Unet). In the discriminator part, a patch GAN is used to judge whether each N*N patches is fake or true, thus can improve the reality of the generated image. + +**Generator(Unet-Based) architectures:** + +Encoder: + +C64-C128-C256-C512-C512-C512-C512-C512 + +Decoder: + +CD512-CD1024-CD1024-C1024-C1024-C512-C256-C128 + +**Discriminator(70 × 70 discriminator) architectures:** + +C64-C128-C256-C512 + +**Note:** Let Ck denote a Convolution-BatchNorm-ReLU layer with k filters. CDk denotes a Convolution-BatchNorm-Dropout-ReLU layer with a dropout rate of 50%. + +# [Dataset](#contents) + +Dataset_1 used: [facades](http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/facades.tar.gz) + +```markdown + Dataset size: 29M, 606 images + 400 train images + 100 validation images + 106 test images + Data format：.jpg images +``` + +Dataset_2 used: [maps](http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/maps.tar.gz) + +```markdown + Dataset size: 239M, 2194 images + 1096 train images + 1098 validation images + Data format：.jpg images +``` + +**Note:** We provide data/download_Pix2Pix_dataset.sh to download the datasets. + +# [Environment Requirements](#contents) + +- Hardware（Ascend） + - Prepare hardware environment with Ascend processor. +- Framework + - [MindSpore](https://www.mindspore.cn/install/en) +- For more information, please check the resources below： + - [MindSpore tutorials](https://www.mindspore.cn/tutorial/training/en/master/index.html) + - [MindSpore Python API](https://www.mindspore.cn/doc/api_python/en/master/index.html) + +## [Dependences](#contents) + +- Python==3.8.5 +- Mindspore==1.2 + +# [Script Description](#contents) + +## [Script and Sample Code](#contents) + +The entire code structure is as following: + +```markdown +.Pix2Pix +├─ README.md # descriptions about Pix2Pix +├─ data + └─download_Pix2Pix_dataset.sh # download dataset +├── scripts + └─run_train_ascend.sh # launch ascend training(1 pcs) + └─run_eval_ascend.sh # launch ascend eval +├─ imgs + └─Pix2Pix-examples.jpg # Pix2Pix Imgs +├─ src + ├─ __init__.py # init file + ├─ dataset + ├─ __init__.py # init file + ├─ pix2pix_dataset.py # create pix2pix dataset + ├─ models + ├─ __init__.py # init file + ├─ discriminator_model.py # define discriminator model——Patch GAN + ├─ generator_model.py # define generator model——Unet-based Generator + ├─ init_w.py # initialize network weights + ├─ loss.py # define losses + └─ pix2pix.py # define Pix2Pix model + └─ utils + ├─ __init__.py # init file + ├─ config.py # parse args + ├─ tools.py # tools for Pix2Pix model +├─ eval.py # evaluate Pix2Pix Model +├─ train.py # train script +└─ export.py # export mindir script +``` + +## [Script Parameters](#contents) + +Major parameters in train.py and config.py as follows: + +```python +"device_target": Ascend # run platform, only support Ascend. +"device_num": 1 # device num, default is 1. +"device_id": 0 # device id, default is 0. +"save_graphs": False # whether save graphs, default is False. +"init_type": normal # network initialization, default is normal. +"init_gain": 0.02 # scaling factor for normal, xavier and orthogonal, default is 0.02. +"load_size": 286 # scale images to this size, default is 286. +"batch_size": 1 # batch_size, default is 1. +"LAMBDA_Dis": 0.5 # weight for Discriminator Loss, default is 0.5. +"LAMBDA_GAN": 1 # weight for GAN Loss, default is 1. +"LAMBDA_L1": 100 # weight for L1 Loss, default is 100. +"beta1": 0.5 # adam beta1, default is 0.5. +"beta2": 0.999 # adam beta2, default is 0.999. +"lr": 0.0002 # the initial learning rate, default is 0.0002. +"lr_policy": linear # learning rate policy, default is linear. +"epoch_num": 200 # epoch number for training, default is 200. +"n_epochs": 100 # number of epochs with the initial learning rate, default is 100. +"n_epochs_decay": 100 # number of epochs with the dynamic learning rate, default is 100. +"dataset_size": 400 # for Facade_dataset,the number is 400; for Maps_dataset,the number is 1096. +"train_data_dir": None # the file path of input data during training. +"val_data_dir": None # the file path of input data during validating. +"train_fakeimg_dir": ./results/fake_img/ # during training, the file path of stored fake img. +"loss_show_dir": ./results/loss_show # during training, the file path of stored loss img. +"ckpt_dir": ./results/ckpt # during training, the file path of stored CKPT. +"ckpt": None # during validating, the file path of the CKPT used. +"predict_dir": ./results/predict/ # during validating, the file path of Generated images. +``` + +## [Training](#contents) + +- running on Ascend with default parameters + +```python +python train.py --device_target [Ascend] --device_id [0] --train_data_dir [./data/facades/train] +``` + +## [Evaluation](#contents) + +```python +python eval.py --device_target [Ascend] --device_id [0] --val_data_dir [./data/facades/test] --ckpt [./results/ckpt/Generator_200.ckpt] +``` + +**Note:**: Before training and evaluating, create folders like "./results/...". Then you will get the results as following in "./results/predict". + +# [Model Description](#contents) + +## [Performance](#contents) + +### Training Performance + +| Parameters | single Ascend | +| -------------------------- | ----------------------------------------------------------- | +| Model Version | Pix2Pix | +| Resource | Ascend 910 | +| MindSpore Version | 1.2 | +| Dataset | facades | +| Training Parameters | epoch=200, steps=400, batch_size=1, lr=0.0002 | +| Optimizer | Adam | +| Loss Function | SigmoidCrossEntropyWithLogits Loss & L1 Loss | +| outputs | probability | +| Speed | 1pc(Ascend): 10 ms/step | +| Total time | 1pc(Ascend): 0.3h | +| Checkpoint for Fine tuning | 207M (.ckpt file) | + +| Parameters | single Ascend | +| -------------------------- | ----------------------------------------------------------- | +| Model Version | Pix2Pix | +| Resource | Ascend 910 | +| MindSpore Version | 1.2 | +| Dataset | maps | +| Training Parameters | epoch=200, steps=1096, batch_size=1, lr=0.0002 | +| Optimizer | Adam | +| Loss Function | SigmoidCrossEntropyWithLogits Loss & L1 Loss | +| outputs | probability | +| Speed | 1pc(Ascend): 20 ms/step | +| Total time | 1pc(Ascend): 1.58h | +| Checkpoint for Fine tuning | 207M (.ckpt file) | + +### Evaluation Performance + +| Parameters | single Ascend | +| ------------------- | --------------------------- | +| Model Version | Pix2Pix | +| Resource | Ascend 910 | +| MindSpore Version | 1.2 | +| Dataset | facades / maps | +| batch_size | 1 | +| outputs | probability | + +# [ModelZoo Homepage](#contents) + +Please check the official [homepage](https://gitee.com/mindspore/mindspore/tree/master/model_zoo). diff --git a/model_zoo/research/cv/Pix2Pix/data/download_Pix2Pix_dataset.sh b/model_zoo/research/cv/Pix2Pix/data/download_Pix2Pix_dataset.sh new file mode 100644 index 00000000000..6b38aa502e1 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/data/download_Pix2Pix_dataset.sh @@ -0,0 +1,36 @@ +#!/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. + +FILE=$1 + +if [[ $FILE != "maps" && $FILE != "facades" ]]; then + echo "Available datasets are: maps, cityscapes, facades" + exit 1 +fi + +if [[ $FILE == "cityscapes" ]]; then + echo "Due to license issue, we cannot provide the Cityscapes dataset from our repository. Please download the Cityscapes dataset from https://cityscapes-dataset.com, and use the script ./datasets/prepare_cityscapes_dataset.py." + echo "You need to download gtFine_trainvaltest.zip and leftImg8bit_trainvaltest.zip. For further instruction, please read ./datasets/prepare_cityscapes_dataset.py" + exit 1 +fi + +echo "Specified [$FILE]" +URL=http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/$FILE.tar.gz +ZIP_FILE=./datasets/$FILE.tar.gz +TARGET_DIR=./datasets/$FILE/ +wget -N $URL -O $ZIP_FILE +mkdir $TARGET_DIR +tar -zxvf $ZIP_FILE +rm $ZIP_FILE diff --git a/model_zoo/research/cv/Pix2Pix/eval.py b/model_zoo/research/cv/Pix2Pix/eval.py new file mode 100644 index 00000000000..40477aef26b --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/eval.py @@ -0,0 +1,74 @@ +# 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. +# =========================================================================== + +""" + Evaluate Pix2Pix Model. +""" + +from mindspore import Tensor, nn +from mindspore.train.serialization import load_checkpoint +from mindspore.train.serialization import load_param_into_net +from src.dataset.pix2pix_dataset import pix2pixDataset_val, create_val_dataset +from src.models.pix2pix import Pix2Pix, get_generator, get_discriminator +from src.models.loss import D_Loss, D_WithLossCell, G_Loss, G_WithLossCell, TrainOneStepCell +from src.utils.tools import save_image, get_lr +from src.utils.config import get_args + +if __name__ == '__main__': + + args = get_args() + + # Preprocess the data for evaluating + dataset_val = pix2pixDataset_val(root_dir=args.val_data_dir) + ds_val = create_val_dataset(dataset_val) + print("ds:", ds_val.get_dataset_size()) + print("ds:", ds_val.get_col_names()) + print("ds.shape:", ds_val.output_shapes()) + + steps_per_epoch = ds_val.get_dataset_size() + + netG = get_generator() + netD = get_discriminator() + + pix2pix = Pix2Pix(generator=netG, discriminator=netD) + + d_loss_fn = D_Loss() + g_loss_fn = G_Loss() + d_loss_net = D_WithLossCell(backbone=pix2pix, loss_fn=d_loss_fn) + g_loss_net = G_WithLossCell(backbone=pix2pix, loss_fn=g_loss_fn) + + d_opt = nn.Adam(pix2pix.netD.trainable_params(), learning_rate=get_lr(), + beta1=args.beta1, beta2=args.beta2, loss_scale=1) + g_opt = nn.Adam(pix2pix.netG.trainable_params(), learning_rate=get_lr(), + beta1=args.beta1, beta2=args.beta2, loss_scale=1) + + train_net = TrainOneStepCell(loss_netD=d_loss_net, loss_netG=g_loss_net, optimizerD=d_opt, optimizerG=g_opt, sens=1) + train_net.set_train() + + # Evaluating loop + ckpt_url = args.ckpt + print("CKPT:", ckpt_url) + param_G = load_checkpoint(ckpt_url) + load_param_into_net(netG, param_G) + + data_loader_val = ds_val.create_dict_iterator(output_numpy=True, num_epochs=args.epoch_num) + print("=======Starting evaluating Loop=======") + for i, data in enumerate(data_loader_val): + input_image = Tensor(data["input_images"]) + target_image = Tensor(data["target_images"]) + + fake_image = netG(input_image) + save_image(fake_image, args.predict_dir + str(i + 1)) + print("=======image", i + 1, "saved success=======") diff --git a/model_zoo/research/cv/Pix2Pix/export.py b/model_zoo/research/cv/Pix2Pix/export.py new file mode 100644 index 00000000000..bd6650dd8a8 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/export.py @@ -0,0 +1,75 @@ +# 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. +# ============================================================================ + +""" + export checkpoint file into air, onnx, mindir models +""" +import argparse +import ast +import numpy as np +from mindspore import Tensor, nn, context +from mindspore.train.serialization import export +from mindspore.train.serialization import load_checkpoint +from mindspore.train.serialization import load_param_into_net +from src.models.pix2pix import Pix2Pix, get_generator, get_discriminator +from src.models.loss import D_Loss, D_WithLossCell, G_Loss, G_WithLossCell, TrainOneStepCell +from src.utils.tools import get_lr + +parser = argparse.ArgumentParser(description='export') +parser.add_argument("--run_modelart", type=ast.literal_eval, default=False, help="Run on modelArt, default is false.") +parser.add_argument("--device_id", type=int, default=0, help="device id, default is 0.") +parser.add_argument("--batch_size", type=int, default=1, help="batch_size, default is 1.") +parser.add_argument("--image_size", type=int, default=256, help="images size, default is 256.") +parser.add_argument('--ckpt_dir', type=str, default='./results/ckpt', + help='during training, the file path of stored CKPT.') +parser.add_argument("--ckpt", type=str, default=None, help="during validating, the file path of the CKPT used.") +parser.add_argument('--train_data_dir', type=str, default=None, help='the file path of input data during training.') +parser.add_argument("--file_name", type=str, default="Pix2Pix", help="output file name.") +parser.add_argument("--file_format", type=str, default="AIR", choices=["AIR", "ONNX", "MINDIR"], help="file format") +parser.add_argument('--device_target', type=str, default='Ascend', choices=('Ascend', 'GPU'), + help='device where the code will be implemented (default: Ascend)') +args = parser.parse_args() + +context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) + +if __name__ == '__main__': + netG = get_generator() + netD = get_discriminator() + + pix2pix = Pix2Pix(generator=netG, discriminator=netD) + + d_loss_fn = D_Loss() + g_loss_fn = G_Loss() + d_loss_net = D_WithLossCell(backbone=pix2pix, loss_fn=d_loss_fn) + g_loss_net = G_WithLossCell(backbone=pix2pix, loss_fn=g_loss_fn) + + d_opt = nn.Adam(pix2pix.netD.trainable_params(), learning_rate=get_lr(), beta1=0.5, beta2=0.999, loss_scale=1) + g_opt = nn.Adam(pix2pix.netG.trainable_params(), learning_rate=get_lr(), beta1=0.5, beta2=0.999, loss_scale=1) + + train_net = TrainOneStepCell(loss_netD=d_loss_net, loss_netG=g_loss_net, optimizerD=d_opt, optimizerG=g_opt, sens=1) + train_net.set_train() + train_net = train_net.loss_netG + + ckpt_url = args.ckpt + param_G = load_checkpoint(ckpt_url) + load_param_into_net(netG, param_G) + + input_shp = [args.batch_size, 3, args.image_size, args.image_size] + input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shp).astype(np.float32)) + target_shp = [args.batch_size, 3, args.image_size, args.image_size] + target_array = Tensor(np.random.uniform(-1.0, 1.0, size=target_shp).astype(np.float32)) + inputs = [input_array, target_array] + file = f"{args.file_name}" + export(train_net, *inputs, file_name=file, file_format=args.file_format) diff --git a/model_zoo/research/cv/Pix2Pix/imgs/Pix2Pix-examples.jpg b/model_zoo/research/cv/Pix2Pix/imgs/Pix2Pix-examples.jpg new file mode 100644 index 00000000000..b1f24d5ef63 Binary files /dev/null and b/model_zoo/research/cv/Pix2Pix/imgs/Pix2Pix-examples.jpg differ diff --git a/model_zoo/research/cv/Pix2Pix/scripts/run_eval_ascend.sh b/model_zoo/research/cv/Pix2Pix/scripts/run_eval_ascend.sh new file mode 100644 index 00000000000..b0614d4c174 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/scripts/run_eval_ascend.sh @@ -0,0 +1,22 @@ +#!/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. + +echo "=======================================================================================================================================" +echo "Please run the eval as: " +echo "python eval.py device_target device_id val_data_dir ckpt" +echo "for example: python eval.py --device_target Ascend --device_id 0 --val_data_dir ./facades/test --ckpt ./results/ckpt/Generator_200.ckpt" +echo "=======================================================================================================================================" + +python eval.py --device_target Ascend --device_id 0 --val_data_dir ./facades/test --ckpt ./results/ckpt/Generator_200.ckpt diff --git a/model_zoo/research/cv/Pix2Pix/scripts/run_train_ascend.sh b/model_zoo/research/cv/Pix2Pix/scripts/run_train_ascend.sh new file mode 100644 index 00000000000..34415c02db4 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/scripts/run_train_ascend.sh @@ -0,0 +1,22 @@ +#!/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. + +echo "=====================================================================================================================" +echo "Please run the train as: " +echo "python train.py device_target device_id dataset_size train_data_dir" +echo "for example: python train.py --device_target Ascend --device_id 0 --dataset_size 400 --train_data_dir ./facades/train" +echo "=====================================================================================================================" + +python train.py --device_target Ascend --device_id 0 --dataset_size 400 --train_data_dir ./facades/train diff --git a/model_zoo/research/cv/Pix2Pix/src/__init__.py b/model_zoo/research/cv/Pix2Pix/src/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/research/cv/Pix2Pix/src/dataset/__init__.py b/model_zoo/research/cv/Pix2Pix/src/dataset/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/research/cv/Pix2Pix/src/dataset/pix2pix_dataset.py b/model_zoo/research/cv/Pix2Pix/src/dataset/pix2pix_dataset.py new file mode 100644 index 00000000000..382fc114f03 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/dataset/pix2pix_dataset.py @@ -0,0 +1,173 @@ +# 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. +# =========================================================================== + +''' + Preprocess Pix2Pix dataset. +''' + +import os +import numpy as np +from PIL import Image +import mindspore +from mindspore import dataset as de +import mindspore.dataset.vision.c_transforms as C +from ..utils.config import get_args + +args = get_args() + +class pix2pixDataset(): + ''' + Define train dataset. + ''' + def __init__(self, root_dir): + self.root_dir = root_dir + self.list_files = os.listdir(self.root_dir) + self.list_files.sort(key=lambda x: int(x[:-4])) + print(self.list_files) + + def __len__(self): + return len(self.list_files) + + def __getitem__(self, index): + img_file = self.list_files[index] + img_path = os.path.join(self.root_dir, img_file) + AB = Image.open(img_path).convert('RGB') + w, h = AB.size + w2 = int(w / 2) + + A = AB.crop((w2, 0, w, h)) + B = AB.crop((0, 0, w2, h)) + + A = A.resize((args.load_size, args.load_size)) + B = B.resize((args.load_size, args.load_size)) + + transform_params = get_params(A.size) + A_crop = crop(A, transform_params, size=256) + B_crop = crop(B, transform_params, size=256) + + return A_crop, B_crop + +def get_params(size=(256, 256)): + ''' + Get parameters from images. + ''' + w, h = size + new_h = h + new_w = w + new_h = new_w = args.load_size # args.load_size + + x = np.random.randint(0, np.maximum(0, new_w - 256)) + y = np.random.randint(0, np.maximum(0, new_h - 256)) + + return (x, y) + +def crop(img, pos, size=256): + ''' + Crop the images. + ''' + ow = oh = args.load_size + x1, y1 = pos + tw = th = size + if (ow > tw or oh > th): + img = img.crop((x1, y1, x1 + tw, y1 + th)) + return img + return img + +def sync_random_Horizontal_Flip(input_images, target_images): + ''' + Randomly flip the input images and the target images. + ''' + seed = np.random.randint(0, 2000000000) + mindspore.set_seed(seed) + op = C.RandomHorizontalFlip(prob=0.5) + out_input = op(input_images) + mindspore.set_seed(seed) + op = C.RandomHorizontalFlip(prob=0.5) + out_target = op(target_images) + return out_input, out_target + +def create_train_dataset(dataset): + ''' + Create train dataset. + ''' + + mean = [0.5 * 255] * 3 + std = [0.5 * 255] * 3 + + trans = [ + C.Normalize(mean=mean, std=std), + C.HWC2CHW() + ] + + train_ds = de.GeneratorDataset(dataset, column_names=["input_images", "target_images"], shuffle=False) + + train_ds = train_ds.map(operations=[sync_random_Horizontal_Flip], input_columns=["input_images", "target_images"]) + + train_ds = train_ds.map(operations=trans, input_columns=["input_images"]) + train_ds = train_ds.map(operations=trans, input_columns=["target_images"]) + + train_ds = train_ds.batch(1, drop_remainder=True) + train_ds = train_ds.repeat(1) + + return train_ds + +class pix2pixDataset_val(): + ''' + Define val dataset. + ''' + def __init__(self, root_dir): + self.root_dir = root_dir + self.list_files = os.listdir(self.root_dir) + self.list_files.sort(key=lambda x: int(x[:-4])) + print(self.list_files) + + def __len__(self): + return len(self.list_files) + + def __getitem__(self, index): + img_file = self.list_files[index] + img_path = os.path.join(self.root_dir, img_file) + + AB = Image.open(img_path).convert('RGB') + w, h = AB.size + + w2 = int(w / 2) + A = AB.crop((w2, 0, w, h)) + B = AB.crop((0, 0, w2, h)) + + return A, B + +def create_val_dataset(dataset): + ''' + Create val dataset. + ''' + + mean = [0.5 * 255] * 3 + std = [0.5 * 255] * 3 + + trans = [ + C.Resize((256, 256)), + C.Normalize(mean=mean, std=std), + C.HWC2CHW() + ] + + val_ds = de.GeneratorDataset(dataset, column_names=["input_images", "target_images"], shuffle=False) + + val_ds = val_ds.map(operations=trans, input_columns=["input_images"]) + val_ds = val_ds.map(operations=trans, input_columns=["target_images"]) + val_ds = val_ds.batch(1, drop_remainder=True) + val_ds = val_ds.repeat(1) + + return val_ds diff --git a/model_zoo/research/cv/Pix2Pix/src/models/__init__.py b/model_zoo/research/cv/Pix2Pix/src/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/research/cv/Pix2Pix/src/models/discriminator_model.py b/model_zoo/research/cv/Pix2Pix/src/models/discriminator_model.py new file mode 100644 index 00000000000..b7dc62ba81c --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/models/discriminator_model.py @@ -0,0 +1,118 @@ +# 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. +# =========================================================================== + +''' + Define discriminator model——Patch GAN. +''' + +import mindspore.nn as nn +from mindspore.ops import Concat + +class ConvNormReLU(nn.Cell): + """ + Convolution fused with BatchNorm/InstanceNorm and ReLU/LackyReLU block definition. + + Args: + in_planes (int): Input channel. + out_planes (int): Output channel. + kernel_size (int): Input kernel size. Default: 4. + stride (int): Stride size for the first convolutional layer. Default: 2. + alpha (float): Slope of LackyReLU. Default: 0.2. + norm_mode (str): Specifies norm method. The optional values are "batch", "instance". + pad_mode (str): Specifies padding mode. The optional values are "CONSTANT", "REFLECT", "SYMMETRIC". + Default: "CONSTANT". + use_relu (bool): Use relu or not. Default: True. + padding (int): Pad size, if it is None, it will calculate by kernel_size. Default: None. + + Returns: + Tensor, output tensor. + """ + def __init__(self, + in_planes, + out_planes, + kernel_size=4, + stride=2, + alpha=0.2, + norm_mode='batch', + pad_mode='CONSTANT', + use_relu=True, + padding=None): + super(ConvNormReLU, self).__init__() + norm = nn.BatchNorm2d(out_planes) + if norm_mode == 'instance': + # Use BatchNorm2d with batchsize=1, affine=False, training=True instead of InstanceNorm2d + norm = nn.BatchNorm2d(out_planes, affine=False) + has_bias = (norm_mode == 'instance') + if padding is None: + padding = (kernel_size - 1) // 2 + if pad_mode == 'CONSTANT': + conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, pad_mode='pad', + has_bias=has_bias, padding=padding) + layers = [conv, norm] + else: + paddings = ((0, 0), (0, 0), (padding, padding), (padding, padding)) + pad = nn.Pad(paddings=paddings, mode=pad_mode) + conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, pad_mode='pad', has_bias=has_bias) + layers = [pad, conv, norm] + if use_relu: + relu = nn.ReLU() + if alpha > 0: + relu = nn.LeakyReLU(alpha) + layers.append(relu) + self.features = nn.SequentialCell(layers) + + def construct(self, x): + output = self.features(x) + return output + + +class Discriminator(nn.Cell): + """ + Discriminator of Model. + + Args: + in_planes (int): Input channel. + ndf (int): the number of filters in the last conv layer + n_layers (int): The number of ConvNormReLU blocks. + alpha (float): LeakyRelu slope. Default: 0.2. + norm_mode (str): Specifies norm method. The optional values are "batch", "instance". + + Returns: + Tensor, output tensor. + """ + def __init__(self, in_planes=3, ndf=64, n_layers=3, alpha=0.2, norm_mode='batch'): + super(Discriminator, self).__init__() + kernel_size = 4 + layers = [ + nn.Conv2d(in_planes, ndf, kernel_size, 2, pad_mode='pad', padding=1), + nn.LeakyReLU(alpha) + ] + nf_mult = ndf + for i in range(1, n_layers): + nf_mult_prev = nf_mult + nf_mult = min(2 ** i, 8) * ndf + layers.append(ConvNormReLU(nf_mult_prev, nf_mult, kernel_size, 2, alpha, norm_mode, padding=1)) + nf_mult_prev = nf_mult + nf_mult = min(2 ** n_layers, 8) * ndf + layers.append(ConvNormReLU(nf_mult_prev, nf_mult, kernel_size, 1, alpha, norm_mode, padding=1)) + layers.append(nn.Conv2d(nf_mult, 1, kernel_size, 1, pad_mode='pad', padding=1)) + + self.features = nn.SequentialCell(layers) + self.concat = Concat(axis=1) + + def construct(self, x, y): + x_y = self.concat((x, y)) + output = self.features(x_y) + return output diff --git a/model_zoo/research/cv/Pix2Pix/src/models/generator_model.py b/model_zoo/research/cv/Pix2Pix/src/models/generator_model.py new file mode 100644 index 00000000000..728fa49e690 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/models/generator_model.py @@ -0,0 +1,127 @@ +# 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. +# =========================================================================== + +''' + Define generator model——Unet-based Generator. +''' + +import mindspore.nn as nn +import mindspore.ops as ops + +class UnetGenerator(nn.Cell): + """ + Unet-based generator. + + Args: + in_planes (int): the number of channels in input images. + out_planes (int): the number of channels in output images. + ngf (int): the number of filters in the last conv layer. + n_layers (int): the number of downsamplings in UNet. + alpha (float): LeakyRelu slope. Default: 0.2. + norm_mode (str): Specifies norm method. The optional values are "batch", "instance". + dropout (bool): Use dropout or not. Default: False. + + Returns: + Tensor, output tensor. + """ + def __init__(self, in_planes, out_planes, ngf=64, n_layers=8, alpha=0.2, norm_mode='bn', dropout=False): + super(UnetGenerator, self).__init__() + # construct unet structure + unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, in_planes=None, submodule=None, + norm_mode=norm_mode, innermost=True) + for _ in range(n_layers - 5): + unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, in_planes=None, submodule=unet_block, + norm_mode=norm_mode, dropout=dropout) + # gradually reduce the number of filters from ngf * 8 to ngf + unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, in_planes=None, submodule=unet_block, + norm_mode=norm_mode) + unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, in_planes=None, submodule=unet_block, + norm_mode=norm_mode) + unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, in_planes=None, submodule=unet_block, norm_mode=norm_mode) + self.model = UnetSkipConnectionBlock(out_planes, ngf, in_planes=in_planes, submodule=unet_block, + outermost=True, norm_mode=norm_mode) + + def construct(self, x): + return self.model(x) + + +class UnetSkipConnectionBlock(nn.Cell): + """Unet submodule with skip connection. + + Args: + outer_nc (int): The number of filters in the outer conv layer + inner_nc (int): The number of filters in the inner conv layer + in_planes (int): The number of channels in input images/features + dropout (bool): Use dropout or not. Default: False. + submodule (Cell): Previously defined submodules + outermost (bool): If this module is the outermost module + innermost (bool): If this module is the innermost module + alpha (float): LeakyRelu slope. Default: 0.2. + norm_mode (str): Specifies norm method. The optional values are "batch", "instance". + + Returns: + Tensor, output tensor. + """ + def __init__(self, outer_nc, inner_nc, in_planes=None, dropout=False, + submodule=None, outermost=False, innermost=False, alpha=0.2, norm_mode='batch'): + super(UnetSkipConnectionBlock, self).__init__() + downnorm = nn.BatchNorm2d(inner_nc) + upnorm = nn.BatchNorm2d(outer_nc) + use_bias = False + if norm_mode == 'instance': + downnorm = nn.BatchNorm2d(inner_nc, affine=False) + upnorm = nn.BatchNorm2d(outer_nc, affine=False) + use_bias = True + if in_planes is None: + in_planes = outer_nc + downconv = nn.Conv2d(in_planes, inner_nc, kernel_size=4, + stride=2, padding=1, has_bias=use_bias, pad_mode='pad') + downrelu = nn.LeakyReLU(alpha) + uprelu = nn.ReLU() + + if outermost: + upconv = nn.Conv2dTranspose(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1, pad_mode='pad') + down = [downconv] + up = [uprelu, upconv, nn.Tanh()] + model = down + [submodule] + up + elif innermost: + upconv = nn.Conv2dTranspose(inner_nc, outer_nc, + kernel_size=4, stride=2, + padding=1, has_bias=use_bias, pad_mode='pad') + down = [downrelu, downconv] + up = [uprelu, upconv, upnorm] + model = down + up + else: + upconv = nn.Conv2dTranspose(inner_nc * 2, outer_nc, + kernel_size=4, stride=2, + padding=1, has_bias=use_bias, pad_mode='pad') + down = [downrelu, downconv, downnorm] + up = [uprelu, upconv, upnorm] + + model = down + [submodule] + up + if dropout: + model.append(nn.Dropout(0.5)) + + self.model = nn.SequentialCell(model) + self.skip_connections = not outermost + self.concat = ops.Concat(axis=1) + + def construct(self, x): + out = self.model(x) + if self.skip_connections: + out = self.concat((out, x)) + return out diff --git a/model_zoo/research/cv/Pix2Pix/src/models/init_w.py b/model_zoo/research/cv/Pix2Pix/src/models/init_w.py new file mode 100644 index 00000000000..f05e4ba5235 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/models/init_w.py @@ -0,0 +1,43 @@ +# 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. +# =========================================================================== + +""" + Initialize network weights. +""" + +import mindspore.nn as nn +from mindspore.common import initializer as init + +def init_weights(net, init_type='normal', init_gain=0.02): + """ + Initialize network weights. + Parameters: + net (Cell): Network to be initialized + init_type (str): The name of an initialization method: normal | xavier. + init_gain (float): Gain factor for normal and xavier. + """ + for _, cell in net.cells_and_names(): + if isinstance(cell, (nn.Conv2d, nn.Conv2dTranspose)): + if init_type == 'normal': + cell.weight.set_data(init.initializer(init.Normal(init_gain), cell.weight.shape)) + elif init_type == 'xavier': + cell.weight.set_data(init.initializer(init.XavierUniform(init_gain), cell.weight.shape)) + elif init_type == 'constant': + cell.weight.set_data(init.initializer(0.001, cell.weight.shape)) + else: + raise NotImplementedError('initialization method [%s] is not implemented' % init_type) + elif isinstance(cell, nn.BatchNorm2d): + cell.gamma.set_data(init.initializer('ones', cell.gamma.shape)) + cell.beta.set_data(init.initializer('zeros', cell.beta.shape)) diff --git a/model_zoo/research/cv/Pix2Pix/src/models/loss.py b/model_zoo/research/cv/Pix2Pix/src/models/loss.py new file mode 100644 index 00000000000..925792f3348 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/models/loss.py @@ -0,0 +1,158 @@ +# 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. +# =========================================================================== + +""" + Define losses. +""" + +import mindspore.nn as nn +import mindspore.ops as ops +from mindspore.ops import functional as F +import mindspore.ops.operations as P +from mindspore.parallel._utils import (_get_device_num, _get_gradients_mean, _get_parallel_mode) +from mindspore.context import ParallelMode +from mindspore.nn.wrap.grad_reducer import DistributedGradReducer +from mindspore.nn.loss.loss import _Loss +from src.utils.config import get_args + +args = get_args() + +class SigmoidCrossEntropyWithLogits(_Loss): + def __init__(self): + super(SigmoidCrossEntropyWithLogits, self).__init__() + self.cross_entropy = P.SigmoidCrossEntropyWithLogits() + + def construct(self, data, label): + x = self.cross_entropy(data, label) + return self.get_loss(x) + +class D_Loss(_Loss): + """ + Define Dloss. + """ + def __init__(self, reduction="mean"): + super(D_Loss, self).__init__(reduction) + self.sig = SigmoidCrossEntropyWithLogits() + self.ones = ops.OnesLike() + self.zeros = ops.ZerosLike() + self.LAMBDA_Dis = args.LAMBDA_Dis + + def construct(self, pred1, pred0): + loss = self.sig(pred1, self.ones(pred1)) + self.sig(pred0, self.zeros(pred0)) + dis_loss = loss * self.LAMBDA_Dis + return dis_loss + +class D_WithLossCell(nn.Cell): + """ + Define D_WithLossCell. + """ + def __init__(self, backbone, loss_fn): + super(D_WithLossCell, self).__init__(auto_prefix=True) + self.netD = backbone.netD + self.netG = backbone.netG + self._loss_fn = loss_fn + + def construct(self, realA, realB): + fakeB = self.netG(realA) + pred1 = self.netD(realA, realB) + pred0 = self.netD(realA, fakeB) + return self._loss_fn(pred1, pred0) + +class G_Loss(_Loss): + """ + Define Gloss. + """ + def __init__(self, reduction="mean"): + super(G_Loss, self).__init__(reduction) + self.sig = SigmoidCrossEntropyWithLogits() + self.l1_loss = nn.L1Loss() + self.ones = ops.OnesLike() + self.LAMBDA_GAN = args.LAMBDA_GAN + self.LAMBDA_L1 = args.LAMBDA_L1 + + def construct(self, fakeB, realB, pred0): + loss_1 = self.sig(pred0, self.ones(pred0)) + loss_2 = self.l1_loss(fakeB, realB) + loss = loss_1 * self.LAMBDA_GAN + loss_2 * self.LAMBDA_L1 + return loss + +class G_WithLossCell(nn.Cell): + """ + Define G_WithLossCell. + """ + def __init__(self, backbone, loss_fn): + super(G_WithLossCell, self).__init__(auto_prefix=True) + self.netD = backbone.netD + self.netG = backbone.netG + self._loss_fn = loss_fn + + def construct(self, realA, realB): + fakeB = self.netG(realA) + pred0 = self.netD(realA, fakeB) + return self._loss_fn(fakeB, realB, pred0) + +class TrainOneStepCell(nn.Cell): + """ + Define TrainOneStepCell. + """ + def __init__(self, loss_netD, loss_netG, optimizerD, optimizerG, sens=1, auto_prefix=True): + super(TrainOneStepCell, self).__init__(auto_prefix=auto_prefix) + self.loss_netD = loss_netD # loss network + self.loss_netD.set_grad() + self.loss_netD.add_flags(defer_inline=True) + + self.loss_netG = loss_netG + self.loss_netG.set_grad() + self.loss_netG.add_flags(defer_inline=True) + + self.weights_G = optimizerG.parameters + self.optimizerG = optimizerG + self.weights_D = optimizerD.parameters + self.optimizerD = optimizerD + + self.grad = ops.GradOperation(get_by_list=True, sens_param=True) + self.sens = sens + + # 并行处理的定义 + self.reducer_flag = False + self.grad_reducer_G = F.identity + self.grad_reducer_D = F.identity + self.parallel_mode = _get_parallel_mode() + if self.parallel_mode in (ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL): + self.reducer_flag = True + if self.reducer_flag: + mean = _get_gradients_mean() + degree = _get_device_num() + self.grad_reducer_G = DistributedGradReducer(self.weights_G, mean, degree) + self.grad_reducer_D = DistributedGradReducer(self.weights_D, mean, degree) + + def set_sens(self, value): + self.sens = value + + def construct(self, realA, realB): + """ + Define TrainOneStepCell. + """ + d_loss = self.loss_netD(realA, realB) + g_loss = self.loss_netG(realA, realB) + + d_sens = ops.Fill()(ops.DType()(d_loss), ops.Shape()(d_loss), self.sens) + d_grads = self.grad(self.loss_netD, self.weights_D)(realA, realB, d_sens) + d_res = ops.depend(d_loss, self.optimizerD(d_grads)) + + g_sens = ops.Fill()(ops.DType()(g_loss), ops.Shape()(g_loss), self.sens) + g_grads = self.grad(self.loss_netG, self.weights_G)(realA, realB, g_sens) + g_res = ops.depend(g_loss, self.optimizerG(g_grads)) + return d_res, g_res diff --git a/model_zoo/research/cv/Pix2Pix/src/models/pix2pix.py b/model_zoo/research/cv/Pix2Pix/src/models/pix2pix.py new file mode 100644 index 00000000000..b501e2ac485 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/models/pix2pix.py @@ -0,0 +1,53 @@ +# 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. +# =========================================================================== + +""" + Define Pix2Pix model. +""" + +import mindspore.nn as nn +from .generator_model import UnetGenerator +from .discriminator_model import Discriminator +from .init_w import init_weights +from ..utils.config import get_args + +args = get_args() + +class Pix2Pix(nn.Cell): + def __init__(self, discriminator, generator): + super(Pix2Pix, self).__init__(auto_prefix=True) + self.netD = discriminator + self.netG = generator + + def construct(self, realA, realB): + fakeB = self.netG(realA) + return fakeB + +def get_generator(): + """ + Return a generator by args. + """ + netG = UnetGenerator(in_planes=3, out_planes=3) + init_weights(netG, init_type=args.init_type, init_gain=args.init_gain) + return netG + + +def get_discriminator(): + """ + Return a discriminator by args. + """ + netD = Discriminator(in_planes=6, ndf=64, n_layers=3) + init_weights(netD, init_type=args.init_type, init_gain=args.init_gain) + return netD diff --git a/model_zoo/research/cv/Pix2Pix/src/utils/__init__.py b/model_zoo/research/cv/Pix2Pix/src/utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/research/cv/Pix2Pix/src/utils/config.py b/model_zoo/research/cv/Pix2Pix/src/utils/config.py new file mode 100644 index 00000000000..411a36c8400 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/utils/config.py @@ -0,0 +1,74 @@ +# 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. +# =========================================================================== + +""" + Define the common options that are used in both training and test. +""" + +import argparse +import ast +from mindspore import context + + +def get_args(): + ''' + get args. + ''' + parser = argparse.ArgumentParser(description='Pix2Pix Model') + + # parameters + parser.add_argument('--device_target', type=str, default='Ascend', choices=('Ascend', 'GPU'), + help='device where the code will be implemented (default: Ascend)') + parser.add_argument('--device_num', type=int, default=1, help='device num, default is 1.') + parser.add_argument('--device_id', type=int, default=6, help='device id, default is 0.') + parser.add_argument('--save_graphs', type=ast.literal_eval, default=False, + help='whether save graphs, default is False.') + parser.add_argument('--init_type', type=str, default='normal', help='network initialization, default is normal.') + parser.add_argument('--init_gain', type=float, default=0.02, + help='scaling factor for normal, xavier and orthogonal, default is 0.02.') + parser.add_argument('--load_size', type=int, default=286, help='scale images to this size, default is 286.') + parser.add_argument('--batch_size', type=int, default=1, help='batch_size, default is 1.') + parser.add_argument('--LAMBDA_Dis', type=float, default=0.5, help='weight for Discriminator Loss, default is 0.5.') + parser.add_argument('--LAMBDA_GAN', type=int, default=1, help='weight for GAN Loss, default is 1.') + parser.add_argument('--LAMBDA_L1', type=int, default=100, help='weight for L1 Loss, default is 100.') + parser.add_argument('--beta1', type=float, default=0.5, help='adam beta1, default is 0.5.') + parser.add_argument('--beta2', type=float, default=0.999, help='adam beta2, default is 0.999.') + parser.add_argument('--lr', type=float, default=0.0002, help='the initial learning rate, default is 0.0002.') + parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy, default is linear.') + parser.add_argument('--epoch_num', type=int, default=200, help='epoch number for training, default is 200.') + parser.add_argument('--n_epochs', type=int, default=100, + help='number of epochs with the initial learning rate, default is 100.') + parser.add_argument('--n_epochs_decay', type=int, default=100, + help='number of epochs with the dynamic learning rate, default is 100.') + parser.add_argument('--dataset_size', type=int, default=400, choices=(400, 1096), + help='for Facade_dataset,the number is 400; for Maps_dataset,the number is 1096.') + + # The location of input and output data + parser.add_argument('--train_data_dir', type=str, default=None, help='the file path of input data during training.') + parser.add_argument('--val_data_dir', type=str, default=None, help='the file path of input data during validating.') + parser.add_argument('--train_fakeimg_dir', type=str, default='./results/fake_img/', + help='during training, the file path of stored fake img.') + parser.add_argument('--loss_show_dir', type=str, default='./results/loss_show', + help='during training, the file path of stored loss img.') + parser.add_argument('--ckpt_dir', type=str, default='./results/ckpt', + help='during training, the file path of stored CKPT.') + parser.add_argument('--ckpt', type=str, default=None, help='during validating, the file path of the CKPT used.') + parser.add_argument('--predict_dir', type=str, default='./results/predict/', + help='during validating, the file path of Generated image.') + args = parser.parse_args() + + context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) + + return args diff --git a/model_zoo/research/cv/Pix2Pix/src/utils/tools.py b/model_zoo/research/cv/Pix2Pix/src/utils/tools.py new file mode 100644 index 00000000000..8e60b05dfe5 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/src/utils/tools.py @@ -0,0 +1,74 @@ +# 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. +# =========================================================================== + +""" + Tools for Pix2Pix model. +""" + +import matplotlib.pyplot as plt +import numpy as np +from PIL import Image +from mindspore import Tensor +from src.utils.config import get_args + +args = get_args() + +def save_losses(G_losses, D_losses, idx): + plt.figure(figsize=(10, 5)) + plt.title("Generator and Discriminator Loss During Training") + plt.plot(G_losses, label="G") + plt.plot(D_losses, label="D") + plt.xlabel("iterations") + plt.ylabel("Losses") + plt.legend() + plt.savefig(args.loss_show_dir+"/{}.png".format(idx)) + + +def save_image(img, img_path): + """Save a numpy image to the disk + + Parameters: + img (numpy array / Tensor): image to save. + image_path (str): the path of the image. + """ + if isinstance(img, Tensor): + img = decode_image(img) + elif not isinstance(img, np.ndarray): + raise ValueError("img should be Tensor or numpy array, but get {}".format(type(img))) + + img_pil = Image.fromarray(img) + img_pil.save(img_path+".jpg") + +def decode_image(img): + """Decode a [1, C, H, W] Tensor to image numpy array.""" + mean = 0.5 * 255 + std = 0.5 * 255 + + return (img.asnumpy()[0] * std + mean).astype(np.uint8).transpose((1, 2, 0)) # ——>（256，256，3） + + +def get_lr(): + """ + Linear learning-rate generator. + Keep the same learning rate for the first epochs + and linearly decay the rate to zero over the next epochs. + """ + lrs = [args.lr] * args.dataset_size * args.n_epochs + lr_epoch = 0 + for epoch in range(args.n_epochs_decay): + lr_epoch = args.lr * (args.n_epochs_decay - epoch) / args.n_epochs_decay + lrs += [lr_epoch] * args.dataset_size + lrs += [lr_epoch] * args.dataset_size * (args.epoch_num - args.n_epochs_decay - args.n_epochs) + return Tensor(np.array(lrs).astype(np.float32)) diff --git a/model_zoo/research/cv/Pix2Pix/train.py b/model_zoo/research/cv/Pix2Pix/train.py new file mode 100644 index 00000000000..42ed87a1294 --- /dev/null +++ b/model_zoo/research/cv/Pix2Pix/train.py @@ -0,0 +1,102 @@ +# 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. +# =========================================================================== + +''' + Train Pix2Pix model +''' + +import os +import datetime +import mindspore.nn as nn +from mindspore import Tensor +from mindspore.train.serialization import save_checkpoint +from src.models.loss import D_Loss, D_WithLossCell, G_Loss, G_WithLossCell, TrainOneStepCell +from src.models.pix2pix import Pix2Pix, get_generator, get_discriminator +from src.dataset.pix2pix_dataset import pix2pixDataset, create_train_dataset +from src.utils.config import get_args +from src.utils.tools import save_losses, save_image, get_lr + +if __name__ == '__main__': + + args = get_args() + + # Preprocess the data for training + dataset = pix2pixDataset(root_dir=args.train_data_dir) + ds = create_train_dataset(dataset) + print("ds:", ds.get_dataset_size()) + print("ds:", ds.get_col_names()) + print("ds.shape:", ds.output_shapes()) + + steps_per_epoch = ds.get_dataset_size() + + netG = get_generator() + netD = get_discriminator() + + pix2pix = Pix2Pix(generator=netG, discriminator=netD) + + d_loss_fn = D_Loss() + g_loss_fn = G_Loss() + d_loss_net = D_WithLossCell(backbone=pix2pix, loss_fn=d_loss_fn) + g_loss_net = G_WithLossCell(backbone=pix2pix, loss_fn=g_loss_fn) + + d_opt = nn.Adam(pix2pix.netD.trainable_params(), learning_rate=get_lr(), + beta1=args.beta1, beta2=args.beta2, loss_scale=1) + g_opt = nn.Adam(pix2pix.netG.trainable_params(), learning_rate=get_lr(), + beta1=args.beta1, beta2=args.beta2, loss_scale=1) + + train_net = TrainOneStepCell(loss_netD=d_loss_net, loss_netG=g_loss_net, optimizerD=d_opt, optimizerG=g_opt, sens=1) + train_net.set_train() + + # Training loop + G_losses = [] + D_losses = [] + + data_loader = ds.create_dict_iterator(output_numpy=True, num_epochs=args.epoch_num) + print("Starting Training Loop...") + for epoch in range(args.epoch_num): + for i, data in enumerate(data_loader): + input_image = Tensor(data["input_images"]) + target_image = Tensor(data["target_images"]) + + dis_loss, gen_loss = train_net(input_image, target_image) + + if i % 100 == 0: + print("================start===================") + print("Date time: ", datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')) + print("epoch: ", epoch + 1, "/", args.epoch_num) + print("step: ", i, "/", steps_per_epoch) + print("Dloss: ", dis_loss) + print("Gloss: ", gen_loss) + print("=================end====================") + + # Save fake_imgs + if i == steps_per_epoch - 1: + fake_image = netG(input_image) + save_image(fake_image, args.train_fakeimg_dir + str(epoch + 1)) + print("image generated from epoch", epoch + 1, "saved") + print("The learning rate at this point is：", get_lr()[epoch*i]) + + D_losses.append(dis_loss.asnumpy()) + G_losses.append(gen_loss.asnumpy()) + + print("epoch", epoch + 1, "saved") + # Save losses + save_losses(G_losses, D_losses, epoch + 1) + print("epoch", epoch + 1, "D&G_Losses saved") + print("epoch", epoch + 1, "finished") + # Save checkpoint + if (epoch+1) % 50 == 0: + save_checkpoint(netG, os.path.join(args.ckpt_dir, f"Generator_{epoch+1}.ckpt")) + print("ckpt generated from epoch", epoch + 1, "saved")