README.md

@@ -0,0 +1,295 @@
+# 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/tutorials/en/master/index.html)
+    - [MindSpore Python API](https://www.mindspore.cn/docs/api/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_infer_310.sh                   # launch ascend 310 inference
+  └─run_train_ascend.sh                # launch ascend training(1 pcs)
+  └─run_distribute_train_ascend.sh     # launch ascend training(8 pcs)
+  └─run_eval_ascend.sh                 # launch ascend eval
+  └─run_train_gpu.sh                   # launch gpu training(1 pcs)
+  └─run_distribute_train_gpu.sh        # launch gpu training(8 pcs)
+  └─run_eval_gpu.sh                    # launch gpu 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]
+```
+
+- running distributed trainning on Ascend with fixed parameters
+
+```python
+bash run_distribute_train_ascend.sh [DEVICE_NUM] [DISTRIBUTE] [RANK_TABLE_FILE] [DATASET_PATH] [DATASET_NAME]
+```
+
+- running on GPU with fixed parameters
+
+```python
+python train.py --device_target [GPU] --run_distribute [1] --device_num [8] --dataset_size 400 --train_data_dir [./data/facades/train] --pad_mode REFLECT
+OR
+bash scripts/run_train_gpu.sh [DATASET_PATH] [DATASET_NAME]
+```
+
+- running distributed trainning on GPU with fixed parameters
+
+```python
+bash run_distribute_train_gpu.sh [DATASET_PATH] [DATASET_NAME] [DEVICE_NUM]
+```
+
+## [Evaluation](#contents)
+
+- running on Ascend
+
+```python
+python eval.py --device_target [Ascend] --device_id [0] --val_data_dir [./data/facades/test] --ckpt [./results/ckpt/Generator_200.ckpt] --pad_mode REFLECT
+OR
+bash scripts/run_eval_ascend.sh [DATASET_PATH] [DATASET_NAME] [CKPT_PATH] [RESULT_DIR]
+```
+
+- running on GPU
+
+```python
+python eval.py --device_target [GPU] --device_id [0] --val_data_dir [./data/facades/test] --ckpt [./train/results/ckpt/Generator_200.ckpt] --predict_dir [./train/results/predict/] \
+--dataset_size 1096 --pad_mode REFLECT
+OR
+bash scripts/run_eval_gpu.sh [DATASET_PATH] [DATASET_NAME] [CKPT_PATH] [RESULT_PATH]
+```
+
+**Note:**: Before training and evaluating, create folders like "./results/...". Then you will get the results as following in "./results/predict".
+
+## [310 infer](#contents)
+
+```python
+bash run_infer_310.sh [The path of the MINDIR for 310 infer] [The path of the dataset for 310 infer] y Ascend 0
+```
+
+**Note:**: Before executing 310 infer, create the MINDIR/AIR model using "python export.py --ckpt [The path of the CKPT for exporting] --train_data_dir [The path of the training dataset]".
+
+# [Model Description](#contents)
+
+## [Performance](#contents)
+
+### Training Performance on single device
+
+| Parameters                 | single Ascend                                            | single GPU                                                         |
+| -------------------------- | ----------------------------------------------------------- | --------------------------------------------------------------- |
+| Model Version              | Pix2Pix                                                    | Pix2Pix                                                          |
+| Resource                   | Ascend 910                                               | PCIE V100-32G                                                      |
+| MindSpore Version          | 1.2                                                         | 1.3.0                                                           |
+| Dataset                    | facades                                                  | facades                                                |
+| Training Parameters        | epoch=200, steps=400, batch_size=1, lr=0.0002               | epoch=200, steps=400, batch_size=1, lr=0.0002, pad_mode=REFLECT |
+| Optimizer                  | Adam                                                        | Adam                                                            |
+| Loss Function              | SigmoidCrossEntropyWithLogits Loss & L1 Loss                                   | SigmoidCrossEntropyWithLogits Loss & L1 Loss |
+| outputs                    | probability                                                 | probability                                                     |
+| Speed                      | 1pc(Ascend): 10 ms/step                                  | 1pc(GPU): 40 ms/step                                     |
+| Total time                 | 1pc(Ascend): 0.3h                                        | 1pc(GPU): 0.8 h                                     |
+| Checkpoint for Fine tuning | 207M (.ckpt file)                                            | 207M (.ckpt file)                                              |
+
+| Parameters                 | single Ascend                                            | single GPU                                                         |
+| -------------------------- | ----------------------------------------------------------- | --------------------------------------------------------------- |
+| Model Version              | Pix2Pix                                                    | Pix2Pix                                                          |
+| Resource                   | Ascend 910                                               |
+| MindSpore Version          | 1.2                                                         | 1.3.0                                                           |
+| Dataset                    | maps                                                     | maps                                                               |
+| Training Parameters        | epoch=200, steps=1096, batch_size=1, lr=0.0002              | epoch=200, steps=400, batch_size=1, lr=0.0002, pad_mode=REFLECT |
+| Optimizer                  | Adam                                                        | Adam                                                            |
+| Loss Function              | SigmoidCrossEntropyWithLogits Loss & L1 Loss                                   | SigmoidCrossEntropyWithLogits Loss & L1 Loss |
+| outputs                    | probability                                                 | probability                                                     |
+| Speed                      | 1pc(Ascend): 20 ms/step                                  | 1pc(GPU): 90 ms/step                                     |
+| Total time                 | 1pc(Ascend): 1.58h                                       | 1pc(GPU): 3.3h                                     |
+| Checkpoint for Fine tuning | 207M (.ckpt file)                                            | 207M (.ckpt file)                                              |
+
+### Distributed Training Performance
+
+| Parameters                 | Ascend    (8pcs)                                        | GPU   (8pcs)                                                      |
+| -------------------------- | ----------------------------------------------------------- | --------------------------------------------------------------- |
+| Model Version              | Pix2Pix                                                    | Pix2Pix                                                          |
+| Resource                   | Ascend 910                                               | PCIE V100-32G                                                      |
+| MindSpore Version          | 1.4.1                                                         | 1.3.0                                                           |
+| Dataset                    | facades                                                  | facades                                                |
+| Training Parameters        | epoch=200, steps=400, batch_size=1, lr=0.0002               | epoch=200, steps=400, batch_size=1, lr=0.0002, pad_mode=REFLECT |
+| Optimizer                  | Adam                                                        | Adam                                                            |
+| Loss Function              | SigmoidCrossEntropyWithLogits Loss & L1 Loss                                   | SigmoidCrossEntropyWithLogits Loss & L1 Loss |
+| outputs                    | probability                                                 | probability                                                     |
+| Speed                      | 8pc(Ascend): 15 ms/step                                  | 8pc(GPU): 30 ms/step                                     |
+| Total time                 | 8pc(Ascend): 0.5h                                        | 8pc(GPU): 1 h                                     |
+| Checkpoint for Fine tuning | 207M (.ckpt file)                                            | 207M (.ckpt file)                                              |
+
+| Parameters                 | Ascend    (8pcs)                                            | GPU   (8pcs)                                                         |
+| -------------------------- | ----------------------------------------------------------- | --------------------------------------------------------------- |
+| Model Version              | Pix2Pix                                                    | Pix2Pix                                                          |
+| Resource                   | Ascend 910                                               | PCIE V100-32G                                                    |
+| MindSpore Version          | 1.4.1                                                         | 1.3.0                                                           |
+| Dataset                    | maps                                                     | maps                                                               |
+| Training Parameters        | epoch=200, steps=1096, batch_size=1, lr=0.0002              | epoch=200, steps=400, batch_size=1, lr=0.0002, pad_mode=REFLECT |
+| Optimizer                  | Adam                                                        | Adam                                                            |
+| Loss Function              | SigmoidCross55EntropyWithLogits Loss & L1 Loss                                   | SigmoidCrossEntropyWithLogits Loss & L1 Loss |
+| outputs                    | probability                                                 | probability                                                     |
+| Speed                      | 8pc(Ascend): 20 ms/step                                  | 8pc(GPU): 40 ms/step                                     |
+| Total time                 | 8pc(Ascend): 1.2h                                       | 8pc(GPU): 2.8h                                     |
+| Checkpoint for Fine tuning | 207M (.ckpt file)                                            | 207M (.ckpt file)                                              |
+
+### Evaluation Performance
+
+| Parameters          | single Ascend               | single GPU                  |
+| ------------------- | --------------------------- | --------------------------- |
+| Model Version       | Pix2Pix                     | Pix2Pix                     |
+| Resource            | Ascend 910                  | PCIE V100-32G               |
+| MindSpore Version   | 1.2                         | 1.3.0                       |
+| Dataset             | facades / maps              | facades / maps              |
+| batch_size          | 1                           | 1                           |
+| outputs             | probability                 | probability                 |
+
+# [ModelZoo Homepage](#contents)
+
+Please check the official [homepage](https://gitee.com/mindspore/models).

ascend310_infer/CMakeLists.txt

@@ -0,0 +1,15 @@
+cmake_minimum_required(VERSION 3.14.1)
+project(Ascend310Infer)
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O2 -g -std=c++17 -Werror -Wall -fPIE -Wl,--allow-shlib-undefined")
+set(PROJECT_SRC_ROOT ${CMAKE_CURRENT_LIST_DIR}/)
+option(MINDSPORE_PATH "mindspore install path" "")
+include_directories(${MINDSPORE_PATH})
+include_directories(${MINDSPORE_PATH}/include)
+include_directories(${PROJECT_SRC_ROOT})
+find_library(MS_LIB libmindspore.so ${MINDSPORE_PATH}/lib)
+file(GLOB_RECURSE MD_LIB ${MINDSPORE_PATH}/_c_dataengine*)
+
+add_executable(main src/main.cc src/utils.cc)
+target_link_libraries(main ${MS_LIB} ${MD_LIB} gflags)
+

ascend310_infer/build.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+if [ -d out ]; then
+    rm -rf out
+fi
+
+mkdir out
+cd out || exit
+
+if [ -f "Makefile" ]; then
+  make clean
+fi
+
+cmake .. \
+    -DMINDSPORE_PATH="`pip show mindspore-ascend | grep Location | awk '{print $2"/mindspore"}' | xargs realpath`"
+make

ascend310_infer/fusion_switch.cfg

@@ -0,0 +1 @@
+ConvBatchnormFusionPass:off

ascend310_infer/inc/utils.h

@@ -0,0 +1,32 @@
+/**
+ * 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.
+ */
+
+#ifndef MINDSPORE_INFERENCE_UTILS_H_
+#define MINDSPORE_INFERENCE_UTILS_H_
+
+#include <sys/stat.h>
+#include <dirent.h>
+#include <vector>
+#include <string>
+#include <memory>
+#include "include/api/types.h"
+
+std::vector<std::string> GetAllFiles(std::string_view dirName);
+DIR *OpenDir(std::string_view dirName);
+std::string RealPath(std::string_view path);
+mindspore::MSTensor ReadFileToTensor(const std::string &file);
+int WriteResult(const std::string& imageFile, const std::vector<mindspore::MSTensor> &outputs);
+#endif

ascend310_infer/src/main.cc

@@ -0,0 +1,140 @@
+/**
+ * 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.
+ */
+#include <sys/time.h>
+#include <gflags/gflags.h>
+#include <dirent.h>
+#include <iostream>
+#include <string>
+#include <algorithm>
+#include <iosfwd>
+#include <vector>
+#include <fstream>
+#include <sstream>
+
+#include "include/api/model.h"
+#include "include/api/context.h"
+#include "include/api/types.h"
+#include "include/api/serialization.h"
+#include "minddata/dataset/include/execute.h"
+#include "minddata/dataset/include/vision.h"
+#include "inc/utils.h"
+
+using mindspore::Context;
+using mindspore::Serialization;
+using mindspore::Model;
+using mindspore::Status;
+using mindspore::MSTensor;
+using mindspore::dataset::Execute;
+using mindspore::ModelType;
+using mindspore::GraphCell;
+using mindspore::kSuccess;
+
+DEFINE_string(mindir_path, "", "mindir path");
+DEFINE_string(input0_path, ".", "input0 path");
+DEFINE_string(fusion_switch_path, ".", "fusion switch path");
+DEFINE_int32(device_id, 0, "device id");
+
+int main(int argc, char **argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  if (RealPath(FLAGS_mindir_path).empty()) {
+    std::cout << "Invalid mindir" << std::endl;
+    return 1;
+  }
+
+  if (RealPath(FLAGS_fusion_switch_path).empty()) {
+    std::cout << "Invalid fusion switch path" << std::endl;
+    return 1;
+  }
+
+  auto context = std::make_shared<Context>();
+  auto ascend310 = std::make_shared<mindspore::Ascend310DeviceInfo>();
+  ascend310->SetDeviceID(FLAGS_device_id);
+  context->MutableDeviceInfo().push_back(ascend310);
+  mindspore::Graph graph;
+  Serialization::Load(FLAGS_mindir_path, ModelType::kMindIR, &graph);
+  if (!FLAGS_fusion_switch_path.empty()) {
+    ascend310->SetFusionSwitchConfigPath(FLAGS_fusion_switch_path);
+  }
+
+  Model model;
+  Status ret = model.Build(GraphCell(graph), context);
+  if (ret != kSuccess) {
+    std::cout << "ERROR: Build failed." << std::endl;
+    return 1;
+  }
+
+  std::vector<MSTensor> model_inputs = model.GetInputs();
+  if (model_inputs.empty()) {
+    std::cout << "Invalid model, inputs is empty." << std::endl;
+    return 1;
+  }
+
+  auto input0_files = GetAllFiles(FLAGS_input0_path);
+
+  if (input0_files.empty()) {
+    std::cout << "ERROR: input data empty." << std::endl;
+    return 1;
+  }
+
+  std::map<double, double> costTime_map;
+  size_t size = input0_files.size();
+
+  for (size_t i = 0; i < size; ++i) {
+    struct timeval start = {0};
+    struct timeval end = {0};
+    double startTimeMs;
+    double endTimeMs;
+    std::vector<MSTensor> inputs;
+    std::vector<MSTensor> outputs;
+    std::cout << "Start predict input files:" << input0_files[i] << std::endl;
+
+    auto input0 = ReadFileToTensor(input0_files[i]);
+
+    inputs.emplace_back(model_inputs[0].Name(), model_inputs[0].DataType(), model_inputs[0].Shape(),
+                        input0.Data().get(), input0.DataSize());
+
+    gettimeofday(&start, nullptr);
+    ret = model.Predict(inputs, &outputs);
+    gettimeofday(&end, nullptr);
+    if (ret != kSuccess) {
+      std::cout << "Predict " << input0_files[i] << " failed." << std::endl;
+      return 1;
+    }
+    startTimeMs = (1.0 * start.tv_sec * 1000000 + start.tv_usec) / 1000;
+    endTimeMs = (1.0 * end.tv_sec * 1000000 + end.tv_usec) / 1000;
+    costTime_map.insert(std::pair<double, double>(startTimeMs, endTimeMs));
+    WriteResult(input0_files[i], outputs);
+  }
+  double average = 0.0;
+  int inferCount = 0;
+
+  for (auto iter = costTime_map.begin(); iter != costTime_map.end(); iter++) {
+    double diff = 0.0;
+    diff = iter->second - iter->first;
+    average += diff;
+    inferCount++;
+  }
+  average = average / inferCount;
+  std::stringstream timeCost;
+  timeCost << "NN inference cost average time: " << average << " ms of infer_count " << inferCount << std::endl;
+  std::cout << "NN inference cost average time: " << average << "ms of infer_count " << inferCount << std::endl;
+  std::string fileName = "./time_Result" + std::string("/test_perform_static.txt");
+  std::ofstream fileStream(fileName.c_str(), std::ios::trunc);
+  fileStream << timeCost.str();
+  fileStream.close();
+  costTime_map.clear();
+  return 0;
+}

ascend310_infer/src/utils.cc

@@ -0,0 +1,128 @@
+/**
+ * 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.
+ */
+
+#include <fstream>
+#include <algorithm>
+#include <iostream>
+#include "inc/utils.h"
+
+using mindspore::MSTensor;
+using mindspore::DataType;
+
+std::vector<std::string> GetAllFiles(std::string_view dirName) {
+  struct dirent *filename;
+  DIR *dir = OpenDir(dirName);
+  if (dir == nullptr) {
+    return {};
+  }
+  std::vector<std::string> res;
+  while ((filename = readdir(dir)) != nullptr) {
+    std::string dName = std::string(filename->d_name);
+    if (dName == "." || dName == ".." || filename->d_type != DT_REG) {
+      continue;
+    }
+    res.emplace_back(std::string(dirName) + "/" + filename->d_name);
+  }
+  std::sort(res.begin(), res.end());
+  for (auto &f : res) {
+    std::cout << "image file: " << f << std::endl;
+  }
+  return res;
+}
+
+int WriteResult(const std::string& imageFile, const std::vector<MSTensor> &outputs) {
+  std::string homePath = "./result_Files";
+  for (size_t i = 0; i < outputs.size(); ++i) {
+    size_t outputSize;
+    std::shared_ptr<const void> netOutput;
+    netOutput = outputs[i].Data();
+    outputSize = outputs[i].DataSize();
+    int pos = imageFile.rfind('/');
+    std::string fileName(imageFile, pos + 1);
+    fileName.replace(fileName.find('.'), fileName.size() - fileName.find('.'), '_' + std::to_string(i) + ".bin");
+    std::string outFileName = homePath + "/" + fileName;
+    FILE * outputFile = fopen(outFileName.c_str(), "wb");
+    fwrite(netOutput.get(), outputSize, sizeof(char), outputFile);
+    fclose(outputFile);
+    outputFile = nullptr;
+  }
+  return 0;
+}
+
+mindspore::MSTensor ReadFileToTensor(const std::string &file) {
+  if (file.empty()) {
+    std::cout << "Pointer file is nullptr" << std::endl;
+    return mindspore::MSTensor();
+  }
+
+  std::ifstream ifs(file);
+  if (!ifs.good()) {
+    std::cout << "File: " << file << " is not exist" << std::endl;
+    return mindspore::MSTensor();
+  }
+
+  if (!ifs.is_open()) {
+    std::cout << "File: " << file << "open failed" << std::endl;
+    return mindspore::MSTensor();
+  }
+
+  ifs.seekg(0, std::ios::end);
+  size_t size = ifs.tellg();
+  mindspore::MSTensor buffer(file, mindspore::DataType::kNumberTypeUInt8, {static_cast<int64_t>(size)}, nullptr, size);
+
+  ifs.seekg(0, std::ios::beg);
+  ifs.read(reinterpret_cast<char *>(buffer.MutableData()), size);
+  ifs.close();
+
+  return buffer;
+}
+
+DIR *OpenDir(std::string_view dirName) {
+  if (dirName.empty()) {
+    std::cout << " dirName is null ! " << std::endl;
+    return nullptr;
+  }
+  std::string realPath = RealPath(dirName);
+  struct stat s;
+  lstat(realPath.c_str(), &s);
+  if (!S_ISDIR(s.st_mode)) {
+    std::cout << "dirName is not a valid directory !" << std::endl;
+    return nullptr;
+  }
+  DIR *dir;
+  dir = opendir(realPath.c_str());
+  if (dir == nullptr) {
+    std::cout << "Can not open dir " << dirName << std::endl;
+    return nullptr;
+  }
+  std::cout << "Successfully opened the dir " << dirName << std::endl;
+  return dir;
+}
+
+std::string RealPath(std::string_view path) {
+  char realPathMem[PATH_MAX] = {0};
+  char *realPathRet = nullptr;
+  realPathRet = realpath(path.data(), realPathMem);
+
+  if (realPathRet == nullptr) {
+    std::cout << "File: " << path << " is not exist.";
+    return "";
+  }
+
+  std::string realPath(realPathMem);
+  std::cout << path << " realpath is: " << realPath << std::endl;
+  return realPath;
+}

config.py

@@ -1,90 +0,0 @@
-import albumentations as A
-import mindspore.dataset.vision.py_transforms as py_vision
-
-# DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
-TRAIN_DIR = "maps/train"
-VAL_DIR = "maps/val"
-LEARNING_RATE = 0.0002
-BATCH_SIZE = 1
-NUM_WORKERS = 2
-IMAGE_SIZE = 256
-CHANNELS_IMG = 3
-L1_LAMBDA = 100
-LAMBDA_GP = 10
-NUM_EPOCHS = 100
-LOAD_MODEL = False
-SAVE_MODEL = True
-CHECKPOINT_DISC = "disc.pth.tar"
-CHECKPOINT_GEN = "gen.pth.tar"
-
-both_transform = A.Compose(
-    [A.Resize(width=256, height=256),A.HorizontalFlip(p=0.5)], additional_targets={"image0": "image"},
-)
-
-transform_only_input = A.Compose(
-    [
-        A.HorizontalFlip(p=0.5),
-        A.ColorJitter(p=0.2),
-        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0,),
-        py_vision.ToTensor(),
-    ]
-)
-
-transform_only_mask = A.Compose(
-    [
-        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.0,),
-        py_vision.ToTensor(),
-    ]
-)
-
-
-
-
-
-
-
-
-# 下面是修改后 不使用albumentation库房 的代码
-# import torch
-# # import albumentations as A
-# # from albumentations.pytorch import ToTensorV2
-# import mindspore.dataset.vision.py_transforms as py_vision
-# import mindspore.dataset.transforms.py_transforms as py_transforms
-#
-#
-#
-# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# TRAIN_DIR = "maps/train"
-# VAL_DIR = "maps/val"
-# LEARNING_RATE = 2e-4
-# BATCH_SIZE = 16
-# NUM_WORKERS = 2
-# IMAGE_SIZE = 256
-# CHANNELS_IMG = 3
-# L1_LAMBDA = 100
-# LAMBDA_GP = 10
-# NUM_EPOCHS = 500
-# LOAD_MODEL = False
-# SAVE_MODEL = False
-# CHECKPOINT_DISC = "disc.pth.tar"
-# CHECKPOINT_GEN = "gen.pth.tar"
-#
-# both_transform = py_transforms.Compose(
-#     [py_transforms.Resize(width=256, height=256),], additional_targets={"image0": "image"},
-# )
-#
-# transform_only_input = py_transforms.Compose(
-#     [
-#         py_vision.RandomHorizontalFlip(p=0.5),
-#         py_vision.RandomColorAdjust(p=0.2),      #与ColorJitter()参数不对应
-#         py_vision.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),   #MS中没有 max_pixel_value=255.0 参数
-#         py_vision.ToTensor(),
-#     ]
-# )
-#
-# transform_only_mask = py_transforms.Compose(
-#     [
-#         py_vision.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-#         py_vision.ToTensor(),
-#     ]
-# )

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

dataset/distributed_sampler.py

@@ -1,48 +0,0 @@
-from __future__ import division
-import math
-import numpy as np
-
-
-"""
-用于多机多卡进行训练
-"""
-class DistributedSampler:
-
-    def __init__(self, dataset_size, num_replicas=None, rank=None, shuffle=True):
-        if num_replicas is None:
-            print("***********Setting world_size to 1 since it is not passed in ******************")
-            num_replicas = 1
-        if rank is None:
-            print("***********Setting rank to 0 since it is not passed in ******************")
-            rank = 0
-        self.dataset_size = dataset_size
-        self.num_replicas = num_replicas
-        self.rank = rank
-        self.epoch = 0
-        self.num_samples = int(math.ceil(dataset_size * 1.0 / self.num_replicas))
-        self.total_size = self.num_samples * self.num_replicas
-        self.shuffle = shuffle
-
-    def __iter__(self):
-        # deterministically shuffle based on epoch
-        if self.shuffle:
-            indices = np.random.RandomState(seed=self.epoch).permutation(self.dataset_size)
-            # np.array type. number from 0 to len(dataset_size)-1, used as index of dataset
-            indices = indices.tolist()
-            self.epoch += 1
-            # change to list type
-        else:
-            indices = list(range(self.dataset_size))
-
-        # add extra samples to make it evenly divisible
-        indices += indices[:(self.total_size - len(indices))]
-        assert len(indices) == self.total_size
-
-        # subsample
-        indices = indices[self.rank:self.total_size:self.num_replicas]
-        assert len(indices) == self.num_samples
-
-        return iter(indices)
-
-    def __len__(self):
-        return self.num_samples

dataset/pix2pix_dataset.py

@@ -1,70 +0,0 @@
-import numpy as np
-import os
-from PIL import Image
-from mindspore import dataset as de
-import mindspore.dataset.vision.c_transforms as C
-
-class MapDataset():
-    def __init__(self,root_dir):
-        self.root_dir = root_dir
-        self.list_files = os.listdir(self.root_dir)
-        print(self.list_files)
-        # print(len(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)
-        image = np.array(Image.open(img_path))
-        input_image = image[:, :256, :]
-        target_image = image[:, 256:, :]
-
-
-        # 第三方库albumentation的图像操作
-        # augmentations = config.both_transform(image=input_image, image0=target_image)
-        # input_image = augmentations["image"]
-        # target_image = augmentations["image0"]
-        #
-        # input_image = config.transform_only_input(image=input_image)["image"]
-        # target_image = config.transform_only_mask(image=target_image)["image"]
-
-        return input_image, target_image
-
-def create_dataset(dataset):
-
-    # dataset=MapDataset(root_dir="../maps/train/")
-
-    mean = [0.5 * 255] * 3
-    std = [0.5 * 255] * 3
-
-    trans = [
-        C.RandomResizedCrop(256, scale=(0.5, 1.0), ratio=(0.75, 1.333)),
-        C.RandomHorizontalFlip(prob=0.5),
-        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=trans,input_columns=["input_images"])
-    train_ds=train_ds.map(operations=trans,input_columns=["target_images"])
-    train_ds=train_ds.batch(1)
-    train_ds=train_ds.repeat(1)
-
-    # print("ds:", train_ds.get_dataset_size())
-    # print("ds:", train_ds.get_col_names())
-    # print("ds.shape:", train_ds.output_shapes())
-    # for data in train_ds.create_dict_iterator():
-    #     print(data["input_images"],data["target_images"])
-
-    return train_ds
-
-
-# if __name__ == "__main__":
-#     dataset = MapDataset(root_dir="../maps/train/")
-#     create_dataset(dataset)
-
-
-
-

discriminator_model.py

@@ -1,73 +0,0 @@
-import mindspore
-import mindspore.nn as nn
-from mindspore.ops import Concat
-import numpy as np
-
-
-class CNNBlock(nn.Cell):
-    def __init__(self, in_channels, out_channels, stride=2):
-        super(CNNBlock, self).__init__()
-        self.conv = nn.SequentialCell(
-            nn.Conv2d(
-                in_channels, out_channels, 4, stride, pad_mode="pad", has_bias=False,
-            ),
-            nn.BatchNorm2d(out_channels),
-            nn.LeakyReLU(alpha= 0.2),
-        )
-
-    def construct(self, x):
-        return self.conv(x)
-
-
-class Discriminator(nn.Cell):
-    def __init__(self, in_channels=3, features=[64,128,256,512]):   # features——output_channel
-        super().__init__()
-        self.initial = nn.SequentialCell(
-            nn.Conv2d(
-                in_channels * 2,
-                features[0],
-                kernel_size=4,
-                stride=2,
-                padding=1,
-                pad_mode="pad",
-            ),
-            nn.LeakyReLU(alpha=0.2),
-        )
-
-        layers = []
-        in_channels = features[0]
-        for feature in features[1:]:
-            layers.append(
-                CNNBlock(in_channels=in_channels, out_channels=feature, stride=1 if feature == features[-1] else 2),
-            )
-            in_channels = feature
-
-        layers.append(
-            nn.Conv2d(
-                in_channels, 1, kernel_size=4, stride=1, padding=1, pad_mode="pad"
-            ),
-        )
-
-        self.model = nn.SequentialCell(*layers)
-        self.concat = Concat(axis=1)
-
-    def construct(self, x, y):
-        x_y = self.concat((x,y))
-        x_y = self.initial(x_y)
-        x_y = self.model(x_y)
-        return x_y
-
-
-
-def test():
-    x=mindspore.Tensor(np.random.rand(1,3,256,256), mindspore.float32)
-    y=mindspore.Tensor(np.random.randn(1,3,256,256), mindspore.float32)
-    print(x)
-    model=Discriminator(in_channels=3)
-    print(model)
-    result=model(x,y)
-    print(result)
-
-
-if __name__ == "__main__":
-    test()

eval.py

@@ -0,0 +1,78 @@
+# 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.
+"""
+
+import os
+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)
+
+    if not os.path.isdir(args.predict_dir):
+        os.makedirs(args.predict_dir)
+
+    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=======")

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)

generator_model.py

@@ -1,106 +0,0 @@
-import mindspore
-import mindspore.nn as nn
-from mindspore.ops import Concat
-import numpy as np
-
-class Block(nn.Cell):
-    def __init__(self, in_channels, out_channels, down=True, act="relu", use_dropout=False):
-        super(Block, self).__init__()
-        self.conv = nn.SequentialCell(
-            nn.Conv2d(in_channels, out_channels, 4, 2, padding=1, has_bias=False, pad_mode="pad")
-            if down
-            else nn.Conv2dTranspose(in_channels, out_channels, 4, 2, padding=1, has_bias=False,pad_mode="pad"),
-            nn.BatchNorm2d(out_channels),
-            nn.ReLU() if act == "relu" else nn.LeakyReLU(alpha=0.2),
-        )
-
-        self.use_dropout = use_dropout
-        self.dropout = nn.Dropout(keep_prob=0.5)
-        self.down = down
-
-    def construct(self, x):
-        x = self.conv(x)
-        return self.dropout(x) if self.use_dropout else x
-
-
-class Generator(nn.Cell):
-    def __init__(self, in_channels=3, features=64):
-        super(Generator, self).__init__()
-        self.initial_down = nn.SequentialCell(
-            nn.Conv2d(in_channels, features, 4, 2, padding=0),
-            nn.LeakyReLU(alpha=0.2),
-        )
-        self.down1 = Block(features, features * 2, down=True, act="leaky", use_dropout=False)
-        self.down2 = Block(
-            features * 2, features * 4, down=True, act="leaky", use_dropout=False
-        )
-        self.down3 = Block(
-            features * 4, features * 8, down=True, act="leaky", use_dropout=False
-        )
-        self.down4 = Block(
-            features * 8, features * 8, down=True, act="leaky", use_dropout=False
-        )
-        self.down5 = Block(
-            features * 8, features * 8, down=True, act="leaky", use_dropout=False
-        )
-        self.down6 = Block(
-            features * 8, features * 8, down=True, act="leaky", use_dropout=False
-        )
-        self.bottleneck = nn.SequentialCell(
-            nn.Conv2d(features * 8, features * 8, 4, 2, padding=1,pad_mode="pad"),
-            nn.ReLU()
-        )
-
-        self.up1 = Block(features * 8, features * 8, down=False, act="relu", use_dropout=True)
-        self.up2 = Block(
-            features * 8 * 2, features * 8, down=False, act="relu", use_dropout=True
-        )
-        self.up3 = Block(
-            features * 8 * 2, features * 8, down=False, act="relu", use_dropout=True
-        )
-        self.up4 = Block(
-            features * 8 * 2, features * 8, down=False, act="relu", use_dropout=False
-        )
-        self.up5 = Block(
-            features * 8 * 2, features * 4, down=False, act="relu", use_dropout=False
-        )
-        self.up6 = Block(
-            features * 4 * 2, features * 2, down=False, act="relu", use_dropout=False
-        )
-        self.up7 = Block(features * 2 * 2, features, down=False, act="relu", use_dropout=False)
-        self.final_up = nn.SequentialCell(
-            nn.Conv2dTranspose(features * 2, in_channels, kernel_size=4, stride=2, padding=1,pad_mode="pad"),
-            nn.Tanh(),
-        )
-        self.concat = Concat(axis=1)
-
-    def construct(self, x):
-        d1 = self.initial_down(x)
-        d2 = self.down1(d1)
-        d3 = self.down2(d2)
-        d4 = self.down3(d3)
-        d5 = self.down4(d4)
-        d6 = self.down5(d5)
-        d7 = self.down6(d6)
-        bottleneck = self.bottleneck(d7)
-        up1 = self.up1(bottleneck)
-        up2 = self.up2(self.concat((up1, d7)))
-        up3 = self.up3(self.concat((up2, d6)))
-        up4 = self.up4(self.concat((up3, d5)))
-        up5 = self.up5(self.concat((up4, d4)))
-        up6 = self.up6(self.concat((up5, d3)))
-        up7 = self.up7(self.concat((up6, d2)))
-        return self.final_up(self.concat((up7, d1)))
-
-
-def test():
-    x = mindspore.Tensor(np.random.randn(1,3,256,256), mindspore.float32)
-    print(x)
-    model = Generator(in_channels=3)
-    print(model)
-    preds = model(x)
-    print(preds)
-
-
-if __name__ == "__main__":
-    test()

losses.py

@@ -1,126 +0,0 @@
-import mindspore
-import numpy as np
-import mindspore.nn as nn
-import mindspore.ops as ops
-from mindspore.ops  import functional as F
-
-# 测试引入
-import config
-from discriminator_model import Discriminator
-from generator_model import Generator
-
-
-# Discriminator_With_Loss
-class WithLossCellDis(nn.Cell):
-    def __init__(self,netD,netG,loss_fn1):
-        super(WithLossCellDis, self).__init__(auto_prefix=True)
-        self.netD=netD
-        self.netG=netG
-        self.loss_fn1=loss_fn1
-        self.ones=ops.OnesLike()
-        self.zeros=ops.ZerosLike()
-
-    def construct(self, x, y):
-        y_fake=self.netG(x)
-        # print("y_fake:\n",y_fake)
-        y_fake=F.stop_gradient(y_fake)
-        D_real=self.netD(x,y)
-        # print(" D_real:\n", D_real)
-        D_real_loss=self.loss_fn1(D_real,self.ones(D_real))
-        D_fake = self.netD(x, y_fake)
-        D_fake_loss=self.loss_fn1(D_fake,self.zeros(D_fake))
-        D_loss=(D_real_loss+D_fake_loss)*0.5
-        # print(" D_loss:\n", D_loss)
-        return D_loss
-
-    @property
-    def backbone_network(self):
-        return self.netD
-
-
-
-
-# Generator_With_Loss
-class WithLossCellGen(nn.Cell):
-    def __init__(self,netD,netG,loss_fn1,loss_fn2):    #fn1为BCE_Loss,fn2为L1_Loss
-        super(WithLossCellGen, self).__init__(auto_prefix=True)
-        self.netD=netD
-        self.netG=netG
-        self.loss_fn1=loss_fn1
-        self.loss_fn2=loss_fn2
-        self.ones = ops.OnesLike()
-
-    def construct(self, x, y):
-        y_fake = self.netG(x)
-        # print("y_fake:\n",y_fake)
-        D_fake=self.netD(x,y_fake)
-        # print(" D_real:\n", D_fake)
-        G_fake_loss=self.loss_fn1(D_fake,self.ones(D_fake))
-        L1_loss=self.loss_fn2(y_fake,y)*config.L1_LAMBDA
-        G_loss=G_fake_loss+L1_loss
-        # print(G_loss)
-        return G_loss
-
-    @property
-    def backbone_network(self):
-        return self.netG
-
-
-# def testlossG():
-#     x=mindspore.Tensor(np.random.rand(1,3,256,256), mindspore.float32)
-#     y=mindspore.Tensor(np.random.randn(1,3,256,256), mindspore.float32)
-#     loss_model=WithLossCellGen(netD=Discriminator(),netG=Generator(),loss_fn1=nn.BCEWithLogitsLoss(),loss_fn2=nn.L1Loss())
-#     # print(loss_model)
-#     loss=loss_model(x,y)
-#     print(loss)
-#
-#
-# if __name__ == "__main__":
-#     testlossG()
-
-
-
-
-
-
-
-
-
-
-
-
-
-# class DiscriminatorLoss(nn.Cell):
-#     def __init__(self,disc,gen):
-#         super(DiscriminatorLoss, self).__init__()
-#         self.disc=disc
-#         self.gen=gen
-#         self.BWL_Loss=nn.BCEWithLogitsLoss()
-#
-#     def construct(self, x, y):
-#         y_fake=self.gen(x)
-#         D_real=self.disc(x,y)
-#         D_real_loss=self.BWL_Loss(D_real,ops.OnesLike(D_real))
-#         D_fake=self.disc(x,y_fake)   #.detach()处理
-#         D_fake_loss=self.BWL_Loss(D_fake,ops.ZerosLike(D_fake))
-#         D_loss=(D_real_loss+D_fake_loss)*0.5
-#         return D_loss
-#
-#
-# class GeneratorLoss(nn.Cell):
-#     def __init__(self,disc,gen):
-#         super(GeneratorLoss, self).__init__()
-#         self.disc=disc
-#         self.gen=gen
-#         self.BWL_Loss=nn.BCEWithLogitsLoss()
-#         self.L1_loss=nn.L1Loss()
-#
-#     def construct(self, x, y):
-#         y_fake = self.gen(x)
-#         D_fake = self.disc(x, y_fake)
-#         G_fake_loss=self.BWL_Loss(D_fake,ops.OnesLike(D_fake))
-#         L1_loss=self.L1_loss(y_fake,y)*config.L1_LAMBDA
-#         G_loss=G_fake_loss+L1_loss
-#         return y_fake,G_loss
-#
-#

pix2pix.py

@@ -1,146 +0,0 @@
-import mindspore as ms
-import mindspore.nn as nn
-from mindspore import context
-from mindspore.context import ParallelMode
-from mindspore.parallel._auto_parallel_context import auto_parallel_context
-from mindspore.communication.management import get_group_size
-import mindspore.ops as ops
-from discriminator_model import Discriminator
-from generator_model import Generator
-
-class Pix2Pix(nn.Cell):
-    def __init__(self,TrainOneStepCellDis,TrainOneStepCellGen):
-        super(Pix2Pix, self).__init__(auto_prefix=True)
-        self.TrainOneStepCellDis=TrainOneStepCellDis
-        self.TrainOneStepCellGen=TrainOneStepCellGen
-
-    def construct(self, x, y):
-        output_D=self.TrainOneStepCellDis(x,y).view(-1)
-        netD_loss=output_D.mean()
-        output_G=self.TrainOneStepCellGen(x,y).view(-1)
-        netG_loss=output_G.mean()
-        return netD_loss,netG_loss
-
-
-
-
-
-
-
-
-
-
-
-# class DisWithLossCell(nn.Cell):
-#     def __int__(self, dis_loss_fn):
-#         super(DisWithLossCell, self).__int__(auto_prefix=False)
-#         self.dis_loss_fn = dis_loss_fn
-#
-#     def construct(self, x, y):
-#         dis_loss = self.dis_loss_fn(x, y)
-#
-#         return dis_loss
-
-#
-# class WithLossCell(nn.Cell):
-#     def __int__(self, network):        #network即为GeneratorLoss
-#         super(WithLossCell, self).__int__(auto_prefix=False)
-#         self.network = network
-#
-#     def construct(self, x, y):
-#         fake_image,gen_loss = self.network(x, y)
-#
-#         return gen_loss
-#
-#
-# class TrainOneStepDis(nn.Cell):
-#     """
-#     Encapsulation class of Cycle GAN discriminator network training.
-#
-#     Append an optimizer to the training network after that the construct
-#     function can be called to create the backward graph.
-#
-#     Args:
-#         D (Cell): Discriminator with loss Cell. Note that loss function should have been added.
-#         optimizer (Optimizer): Optimizer for updating the weights.
-#         sens (Number): The adjust parameter. Default: 1.0.
-#     """
-#     def __init__(self,D,optimizer,sens=1.0):   #D为DiscriminatorLoss
-#         super(TrainOneStepDis, self).__init__(auto_prefix=False)
-#         self.optimizer=optimizer
-#         self.D=D
-#         self.D.set_grad()
-#         self.D.set_train()
-#         self.grad=ops.GradOperation(get_by_list=True,sens_param=True)
-#         self.sens=sens
-#         self.weights=ms.ParameterTuple(D.trainaable_params())
-#         self.reducer_flag=False
-#         self.grad_reducer=None
-#         self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
-#         if self.parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
-#             self.reducer_flag = True
-#         if self.reducer_flag:
-#             mean = context.get_auto_parallel_context("gradients_mean")
-#             if auto_parallel_context().get_device_num_is_set():
-#                 degree = context.get_auto_parallel_context("device_num")
-#             else:
-#                 degree = get_group_size()
-#             self.grad_reducer = nn.DistributedGradReducer(optimizer.parameters, mean, degree)
-#
-#     def construct(self, x, y):
-#         weights = self.weights
-#         ld = self.D(x,y)      #discriminator_loss
-#         sens_d = ops.Fill()(ops.DType()(ld), ops.Shape()(ld), self.sens)
-#         grads_d = self.grad(self.D, weights)(x,y,sens_d)
-#         if self.reducer_flag:
-#             # apply grad reducer on grads
-#             grads_d = self.grad_reducer(grads_d)
-#         return ops.depend(ld, self.optimizer(grads_d))
-#
-#
-# class TrainOneStepGen(nn.Cell):
-#     """
-#     Encapsulation class of Cycle GAN generator network training.
-#
-#     Append an optimizer to the training network after that the construct
-#     function can be called to create the backward graph.
-#
-#     Args:
-#         G (Cell): Generator with loss Cell. Note that loss function should have been added.
-#         generator (Cell): Generator of CycleGAN.
-#         optimizer (Optimizer): Optimizer for updating the weights.
-#         sens (Number): The adjust parameter. Default: 1.0.
-#     """
-#     def __init__(self,G,gen,optimizer,sens=1.0):   #G为GeneratorLoss，gen为generator
-#         super(TrainOneStepGen, self).__init__(auto_prefix=False)
-#         self.optimizer = optimizer
-#         self.G = G
-#         self.G.set_grad()
-#         self.G.set_train()
-#         self.grad=ops.GradOperation(get_by_list=True,sens_param=True)
-#         self.sens=sens
-#         self.weights = ms.ParameterTuple(gen.trainable_params())
-#         self.net = WithLossCell(G)       #
-#         self.reducer_flag = False
-#         self.grad_reducer = None
-#         self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
-#         if self.parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
-#             self.reducer_flag = True
-#         if self.reducer_flag:
-#             mean = context.get_auto_parallel_context("gradients_mean")
-#             if auto_parallel_context().get_device_num_is_set():
-#                 degree = context.get_auto_parallel_context("device_num")
-#             else:
-#                 degree = get_group_size()
-#             self.grad_reducer = nn.DistributedGradReducer(optimizer.parameters, mean, degree)
-#
-#     def construct(self, x, y):
-#         weights = self.weights
-#         fake_image,lg= self.G(x, y)   #G为Generatorloss
-#         sens = ops.Fill()(ops.DType()(lg), ops.Shape()(lg), self.sens)
-#         grads_g = self.grad(self.net, weights)(x, y, sens)
-#         if self.reducer_flag:
-#             # apply grad reducer on grads
-#             grads_g = self.grad_reducer(grads_g)
-#
-#         return fake_image,ops.depend(lg, self.optimizer(grads_g))

postprocess.py

@@ -0,0 +1,60 @@
+# 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.
+# ============================================================================
+"""
+    postprocess
+"""
+import os
+import numpy as np
+from PIL import Image
+from src.utils.config import get_args
+from mindspore import Tensor
+
+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 = img.asnumpy()
+    elif not isinstance(img, np.ndarray):
+        raise ValueError("img should be Tensor or numpy array, but get {}".format(type(img)))
+    img = decode_image(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 * std + mean).astype(np.uint8).transpose((1, 2, 0))
+
+if __name__ == '__main__':
+    args = get_args()
+
+    result_dir = "./result_Files"
+    object_imageSize = 256
+    rst_path = result_dir
+
+    for i in range(len(os.listdir(rst_path))):
+        file_name = os.path.join(rst_path, "Pix2Pix_data_bs" + str(args.batch_size) + '_' + str(i) + '_0.bin')
+        output = np.fromfile(file_name, np.float32).reshape(3, object_imageSize, object_imageSize)
+        print(output.shape)
+        save_image(output, './310_infer_img' + str(i + 1))
+        print("=======image", i + 1, "saved success=======")
+    print("Generate images success!")

preprocess.py

@@ -0,0 +1,38 @@
+# 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
+"""
+import os
+import numpy as np
+from src.dataset.pix2pix_dataset import pix2pixDataset_val, create_val_dataset
+from src.utils.config import get_args
+
+if __name__ == '__main__':
+    args = get_args()
+    result_path = "./preprocess_Result/"
+    dataset_val = pix2pixDataset_val(root_dir=args.val_data_dir)
+    ds_val = create_val_dataset(dataset_val)
+    img_path = os.path.join(result_path, "00_data")
+    os.makedirs(img_path)
+    target_images_list = []
+    for i, data in enumerate(ds_val.create_dict_iterator(output_numpy=True)):
+        file_name = "Pix2Pix_data_bs" + str(args.batch_size) + "_" + str(i) + ".bin"
+        file_path = img_path + "/" + file_name
+        data['input_images'].tofile(file_path)
+        target_images_list.append(data['target_images'])
+
+    np.save(result_path + "target_images_ids.npy", target_images_list)
+    print("=" * 20, "export bin files finished", "=" * 20)

scripts/run_distribute_train_ascend.sh

@@ -0,0 +1,64 @@
+#!/bin/bash
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+if [ $# != 5 ]
+then
+    echo "Usage: sh run_distribute_train.sh [DEVICE_NUM] [DISTRIBUTE] [RANK_TABLE_FILE] [DATASET_PATH] [DATASET_NAME]"
+    exit 1
+fi
+
+echo "After running the script, the network runs in the background. The log will be generated in LOGx/log.txt"
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+export RANK_SIZE=$1
+export DISTRIBUTE=$2
+export RANK_TABLE_FILE=$(get_real_path $3)
+export DATASET_PATH=$(get_real_path $4)
+
+for((i=0;i<RANK_SIZE;i++))
+do
+        export DEVICE_ID=$i
+        rm -rf ./train_parallel$i
+        mkdir ./train_parallel$i
+        mkdir ./train_parallel$i/results
+        mkdir ./train_parallel$i/results/fake_img
+        mkdir ./train_parallel$i/results/ckpt
+        mkdir ./train_parallel$i/results/predict
+        mkdir ./train_parallel$i/results/loss_show
+        cp -r ../src ./train_parallel$i
+        cp -r ../*.py ./train_parallel$i
+        cd ./train_parallel$i || exit
+        export RANK_ID=$i
+        echo "start training for rank $i, device $DEVICE_ID"
+        env > env.log
+        if [ $5 == 'facades' ];
+        then
+                python train.py --run_distribute=$DISTRIBUTE --device_target Ascend --device_num=$RANK_SIZE \
+                                --device_id=$DEVICE_ID --dataset_size 400 --train_data_dir $DATASET_PATH &> log &
+        
+        elif [ $5 == 'maps' ];
+        then
+                python train.py --run_distribute=$DISTRIBUTE --device_target Ascend --device_num=$RANK_SIZE \
+                                --device_id=$DEVICE_ID --dataset_size 1096 --train_data_dir $DATASET_PATH &> log &
+        fi
+        cd ..
+done

scripts/run_distribute_train_gpu.sh

@@ -0,0 +1,64 @@
+#!/bin/bash
+# Copyright 2020-2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+if [ $# != 3 ]
+then
+    echo "Usage: bash run_distribute_train_gpu.sh [DATASET_PATH] [DATASET_NAME] [DEVICE_NUM]"
+    exit 1
+fi
+
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+PATH1=$(get_real_path $1)
+
+if [ ! -d $PATH1 ]
+then
+    echo "error: DATASET_PATH=$PATH1 is not a directory"
+    exit 1
+fi
+
+export DEVICE_NUM=$3
+export RANK_SIZE=$3
+
+rm -rf ./train_parallel
+mkdir ./train_parallel
+cp ../*.py ./train_parallel
+cp *.sh ./train_parallel
+cp -r ../src ./train_parallel
+mkdir ./train_parallel/results
+mkdir ./train_parallel/results/ckpt
+mkdir ./train_parallel/results/predict
+mkdir ./train_parallel/results/loss_show
+mkdir ./train_parallel/results/fake_img
+cd ./train_parallel || exit
+
+if [ $2 == 'facades' ];
+then
+    mpirun -n $3 --output-filename log_output --merge-stderr-to-stdout --allow-run-as-root python train.py --device_target GPU \
+    --run_distribute 1 --device_num $3 --dataset_size 400 --train_data_dir $PATH1 --pad_mode REFLECT &> log &
+elif [ $2 == 'maps' ];
+then
+    mpirun --allow-run-as-root -n $3 --output-filename log_output --merge-stderr-to-stdout \
+    python train.py --device_target GPU --device_num $3 --dataset_size 1096 \
+    --run_distribute 1 --train_data_dir $PATH1 --pad_mode REFLECT &> log &
+fi

scripts/run_eval_ascend.sh

@@ -0,0 +1,50 @@
+#!/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 "======================================================================================================================================="
+
+if [ $# != 4 ]
+then
+    echo "Usage: bash run_eval_gpu.sh [DATASET_PATH] [DATASET_NAME] [CKPT_PATH] [RESULT_DIR]"
+    exit 1
+fi
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+PATH1=$(get_real_path $1)
+CKPT_PATH=$(get_real_path $3)
+RESULT_PATH=$4
+if [ ! -d $PATH1 ]
+then
+    echo "error: DATASET_PATH=$PATH1 is not a directory"
+    exit 1
+fi
+
+if [ $2 == 'facades' ]; then
+  python eval.py --device_target Ascend --device_id 0 --val_data_dir $PATH1 --ckpt $CKPT_PATH --predict_dir $RESULT_PATH --dataset_size 400
+elif [ $2 == 'maps' ]; then
+  python eval.py --device_target Ascend --device_id 0 --val_data_dir $PATH1 --ckpt $CKPT_PATH --predict_dir $RESULT_PATH --dataset_size 1096
+fi
+

scripts/run_eval_gpu.sh

@@ -0,0 +1,52 @@
+#!/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 GPU --device_id 0 --val_data_dir ./facades/test --ckpt ./results/ckpt/Generator_200.ckpt"
+echo "======================================================================================================================================="
+
+if [ $# != 4 ]
+then
+    echo "Usage: bash run_eval_gpu.sh [DATASET_PATH] [DATASET_NAME] [CKPT_PATH] [RESULT_DIR]"
+    exit 1
+fi
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+PATH1=$(get_real_path $1)
+CKPT_PATH=$(get_real_path $3)
+RESULT_PATH=$4
+
+if [ ! -d $PATH1 ]
+then
+    echo "error: DATASET_PATH=$PATH1 is not a directory"
+    exit 1
+fi
+
+if [ $2 == 'facades' ]; then
+  python eval.py --device_target GPU --device_id 0 --val_data_dir $PATH1 --ckpt $CKPT_PATH --predict_dir $RESULT_PATH --pad_mod REFLECT \
+  --dataset_size 400
+elif [ $2 == 'maps' ]; then
+  python eval.py --device_target GPU --device_id 0 --val_data_dir $PATH1 --ckpt $CKPT_PATH --predict_dir $RESULT_PATH --dataset_size 1096 \
+  --pad_mode REFLECT
+fi

scripts/run_infer_310.sh

@@ -0,0 +1,130 @@
+#!/bin/bash
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+if [[ $# -lt 4 || $# -gt 5 ]]; then
+    echo "Usage: bash run_infer_310.sh [MINDIR_PATH] [DATASET_PATH] [NEED_PREPROCESS] [DEVICE_TARGET] [DEVICE_ID]
+    DEVICE_TARGET must choose from ['GPU', 'CPU', 'Ascend']
+    NEED_PREPROCESS means weather need preprocess or not, it's value is 'y' or 'n'.
+    DEVICE_ID is optional, it can be set by environment variable device_id, otherwise the value is zero"
+exit 1
+fi
+
+get_real_path(){
+    if [ "${1:0:1}" == "/" ]; then
+        echo "$1"
+    else
+        echo "$(realpath -m $PWD/$1)"
+    fi
+}
+model=$(get_real_path $1)
+dataset_path=$(get_real_path $2)
+
+if [ "$3" == "y" ] || [ "$3" == "n" ];then
+    need_preprocess=$3
+else
+  echo "weather need preprocess or not, it's value must be in [y, n]"
+  exit 1
+fi
+if [ "$4" == "GPU" ] || [ "$4" == "CPU" ] || [ "$4" == "Ascend" ];then
+    device_target=$4
+else
+  echo "device_target must be in  ['GPU', 'CPU', 'Ascend']"
+  exit 1
+fi
+
+device_id=0
+if [ $# == 5 ]; then
+    device_id=$5
+fi
+
+echo "mindir name: "$model
+echo "dataset path: "$dataset_path
+echo "need preprocess: "$need_preprocess
+echo "device_target: "$device_target
+echo "device id: "$device_id
+
+export ASCEND_HOME=/usr/local/Ascend/
+if [ -d ${ASCEND_HOME}/ascend-toolkit ]; then
+    export PATH=$ASCEND_HOME/fwkacllib/bin:$ASCEND_HOME/fwkacllib/ccec_compiler/bin:$ASCEND_HOME/ascend-toolkit/latest/fwkacllib/ccec_compiler/bin:$ASCEND_HOME/ascend-toolkit/latest/atc/bin:$PATH
+    export LD_LIBRARY_PATH=$ASCEND_HOME/fwkacllib/lib64:/usr/local/lib:$ASCEND_HOME/ascend-toolkit/latest/atc/lib64:$ASCEND_HOME/ascend-toolkit/latest/fwkacllib/lib64:$ASCEND_HOME/driver/lib64:$ASCEND_HOME/add-ons:$LD_LIBRARY_PATH
+    export TBE_IMPL_PATH=$ASCEND_HOME/ascend-toolkit/latest/opp/op_impl/built-in/ai_core/tbe
+    export PYTHONPATH=$ASCEND_HOME/fwkacllib/python/site-packages:${TBE_IMPL_PATH}:$ASCEND_HOME/ascend-toolkit/latest/fwkacllib/python/site-packages:$PYTHONPATH
+    export ASCEND_OPP_PATH=$ASCEND_HOME/ascend-toolkit/latest/opp
+else
+    export PATH=$ASCEND_HOME/fwkacllib/bin:$ASCEND_HOME/fwkacllib/ccec_compiler/bin:$ASCEND_HOME/atc/ccec_compiler/bin:$ASCEND_HOME/atc/bin:$PATH
+    export LD_LIBRARY_PATH=$ASCEND_HOME/fwkacllib/lib64:/usr/local/lib:$ASCEND_HOME/atc/lib64:$ASCEND_HOME/acllib/lib64:$ASCEND_HOME/driver/lib64:$ASCEND_HOME/add-ons:$LD_LIBRARY_PATH
+    export PYTHONPATH=$ASCEND_HOME/fwkacllib/python/site-packages:$ASCEND_HOME/atc/python/site-packages:$PYTHONPATH
+    export ASCEND_OPP_PATH=$ASCEND_HOME/opp
+fi
+
+function preprocess_data()
+{
+    if [ -d preprocess_Result ]; then
+        rm -rf ./preprocess_Result
+    fi
+    mkdir preprocess_Result
+    python ../preprocess.py  --val_data_dir=$dataset_path --predict_dir=./preprocess_Result/ >& preprocess.log
+}
+
+function compile_app()
+{
+    cd ../ascend310_infer || exit
+    bash build.sh &> build.log
+}
+
+function infer()
+{
+    cd - || exit
+    if [ -d result_Files ]; then
+        rm -rf ./result_Files
+    fi
+    if [ -d time_Result ]; then
+        rm -rf ./time_Result
+    fi
+    mkdir result_Files
+    mkdir time_Result
+
+    ../ascend310_infer/out/main --mindir_path=$model --input0_path=./preprocess_Result/00_data --device_id=$device_id --fusion_switch_path=../ascend310_infer/fusion_switch.cfg &> infer.log
+
+}
+
+function generate_img()
+{
+    python ../postprocess.py  --predict_dir=./preprocess_Result/ --val_data_dir=./preprocess_Result/target_images_ids.npy --device_target=$device_target &> acc.log
+}
+
+if [ $need_preprocess == "y" ]; then
+    preprocess_data
+    if [ $? -ne 0 ]; then
+        echo "preprocess dataset failed"
+        exit 1
+    fi
+fi
+compile_app
+if [ $? -ne 0 ]; then
+    echo "compile app code failed"
+    exit 1
+fi
+infer
+if [ $? -ne 0 ]; then
+    echo " execute inference failed"
+    exit 1
+fi
+generate_img
+if [ $? -ne 0 ]; then
+    echo "generate images failed"
+    exit 1
+fi

scripts/run_train_ascend.sh

@@ -0,0 +1,48 @@
+#!/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 "====================================================================================================================="
+
+if [ $# != 2 ]
+then
+    echo "Usage: bash run_train_ascend.sh [DATASET_PATH] [DATASET_NAME]"
+    exit 1
+fi
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+PATH1=$(get_real_path $1)
+
+if [ ! -d $PATH1 ]
+then
+    echo "error: DATASET_PATH=$PATH1 is not a directory"
+    exit 1
+fi
+
+if [ $2 == 'facades' ]; then
+    python train.py --device_target Ascend --device_id 0 --dataset_size 400 --train_data_dir $PATH1
+elif [ $2 == 'maps' ]; then
+    python train.py --device_target Ascend --device_id 0 --dataset_size 1096 --train_data_dir $PATH1
+fi 

scripts/run_train_gpu.sh

@@ -0,0 +1,63 @@
+#!/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 GPU --device_id 0 --dataset_size 400 --train_data_dir ./facades/train"
+echo "====================================================================================================================="
+
+if [ $# != 2 ]
+then
+    echo "Usage: bash run_train_gpu.sh [DATASET_PATH] [DATASET_NAME]"
+    exit 1
+fi
+
+get_real_path(){
+  if [ "${1:0:1}" == "/" ]; then
+    echo "$1"
+  else
+    echo "$(realpath -m $PWD/$1)"
+  fi
+}
+
+PATH1=$(get_real_path $1)
+
+if [ ! -d $PATH1 ]
+then
+    echo "error: DATASET_PATH=$PATH1 is not a directory"
+    exit 1
+fi
+
+
+rm -rf ./train
+mkdir ./train
+mkdir ./train/results
+mkdir ./train/results/fake_img
+mkdir ./train/results/loss_show
+mkdir ./train/results/ckpt
+mkdir ./train/results/predict
+cp ./*.py ./train
+cp ./scripts/*.sh ./train
+cp -r ./src ./train
+cd ./train || exit
+
+if [ $2 == 'facades' ]; then
+    mpirun --allow-run-as-root -n 1 --output-filename log_output --merge-stderr-to-stdout \
+    python train.py --device_target GPU --device_id 0 --dataset_size 400 --train_data_dir $PATH1 --pad_mode REFLECT &> log &
+elif [ $2 == 'maps' ]; then
+    mpirun --allow-run-as-root -n 1 --output-filename log_output --merge-stderr-to-stdout \
+    python train.py --device_target GPU --device_id 0 --dataset_size 1096 --train_data_dir $PATH1 --pad_mode REFLECT &> log &
+fi

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

src/models/discriminator_model.py

@@ -0,0 +1,125 @@
+# 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
+from ..utils.config import get_args
+
+args = get_args()
+
+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 args.pad_mode == 'REFLECT':
+            pad_mode = "REFLECT"
+        elif args.pad_mode == "SYMMETRIC":
+            pad_mode = "SYMMETRIC"
+        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

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

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))

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
+
+        # Parallel processing
+        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

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

src/utils/config.py

@@ -0,0 +1,73 @@
+# 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
+
+
+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('--run_distribute', type=int, default=0, help='distributed training, default is 0.')
+    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('--pad_mode', type=str, default='CONSTANT', choices=('CONSTANT', 'REFLECT', 'SYMMETRIC'),
+                        help='scale images to this size, default is CONSTANT.')
+    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()
+    return args

src/utils/tools.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.
+# ===========================================================================
+
+"""
+    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
+
+plt.switch_backend('Agg')
+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))

train.py

@@ -1,113 +1,153 @@
-import argparse
-import matplotlib.pyplot as plt
+# 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 including distributed training
+'''
+
 import os
 import datetime
-from mindspore import context
 import mindspore.nn as nn
 from mindspore import Tensor
-import config
-from generator_model import Generator
-from discriminator_model import Discriminator
-from mindspore.train.callback import CheckpointConfig,ModelCheckpoint,_InternalCallbackParam
-from losses import WithLossCellGen,WithLossCellDis
-from pix2pix import Pix2Pix
-from dataset.pix2pix_dataset import MapDataset,create_dataset
-
-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("Loss")
-    plt.legend()
-    plt.savefig("./loss_show/{}.png".format(idx))
-
+from mindspore.train.serialization import save_checkpoint
+from mindspore.communication.management import init, get_rank, get_group_size
+from mindspore import context
+from mindspore.context import ParallelMode
+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__':
-    parser = argparse.ArgumentParser(description='MindSpore dcgan training')
-    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_id', type=int, default=4, help='device id of GPU or Ascend. (Default: 0)')
-    args = parser.parse_args()
-
-    context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target)
-    # context.set_context(device_id=args.device_id)
-
-    # 数据
-    dataset = MapDataset(root_dir="./maps/train/")
-    ds=create_dataset(dataset)
+
+    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())
-    # print("ds.dtype:", ds)
-    steps_per_epoch = ds.get_dataset_size()
-
-    # 定义网络结构
-    netD=Discriminator()
-    netG=Generator()
-
-    loss_fn1=nn.BCEWithLogitsLoss()
-    loss_fn2=nn.L1Loss()
-
-    netD_With_Loss=WithLossCellDis(netD,netG,loss_fn1)
-    netG_With_Loss=WithLossCellGen(netD,netG,loss_fn1,loss_fn2)
-
-    optimizerD = nn.Adam(netD.trainable_params(), learning_rate=0.0002, beta1=0.5, beta2=0.999)
-    optimizerG = nn.Adam(netG.trainable_params(), learning_rate=0.0002, beta1=0.5, beta2=0.999)
-
-    TrainOneStepDis=nn.TrainOneStepCell(netD_With_Loss,optimizerD,sens=1.0)
-    TrainOneStepGen=nn.TrainOneStepCell(netG_With_Loss,optimizerG,sens=1.0)
 
-    pix2pix_model=Pix2Pix(TrainOneStepDis,TrainOneStepGen)
-    pix2pix_model.set_train()
-
-    # save checkpoints
-    ckpt_config=CheckpointConfig(save_checkpoint_steps=steps_per_epoch,keep_checkpoint_max=10)
-    ckpt_cb=ModelCheckpoint(config=ckpt_config,directory='./ckpt',prefix='pix2pix_model')
-
-    cb_params = _InternalCallbackParam()
-    cb_params.train_network = pix2pix_model
-    cb_params.batch_num = steps_per_epoch
-    cb_params.epoch_num = config.NUM_EPOCHS   #100
-    # For each epoch
-    cb_params.cur_epoch_num = 0
-    cb_params.cur_step_num = 0
-
-    # Training Loop
+    steps_per_epoch = ds.get_dataset_size()
+    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
+    if args.device_target == 'Ascend':
+        if args.run_distribute:
+            print("Ascend distribute")
+            context.set_context(device_id=int(os.getenv("DEVICE_ID")))
+            device_num = args.device_num
+            context.reset_auto_parallel_context()
+            context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
+                                              device_num=device_num)
+            init()
+
+            rank = get_rank()
+        else:
+            context.set_context(device_id=args.device_id)
+    elif args.device_target == 'GPU':
+        if args.run_distribute:
+            print("GPU distribute")
+            init()
+            device_num = args.device_num
+            context.set_auto_parallel_context(device_num=get_group_size(),
+                                              parallel_mode=ParallelMode.DATA_PARALLEL,
+                                              gradients_mean=True)
+        else:
+            context.set_context(device_id=args.device_id)
+    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()
+
+    if not os.path.isdir(args.train_fakeimg_dir):
+        os.makedirs(args.train_fakeimg_dir)
+    if not os.path.isdir(args.loss_show_dir):
+        os.makedirs(args.loss_show_dir)
+    if not os.path.isdir(args.ckpt_dir):
+        os.makedirs(args.ckpt_dir)
+
+    # Training loop
     G_losses = []
     D_losses = []
 
-    data_loader = ds.create_dict_iterator(output_numpy=True, num_epochs=config.NUM_EPOCHS)
+    data_loader = ds.create_dict_iterator(output_numpy=True, num_epochs=args.epoch_num)
     print("Starting Training Loop...")
-    for epoch in range(config.NUM_EPOCHS):
-        # for each batch in the data_loader
-        for i,data in enumerate(data_loader):
-            input_image=Tensor(data["input_images"])
-            target_image=Tensor(data["target_images"])
-            netD_loss, netG_loss = pix2pix_model(input_image, target_image)
-            if i % 20 == 0:
+    if args.run_distribute:
+        rank = get_rank()
+    for epoch in range(args.epoch_num):
+        for i, data in enumerate(data_loader):
+            start_time = datetime.datetime.now()
+            input_image = Tensor(data["input_images"])
+            target_image = Tensor(data["target_images"])
+            dis_loss, gen_loss = train_net(input_image, target_image)
+            end_time = datetime.datetime.now()
+            delta = (end_time - start_time).microseconds
+            if i % 100 == 0:
                 print("================start===================")
-                print("Date time: ", datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
-                print("epoch: ", epoch, "/", config.NUM_EPOCHS)
+                print("Date time: ", start_time)
+                if args.run_distribute:
+                    print("Device ID :", str(rank))
+                print("ms per step :", delta/1000)
+                print("epoch: ", epoch + 1, "/", args.epoch_num)
                 print("step: ", i, "/", steps_per_epoch)
-                print("Dloss: ", netD_loss)
-                print("Gloss: ", netG_loss)
+                print("Dloss: ", dis_loss)
+                print("Gloss: ", gen_loss)
                 print("=================end====================")
-            D_losses.append(netD_loss.asnumpy())
-            G_losses.append(netG_loss.asnumpy())
-            cb_params.cur_step_num = cb_params.cur_step_num + 1
-
-        cb_params.cur_epoch_num = cb_params.cur_epoch_num + 1
-        ckpt_cb._save_ckpt(cb_params, True)
-        print("epoch-",epoch+1," saved")
-        save_losses(G_losses,D_losses,epoch+1)
-        print("epoch-",epoch+1," D&G_Losses saved")
-
-
-
-
-
-
-
 
+            # Save fake_imgs
+            if i == steps_per_epoch - 1:
+                fake_image = netG(input_image)
+                if args.run_distribute:
+                    fakeimg_path = args.train_fakeimg_dir + str(rank) + '/'
+                    if not os.path.isdir(fakeimg_path):
+                        os.makedirs(fakeimg_path)
+                    save_image(fake_image, fakeimg_path + str(epoch + 1))
+                else:
+                    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:
+            if args.run_distribute:
+                save_checkpoint_path = args.ckpt_dir + str(rank) + '/'
+                if not os.path.isdir(save_checkpoint_path):
+                    os.makedirs(save_checkpoint_path)
+                save_checkpoint(netG, os.path.join(save_checkpoint_path, f"Generator_{epoch+1}.ckpt"))
+            else:
+                save_checkpoint(netG, os.path.join(args.ckpt_dir, f"Generator_{epoch+1}.ckpt"))
+            print("ckpt generated from epoch", epoch + 1, "saved")

utils.py

@@ -1,52 +0,0 @@
-# import torch
-# import config
-# from torchvision.utils import save_image
-
-import mindspore
-import config
-from mindspore.ops import functional as F
-import matplotlib.image as mpimg
-import matplotlib.pyplot as plt
-
-
-def save_some_examples(gen, val_loader, epoch, folder):
-    x, y = next(iter(val_loader))
-    # x, y = x.to(config.DEVICE), y.to(config.DEVICE)
-    gen.eval()
-    # with torch.no_grad():
-    y_fake = gen(x)
-    y_fake = F.stop_gradient(y_fake)      #取消了torch.no_grad后加入stop_gradient()
-    y_fake = y_fake * 0.5 + 0.5  # remove normalization#
-
-    # save_image(y_fake, folder + f"/y_gen_{epoch}.png")
-    # save_image(x * 0.5 + 0.5, folder + f"/input_{epoch}.png")
-    # if epoch == 1:
-    #     save_image(y * 0.5 + 0.5, folder + f"/label_{epoch}.png")
-    mpimg.imsave(folder + f"/y_gen_{epoch}.png",y_fake)
-    mpimg.imsave(folder + f"/input_{epoch}.png",x * 0.5 + 0.5)
-    if epoch==1:
-        mpimg.imsave(folder + f"/label_{epoch}.png", y * 0.5 + 0.5)
-    gen.train()
-
-
-def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
-    print("=> Saving checkpoint")
-    checkpoint = [
-        {"state_dict": model.state_dict(),
-         "optimizer": optimizer.state_dict(), }
-    ]
-    mindspore.save_checkpoint(checkpoint, filename)
-
-
-def load_checkpoint(checkpoint_file, model, optimizer, lr):
-    print("=> Loading checkpoint")
-    checkpoint = mindspore.load_checkpoint(checkpoint_file)
-    model.load_param_into_net(checkpoint["state_dict"])
-    optimizer.load_param_into_net(checkpoint["optimizer"])
-
-    # If we don't do this then it will just have learning rate of old checkpoint
-    # and it will lead to many hours of debugging \:
-    for param_group in optimizer.param_groups:
-        param_group["lr"] = lr
-
-

utils/cityscapes_utils.py

@@ -1,82 +0,0 @@
-
-"""cityscape utils."""
-
-import numpy as np
-from PIL import Image
-
-# label name and RGB color map.
-label2color = {
-    'unlabeled': (0, 0, 0),
-    'ego vehicle': (0, 0, 0),
-    'rectification border': (0, 0, 0),
-    'out of roi': (0, 0, 0),
-    'static': (0, 0, 0),
-    'dynamic': (111, 74, 0),
-    'ground': (81, 0, 81),
-    'road': (128, 64, 128),
-    'sidewalk': (244, 35, 232),
-    'parking': (250, 170, 160),
-    'rail track': (230, 150, 140),
-    'building': (70, 70, 70),
-    'wall': (102, 102, 156),
-    'fence': (190, 153, 153),
-    'guard rail': (180, 165, 180),
-    'bridge': (150, 100, 100),
-    'tunnel': (150, 120, 90),
-    'pole': (153, 153, 153),
-    'polegroup': (153, 153, 153),
-    'traffic light': (250, 170, 30),
-    'traffic sign': (220, 220, 0),
-    'vegetation': (107, 142, 35),
-    'terrain': (152, 251, 152),
-    'sky': (70, 130, 180),
-    'person': (220, 20, 60),
-    'rider': (255, 0, 0),
-    'car': (0, 0, 142),
-    'truck': (0, 0, 70),
-    'bus': (0, 60, 100),
-    'caravan': (0, 0, 90),
-    'trailer': (0, 0, 110),
-    'train': (0, 80, 100),
-    'motorcycle': (0, 0, 230),
-    'bicycle': (119, 11, 32),
-    'license plate': (0, 0, 142)
-}
-
-def fast_hist(a, b, n):
-    k = np.where((a >= 0) & (a < n))[0]
-    bc = np.bincount(n * a[k].astype(int) + b[k], minlength=n**2)
-    if len(bc) != n**2:
-        # ignore this example if dimension mismatch
-        return 0
-    return bc.reshape(n, n)
-
-def get_scores(hist):
-    # Mean pixel accuracy
-    acc = np.diag(hist).sum() / (hist.sum() + 1e-12)
-    # Per class accuracy
-    cl_acc = np.diag(hist) / (hist.sum(1) + 1e-12)
-    # Per class IoU
-    iu = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist) + 1e-12)
-    return acc, np.nanmean(cl_acc), np.nanmean(iu), cl_acc, iu
-
-class CityScapes:
-    """CityScapes util class."""
-    def __init__(self):
-        self.classes = ['road', 'sidewalk', 'building', 'wall', 'fence',
-                        'pole', 'traffic light', 'traffic sign', 'vegetation', 'terrain',
-                        'sky', 'person', 'rider', 'car', 'truck',
-                        'bus', 'train', 'motorcycle', 'bicycle', 'unlabeled']
-        self.color_list = []
-        for name in self.classes:
-            self.color_list.append(label2color[name])
-        self.class_num = len(self.classes)
-
-    def get_id(self, img_path):
-        """Get train id by img"""
-        img = np.array(Image.open(img_path).convert("RGB"))
-        w, h, _ = img.shape
-        img_tile = np.tile(img, (1, 1, self.class_num)).reshape(w, h, self.class_num, 3)
-        diff = np.abs(img_tile - self.color_list).sum(axis=-1)
-        ids = diff.argmin(axis=-1)
-        return ids

utils/prepare_cityscapes_dataset.py

@@ -1,84 +0,0 @@
-# Copyright 2020 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.
-# ============================================================================
-"""prepare cityscapes dataset to cyclegan format"""
-import os
-import argparse
-import glob
-from PIL import Image
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--gtFine_dir', type=str, required=True,
-                    help='Path to the Cityscapes gtFine directory.')
-parser.add_argument('--leftImg8bit_dir', type=str, required=True,
-                    help='Path to the Cityscapes leftImg8bit_trainvaltest directory.')
-parser.add_argument('--output_dir', type=str, required=True,
-                    default='./cityscapes',
-                    help='Directory the output images will be written to.')
-opt = parser.parse_args()
-
-def load_resized_img(path):
-    """Load image with RGB and resize to (256, 256)"""
-    return Image.open(path).convert('RGB').resize((256, 256))
-
-def check_matching_pair(segmap_path, photo_path):
-    """Check the segment images and photo images are matched or not."""
-    segmap_identifier = os.path.basename(segmap_path).replace('_gtFine_color', '')
-    photo_identifier = os.path.basename(photo_path).replace('_leftImg8bit', '')
-
-    assert segmap_identifier == photo_identifier, \
-        f"[{segmap_path}] and [{photo_path}] don't seem to be matching. Aborting."
-
-
-def process_cityscapes(gtFine_dir, leftImg8bit_dir, output_dir, phase):
-    """Process citycapes dataset to cyclegan dataset format."""
-    save_phase = 'test' if phase == 'val' else 'train'
-    savedir = os.path.join(output_dir, save_phase)
-    os.makedirs(savedir + 'A', exist_ok=True)
-    os.makedirs(savedir + 'B', exist_ok=True)
-    print(f"Directory structure prepared at {output_dir}")
-
-    segmap_expr = os.path.join(gtFine_dir, phase) + "/*/*_color.png"
-    segmap_paths = glob.glob(segmap_expr)
-    segmap_paths = sorted(segmap_paths)
-
-    photo_expr = os.path.join(leftImg8bit_dir, phase) + "/*/*_leftImg8bit.png"
-    photo_paths = glob.glob(photo_expr)
-    photo_paths = sorted(photo_paths)
-
-    assert len(segmap_paths) == len(photo_paths), \
-        "{} images that match [{}], and {} images that match [{}]. Aborting.".format(
-            len(segmap_paths), segmap_expr, len(photo_paths), photo_expr)
-
-    for i, (segmap_path, photo_path) in enumerate(zip(segmap_paths, photo_paths)):
-        check_matching_pair(segmap_path, photo_path)
-        segmap = load_resized_img(segmap_path)
-        photo = load_resized_img(photo_path)
-
-        # data for cyclegan where the two images are stored at two distinct directories
-        savepath = os.path.join(savedir + 'A', f"{i + 1}.jpg")
-        photo.save(savepath)
-        savepath = os.path.join(savedir + 'B', f"{i + 1}.jpg")
-        segmap.save(savepath)
-
-        if i % (len(segmap_paths) // 10) == 0:
-            print("%d / %d: last image saved at %s, " % (i, len(segmap_paths), savepath))
-
-
-if __name__ == '__main__':
-    print('Preparing Cityscapes Dataset for val phase')
-    process_cityscapes(opt.gtFine_dir, opt.leftImg8bit_dir, opt.output_dir, "val")
-    print('Preparing Cityscapes Dataset for train phase')
-    process_cityscapes(opt.gtFine_dir, opt.leftImg8bit_dir, opt.output_dir, "train")
-    print('Done')

utils/reporter.py

@@ -1,60 +0,0 @@
-import logging
-import os
-import time
-from datetime import datetime
-from mindspore.train.serialization import save_checkpoint
-# from .tools import save_image
-
-class Reporter(logging.Logger):
-    """
-    This class includes several functions that can save images/checkpoints and print/save logging information.
-
-    Args:
-        args (class): Option class.
-    """
-    def __init__(self, args):
-        super(Reporter, self).__init__("cyclegan")
-        self.log_dir = os.path.join(args.outputs_dir, 'log')
-        self.imgs_dir = os.path.join(args.outputs_dir, "imgs")
-        self.ckpts_dir = os.path.join(args.outputs_dir, "ckpt")
-        if not os.path.exists(self.log_dir):
-            os.makedirs(self.log_dir, exist_ok=True)
-        if not os.path.exists(self.imgs_dir):
-            os.makedirs(self.imgs_dir, exist_ok=True)
-        if not os.path.exists(self.ckpts_dir):
-            os.makedirs(self.ckpts_dir, exist_ok=True)
-        self.rank = args.rank
-        self.save_checkpoint_epochs = args.save_checkpoint_epochs
-        self.save_imgs = args.save_imgs
-        # console handler
-        console = logging.StreamHandler()
-        console.setLevel(logging.INFO)
-        formatter = logging.Formatter('%(message)s')
-        console.setFormatter(formatter)
-        self.addHandler(console)
-        # file handler
-        log_name = datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S') + '_rank_{}.log'.format(self.rank)
-        self.log_fn = os.path.join(self.log_dir, log_name)
-        fh = logging.FileHandler(self.log_fn)
-        fh.setLevel(logging.INFO)
-        fh.setFormatter(formatter)
-        self.addHandler(fh)
-        self.save_args(args)
-        self.step = 0
-        self.epoch = 0
-        self.dataset_size = args.dataset_size
-        self.print_iter = args.print_iter
-        self.G_loss = []
-        self.D_loss = []
-
-        def info(self, msg, *args, **kwargs):
-            if self.isEnabledFor(logging.INFO):
-                self._log(logging.INFO, msg, args, **kwargs)
-
-        def epoch_start(self):
-            self.step_start_time = time.time()
-            self.epoch_start_time = time.time()
-            self.step = 0
-            self.epoch += 1
-            self.G_loss = []
-            self.D_loss = []