forked from huawei/mindspore2022
!18702 Pix2Pix pull request
Merge pull request !18702 from vvvvvvPeng/master
This commit is contained in:
Gitee
commit
4564f37262
|
|
@ -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.
|
||||
|
||||
|
||||
|
||||
# [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).
|
||||
|
|
@ -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
|
||||
|
|
@ -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=======")
|
||||
|
|
@ -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)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 469 KiB |
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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 <opt.n_epochs> epochs
|
||||
and linearly decay the rate to zero over the next <opt.n_epochs_decay> 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))
|
||||
|
|
@ -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")
|
||||
Loading…
Reference in New Issue