From c6a08bc7ff659f12137eaea959f7767aa23f9773 Mon Sep 17 00:00:00 2001 From: maijianqiang Date: Sat, 19 Jun 2021 17:21:04 +0800 Subject: [PATCH] merge ctpn --- model_zoo/official/cv/ctpn/README.md | 89 ++++++++- .../official/cv/ctpn/default_config.yaml | 178 ++++++++++++++++++ model_zoo/official/cv/ctpn/eval.py | 51 +++-- model_zoo/official/cv/ctpn/export.py | 43 +++-- .../scripts/run_distribute_train_ascend.sh | 12 +- .../cv/ctpn/scripts/run_eval_ascend.sh | 11 +- .../scripts/run_standalone_train_ascend.sh | 11 +- model_zoo/official/cv/ctpn/src/config.py | 138 -------------- model_zoo/official/cv/ctpn/src/dataset.py | 3 +- model_zoo/official/cv/ctpn/src/eval_utils.py | 9 +- .../cv/ctpn/src/model_utils/__init__.py | 0 .../cv/ctpn/src/model_utils/config.py | 130 +++++++++++++ .../cv/ctpn/src/model_utils/device_adapter.py | 26 +++ .../cv/ctpn/src/model_utils/local_adapter.py | 36 ++++ .../cv/ctpn/src/model_utils/moxing_adapter.py | 124 ++++++++++++ .../cv/ctpn/src/text_connector/detector.py | 2 +- .../src/text_connector/get_successions.py | 2 +- model_zoo/official/cv/ctpn/train.py | 134 +++++++------ 18 files changed, 749 insertions(+), 250 deletions(-) create mode 100644 model_zoo/official/cv/ctpn/default_config.yaml delete mode 100644 model_zoo/official/cv/ctpn/src/config.py create mode 100644 model_zoo/official/cv/ctpn/src/model_utils/__init__.py create mode 100644 model_zoo/official/cv/ctpn/src/model_utils/config.py create mode 100644 model_zoo/official/cv/ctpn/src/model_utils/device_adapter.py create mode 100644 model_zoo/official/cv/ctpn/src/model_utils/local_adapter.py create mode 100644 model_zoo/official/cv/ctpn/src/model_utils/moxing_adapter.py diff --git a/model_zoo/official/cv/ctpn/README.md b/model_zoo/official/cv/ctpn/README.md index 5cc8db418bc..f77668253d5 100644 --- a/model_zoo/official/cv/ctpn/README.md +++ b/model_zoo/official/cv/ctpn/README.md @@ -90,7 +90,11 @@ Here we used 6 datasets for training, and 1 datasets for Evaluation. │   │   ├── proposal_generator.py # proposla generator │   │   ├── rpn.py # region-proposal network │   │   └── vgg16.py # backbone - │   ├── config.py # training configuration + │ ├── model_utils + │ │ ├──config.py // Parameter config + │ │ ├──moxing_adapter.py // modelarts device configuration + │ │ ├──device_adapter.py // Device Config + │ │ ├──local_adapter.py // local device config │   ├── convert_icdar2015.py # convert icdar2015 dataset label │   ├── convert_svt.py # convert svt label │   ├── create_dataset.py # create mindrecord dataset @@ -109,6 +113,7 @@ Here we used 6 datasets for training, and 1 datasets for Evaluation. ├──postprogress.py # post process for 310 inference ├──export.py # script to export AIR,MINDIR model └── train.py # train net + ├── default_config.yaml # config file ``` @@ -221,6 +226,65 @@ Training result will be stored in the example path. Checkpoints will be stored a 424 epoch: 3 step: 229 ,rpn_loss: 0.00910, rpn_cls_loss: 0.00385, rpn_reg_loss: 0.00175, ``` +- running on ModelArts +- If you want to train the model on modelarts, you can refer to the [official guidance document] of modelarts (https://support.huaweicloud.com/modelarts/) + +```python +# Example of using distributed training dpn on modelarts : +# Data set storage method + +# ├── ctpn_dataset # dir +# ├──train # train dir +# ├── pretrain # pretrain dataset dir +# ├── finetune # finetune dataset dir +# ├── backbone # predtrained dir if exists +# ├── eval # eval dir +# ├── ICDAR2013 # ICDAR2013 img dir +# ├── checkpoint # ckpt files dir +# ├── test # ckpt files dir +# ├── ctpn_test.mindrecord # test img of mindrecord +# ├── ctpn_test.mindrecord.db # test img of mindrecord.db + +# (1) Choose either a (modify yaml file parameters) or b (modelArts create training job to modify parameters) 。 +# a. set "enable_modelarts=True" 。 +# set "run_distribute=True" +# set "save_checkpoint_path=/cache/train/checkpoint/" +# set "finetune_dataset_file=/cache/data/finetune/ctpn_finetune.mindrecord0" +# set "pretrain_dataset_file=/cache/data/finetune/ctpn_pretrain.mindrecord0" +# set "task_type=Pretraining" or task_type=Finetune +# set "pre_trained=/cache/data/backbone/pred file name" Without pre-training weights pre_trained="" +# +# b. add "enable_modelarts=True" Parameters are on the interface of modearts。 +# Set the parameters required by method a on the modelarts interface +# Note: The path parameter does not need to be quoted + +# (2) Set the path of the network configuration file "_config_path=/The path of config in default_config.yaml/" +# (3) Set the code path on the modelarts interface "/path/ctpn"。 +# (4) Set the model's startup file on the modelarts interface "train.py" 。 +# (5) Set the data path of the model on the modelarts interface ".../ctpn_dataset/train"(choices ctpn_dataset/train Folder path) , +# The output path of the model "Output file path" and the log path of the model "Job log path" 。 +# (6) start trainning the model。 + +# Example of using model inference on modelarts +# (1) Place the trained model to the corresponding position of the bucket。 +# (2) chocie a or b。 +# a. set "enable_modelarts=True" 。 +# set "dataset_path=/cache/data/test/ctpn_test.mindrecord" +# set "img_dir=/cache/data/ICDAR2013/test" +# set "checkpoint_path=/cache/data/checkpoint/checkpoint file name" + +# b. Add "enable_modelarts=True" parameter on the interface of modearts。 +# Set the parameters required by method a on the modelarts interface +# Note: The path parameter does not need to be quoted + +# (3) Set the path of the network configuration file "_config_path=/The path of config in default_config.yaml/" +# (4) Set the code path on the modelarts interface "/path/ctpn"。 +# (5) Set the model's startup file on the modelarts interface "eval.py" 。 +# (6) Set the data path of the model on the modelarts interface ".../ctpn_dataset/eval"(choices FSNS/eval Folder path) , +# The output path of the model "Output file path" and the log path of the model "Job log path" 。 +# (7) Start model inference。 +``` + ## [Eval process](#contents) ### Usage @@ -264,7 +328,28 @@ Evaluation result will be stored in the example path, you can find result like t ## Model Export ```shell -python export.py --ckpt_file [CKPT_PATH] --device_target [DEVICE_TARGET] --file_format[EXPORT_FORMAT] +python export.py --ckpt_file [CKPT_PATH] --file_format[EXPORT_FORMAT] +``` + +- Export MindIR on Modelarts + +```Modelarts +Export MindIR example on ModelArts +Data storage method is the same as training +# (1) Choose either a (modify yaml file parameters) or b (modelArts create training job to modify parameters)。 +# a. set "enable_modelarts=True" +# set "file_name=/cache/train/cnnctc" +# set "file_format=MINDIR" +# set "ckpt_file=/cache/data/checkpoint file name" + +# b. Add "enable_modelarts=True" parameter on the interface of modearts。 +# Set the parameters required by method a on the modelarts interface +# Note: The path parameter does not need to be quoted +# (2)Set the path of the network configuration file "_config_path=/The path of config in default_config.yaml/" +# (3) Set the code path on the modelarts interface "/path/ctpn"。 +# (4) Set the model's startup file on the modelarts interface "export.py" 。 +# (5) Set the data path of the model on the modelarts interface ".../ctpn_dataset/eval/checkpoint"(choices CNNCTC_Data/eval/checkpoint Folder path) , +# The output path of the model "Output file path" and the log path of the model "Job log path" 。 ``` `EXPORT_FORMAT` should be in ["AIR", "MINDIR"] diff --git a/model_zoo/official/cv/ctpn/default_config.yaml b/model_zoo/official/cv/ctpn/default_config.yaml new file mode 100644 index 00000000000..bc3bc7c3734 --- /dev/null +++ b/model_zoo/official/cv/ctpn/default_config.yaml @@ -0,0 +1,178 @@ +# Builtin Configurations(DO NOT CHANGE THESE CONFIGURATIONS unlesee you know exactly what you are doing) +enable_modelarts: False +# url for modelarts +data_url: "" +train_url: "" +checkpoint_url: "" +# path for local +data_path: "/cache/data" +output_path: "/cache/train" +load_path: "/cache/checkpoint_path" +device_target: "Ascend" +enable_profiling: False +modelarts_home: "/home/work/user-job-dir" +object_name: "ctpn" + +# ====================================================================================== +# common options +run_distribute: False + +# ====================================================================================== +# Training options +img_width: 960 +img_height: 576 +keep_ratio: False +flip_ratio: 0.0 +photo_ratio: 0.0 +expand_ratio: 1.0 + +# anchor +num_anchors: 14 +anchor_base: 16 +anchor_height: [2, 4, 7, 11, 16, 23, 33, 48, 68, 97, 139, 198, 283, 406] +anchor_width: [16] + +# rpn +rpn_in_channels: 256 +rpn_feat_channels: 512 +rpn_loss_cls_weight: 1.0 +rpn_loss_reg_weight: 3.0 +rpn_cls_out_channels: 2 + +# bbox_assign_sampler +neg_iou_thr: 0.5 +pos_iou_thr: 0.7 +min_pos_iou: 0.001 +num_gts: 256 +num_expected_neg: 512 +num_expected_pos: 256 + +# proposal +activate_num_classes: 2 +use_sigmoid_cls: False + +# train proposal +rpn_proposal_nms_across_levels: False +rpn_proposal_nms_pre: 2000 +rpn_proposal_nms_post: 1000 +rpn_proposal_max_num: 1000 +rpn_proposal_nms_thr: 0.7 +rpn_proposal_min_bbox_size: 8 + +# rnn structure +input_size: 512 +hidden_size: 128 + +# training +warmup_mode: "linear" +# batch_size only support 1 +batch_size: 1 +momentum: 0.9 +save_checkpoint: True +save_checkpoint_epochs: 10 +keep_checkpoint_max: 5 +save_checkpoint_path: "./" +use_dropout: False +loss_scale: 1 +weight_decay: 1e-4 +pre_trained: "" +task_type: "Pretraining" +run_eval: False +save_best_ckpt: True +eval_image_path: "" +eval_dataset_path: "" +eval_start_epoch: 10 +eval_interval: 10 + +# text proposal connection +max_horizontal_gap: 60 +text_proposals_min_scores: 0.7 +text_proposals_nms_thresh: 0.2 +min_v_overlaps: 0.7 +min_size_sim: 0.7 +min_ratio: 0.5 +line_min_score: 0.9 +text_proposals_width: 16 +min_num_proposals: 2 + +# create dataset +coco_root: "" +coco_train_data_type: "" +cocotext_json: "" +icdar11_train_path: [] +icdar13_train_path: [] +icdar15_train_path: [] +icdar13_test_path: [] +flick_train_path: [] +svt_train_path: [] +pretrain_dataset_path: "" +finetune_dataset_path: "" +test_dataset_path: "" + +# training dataset +pretraining_dataset_file: "" +finetune_dataset_file: "" + +# pretrain lr +pre_base_lr: 0.0009 +pre_warmup_step: 30000 +pre_warmup_ratio: 1/3 +pre_total_epoch: 100 + +# finetune lr +fine_base_lr: 0.0005 +fine_warmup_step: 300 +fine_warmup_ratio: 1/3 +fine_total_epoch: 50 + +# ====================================================================================== +# Eval options +rpn_nms_pre: 2000 +rpn_nms_post: 1000 +rpn_max_num: 1000 +rpn_nms_thr: 0.7 +rpn_min_bbox_min_size: 8 +test_iou_thr: 0.7 +test_max_per_img: 1000 +test_batch_size: 1 +use_python_proposal: False +dataset_path: "" +image_path: "" +checkpoint_path: "" +img_dir: "" + +# ====================================================================================== +# export options +device_id: 0 +file_name: "cnnctc" +file_format: "MINDIR" +ckpt_file: "" + +# ====================================================================================== +# 310 infer + +--- +# Help description for each configuration +enable_modelarts: "Whether training on modelarts default: False" +data_url: "Url for modelarts" +train_url: "Url for modelarts" +data_path: "The location of input data" +output_pah: "The location of the output file" +device_target: "device id of GPU or Ascend. (Default: None)" +enable_profiling: "Whether enable profiling while training default: False" +file_name: "CNN&CTC output air name" +file_format: "choices [AIR, MINDIR]" +ckpt_file: "CNN&CTC ckpt file" +run_distribute: "Run distribute, default: false." +pre_trained: "Pretrained file path." +device_id: "Device id, default: 0." +task_type: "Pretraining choices [Pretraining, Finetune]" +run_eval: "Run evaluation when training, default is False." +save_best_ckpt: "Save best checkpoint when run_eval is True, default is True." +eval_image_path: "eval image path, when run_eval is True, eval_image_path should be set." +eval_dataset_path: "eval dataset path, when run_eval is True, eval_dataset_path should be set." +eval_start_epoch: "Evaluation start epoch when run_eval is True, default is 10." +eval_interval: "Evaluation interval when run_eval is True, default is 10." +dataset_path: "Dataset path." +image_path: "Image path." +checkpoint_path: "Checkpoint file path." diff --git a/model_zoo/official/cv/ctpn/eval.py b/model_zoo/official/cv/ctpn/eval.py index 77636d2e569..0ffba5bf902 100644 --- a/model_zoo/official/cv/ctpn/eval.py +++ b/model_zoo/official/cv/ctpn/eval.py @@ -14,35 +14,52 @@ # ============================================================================ """Evaluation for CTPN""" -import argparse -from mindspore import context +import os from mindspore.train.serialization import load_checkpoint, load_param_into_net from mindspore.common import set_seed from src.ctpn import CTPN -from src.config import config from src.dataset import create_ctpn_dataset from src.eval_utils import eval_for_ctpn +from src.model_utils.config import config +from src.model_utils.moxing_adapter import moxing_wrapper + + set_seed(1) -parser = argparse.ArgumentParser(description="CTPN evaluation") -parser.add_argument("--dataset_path", type=str, default="", help="Dataset path.") -parser.add_argument("--image_path", type=str, default="", help="Image path.") -parser.add_argument("--checkpoint_path", type=str, default="", help="Checkpoint file path.") -parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.") -args_opt = parser.parse_args() -context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id) -def ctpn_infer_test(dataset_path='', ckpt_path='', img_dir=''): - """ctpn infer.""" - print("ckpt path is {}".format(ckpt_path)) - ds = create_ctpn_dataset(dataset_path, batch_size=config.test_batch_size, repeat_num=1, is_training=False) +context.set_context(mode=context.GRAPH_MODE, device_target=config.device_target, device_id=get_device_id) + + +def modelarts_pre_process(): + pass + + +def modelarts_post_process(): + local_path = os.path.join(config.modelarts_home, config.object_name) + basename = os.path.basename(config.checkpoint_path) + copy_label = 'cd {}&&zip submit_{}.zip ./submit *.txt'.format(local_path, basename) + os.system(copy_label) + os.system('cd {}&&sed -i "s/\r//" scripts/eval_res.sh'.format(local_path)) + os.system('cd {}&& sh scripts/eval_res.sh'.format(local_path)) + + +@moxing_wrapper(pre_process=modelarts_pre_process, post_process=modelarts_post_process) +def ctpn_infer_test(): + config.feature_shapes = [config.img_height // 16, config.img_width // 16] + config.num_bboxes = (config.img_height // 16) * (config.img_width // 16) * config.num_anchors + config.num_step = config.img_width // 16 + config.rnn_batch_size = config.img_height // 16 + + print("ckpt path is {}".format(config.checkpoint_path)) + ds = create_ctpn_dataset(config.dataset_path, batch_size=config.test_batch_size, repeat_num=1, is_training=False) total = ds.get_dataset_size() print("eval dataset size is {}".format(total)) net = CTPN(config, batch_size=config.test_batch_size, is_training=False) - param_dict = load_checkpoint(ckpt_path) + param_dict = load_checkpoint(config.checkpoint_path) load_param_into_net(net, param_dict) net.set_train(False) - eval_for_ctpn(net, ds, img_dir) + eval_for_ctpn(net, ds, config.img_dir) + if __name__ == '__main__': - ctpn_infer_test(args_opt.dataset_path, args_opt.checkpoint_path, img_dir=args_opt.image_path) + ctpn_infer_test() diff --git a/model_zoo/official/cv/ctpn/export.py b/model_zoo/official/cv/ctpn/export.py index c886ea1b62a..0aa61641b0c 100644 --- a/model_zoo/official/cv/ctpn/export.py +++ b/model_zoo/official/cv/ctpn/export.py @@ -13,32 +13,35 @@ # limitations under the License. # ============================================================================ """export checkpoint file into air, onnx, mindir models""" -import argparse import numpy as np - import mindspore as ms from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context - from src.ctpn import CTPN_Infer -from src.config import config +from src.model_utils.config import config +from src.model_utils.moxing_adapter import moxing_wrapper -parser = argparse.ArgumentParser(description='fasterrcnn_export') -parser.add_argument("--device_id", type=int, default=0, help="Device id") -parser.add_argument("--file_name", type=str, default="ctpn", help="output file name.") -parser.add_argument("--file_format", type=str, choices=["AIR", "MINDIR"], default="MINDIR", help="file format") -parser.add_argument("--device_target", type=str, choices=["Ascend", "GPU", "CPU"], default="Ascend", - help="device target") -parser.add_argument('--ckpt_file', type=str, default='', help='ctpn ckpt file.') -args = parser.parse_args() -context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target) -if args.device_target == "Ascend": - context.set_context(device_id=args.device_id) +context.set_context(mode=context.GRAPH_MODE, device_target=config.device_target) + + +if config.device_target == "Ascend": + context.set_context(device_id=config.device_id) + + +def modelarts_pre_process(): + pass + + +@moxing_wrapper(pre_process=modelarts_pre_process) +def model_export(): + config.feature_shapes = [config.img_height // 16, config.img_width // 16] + config.num_bboxes = (config.img_height // 16) * (config.img_width // 16) * config.num_anchors + config.num_step = config.img_width // 16 + config.rnn_batch_size = config.img_height // 16 -if __name__ == '__main__': net = CTPN_Infer(config=config, batch_size=config.test_batch_size) - param_dict = load_checkpoint(args.ckpt_file) + param_dict = load_checkpoint(config.ckpt_file) param_dict_new = {} for key, value in param_dict.items(): @@ -48,4 +51,8 @@ if __name__ == '__main__': img = Tensor(np.zeros([config.test_batch_size, 3, config.img_height, config.img_width]), ms.float16) - export(net, img, file_name=args.file_name, file_format=args.file_format) + export(net, img, file_name=config.file_name, file_format=config.file_format) + + +if __name__ == '__main__': + model_export() diff --git a/model_zoo/official/cv/ctpn/scripts/run_distribute_train_ascend.sh b/model_zoo/official/cv/ctpn/scripts/run_distribute_train_ascend.sh index 7669552e361..abfce7cbde7 100644 --- a/model_zoo/official/cv/ctpn/scripts/run_distribute_train_ascend.sh +++ b/model_zoo/official/cv/ctpn/scripts/run_distribute_train_ascend.sh @@ -16,7 +16,7 @@ if [ $# -ne 3 ] then - echo "Usage: sh run_distribute_train_ascend.sh [RANK_TABLE_FILE] [TASK_TYPE] [PRETRAINED_PATH]" + echo "Usage: sh scripts/run_distribute_train_ascend.sh [RANK_TABLE_FILE] [TASK_TYPE] [PRETRAINED_PATH]" exit 1 fi @@ -57,13 +57,13 @@ do rm -rf ./train_parallel$i mkdir ./train_parallel$i cp ./*.py ./train_parallel$i - cp ./*.zip ./train_parallel$i - cp ../*.py ./train_parallel$i - cp *.sh ./train_parallel$i - cp -r ../src ./train_parallel$i + cp ./*.py ./train_parallel$i + cp ./*yaml ./train_parallel$i + cp -r ./scripts/ ./train_parallel$i + cp -r ./src ./train_parallel$i cd ./train_parallel$i || exit echo "start training for rank $RANK_ID, device $DEVICE_ID" env > env.log - python train.py --device_id=$i --rank_id=$i --run_distribute=True --device_num=$DEVICE_NUM --task_type=$TASK_TYPE --pre_trained=$PATH2 &> log & + python train.py --run_distribute=True --task_type=$TASK_TYPE --pre_trained=$PATH2 &> log & cd .. done diff --git a/model_zoo/official/cv/ctpn/scripts/run_eval_ascend.sh b/model_zoo/official/cv/ctpn/scripts/run_eval_ascend.sh index 9848459fb43..01259aeb933 100644 --- a/model_zoo/official/cv/ctpn/scripts/run_eval_ascend.sh +++ b/model_zoo/official/cv/ctpn/scripts/run_eval_ascend.sh @@ -16,7 +16,7 @@ if [ $# != 3 ] then - echo "Usage: sh run_eval_ascend.sh [IMAGE_PATH] [DATASET_PATH] [CHECKPOINT_PATH]" + echo "Usage: sh scripts/run_eval_ascend.sh [IMAGE_PATH] [DATASET_PATH] [CHECKPOINT_PATH]" exit 1 fi @@ -64,14 +64,15 @@ do rm -rf ./eval fi mkdir ./eval - cp ../*.py ./eval - cp *.sh ./eval - cp -r ../src ./eval + cp ./*.py ./eval + cp -r ./scripts ./eval + cp -r ./src ./eval + cp ./*yaml ./eval cd ./eval || exit env > env.log CHECKPOINT_FILE_PATH=$file echo "start eval for checkpoint file: ${CHECKPOINT_FILE_PATH}" - python eval.py --device_id=$DEVICE_ID --image_path=$IMAGE_PATH --dataset_path=$DATASET_PATH --checkpoint_path=$CHECKPOINT_FILE_PATH &> log + python eval.py --image_path=$IMAGE_PATH --dataset_path=$DATASET_PATH --checkpoint_path=$CHECKPOINT_FILE_PATH &> log echo "end eval for checkpoint file: ${CHECKPOINT_FILE_PATH}" cd ./submit || exit file_base_name=$(basename $file) diff --git a/model_zoo/official/cv/ctpn/scripts/run_standalone_train_ascend.sh b/model_zoo/official/cv/ctpn/scripts/run_standalone_train_ascend.sh index 394ea845314..f6758c72a03 100644 --- a/model_zoo/official/cv/ctpn/scripts/run_standalone_train_ascend.sh +++ b/model_zoo/official/cv/ctpn/scripts/run_standalone_train_ascend.sh @@ -22,7 +22,7 @@ echo "It is better to use absolute path." echo "==============================================================================================================" if [ $# -ne 3 ] then - echo "Usage: sh run_standalone_train_ascend.sh [TASK_TYPE] [PRETRAINED_PATH] [DEVICE_ID]" + echo "Usage: sh scripts/run_standalone_train_ascend.sh [TASK_TYPE] [PRETRAINED_PATH] [DEVICE_ID]" exit 1 fi @@ -51,11 +51,12 @@ export RANK_SIZE=1 rm -rf ./train mkdir ./train -cp ../*.py ./train -cp *.sh ./train -cp -r ../src ./train +cp ./*.py ./train +cp -r ./scripts ./train +cp ./*yaml ./train +cp -r ./src ./train cd ./train || exit echo "start training for device $DEVICE_ID" env > env.log -python train.py --device_id=$DEVICE_ID --task_type=$TASK_TYPE --pre_trained=$PRETRAINED_PATH &> log & +python train.py --task_type=$TASK_TYPE --pre_trained=$PRETRAINED_PATH &> log & cd .. diff --git a/model_zoo/official/cv/ctpn/src/config.py b/model_zoo/official/cv/ctpn/src/config.py deleted file mode 100644 index 45f3aed9035..00000000000 --- a/model_zoo/official/cv/ctpn/src/config.py +++ /dev/null @@ -1,138 +0,0 @@ -# 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. -# ============================================================================ -"""Network parameters.""" -from easydict import EasyDict -pretrain_config = EasyDict({ - # LR - "base_lr": 0.0009, - "warmup_step": 30000, - "warmup_ratio": 1/3.0, - "total_epoch": 100, -}) -finetune_config = EasyDict({ - # LR - "base_lr": 0.0005, - "warmup_step": 300, - "warmup_ratio": 1/3.0, - "total_epoch": 50, -}) - -config_default = EasyDict({ - "img_width": 960, - "img_height": 576, - "keep_ratio": False, - "flip_ratio": 0.0, - "photo_ratio": 0.0, - "expand_ratio": 1.0, - - # anchor - "num_anchors": 14, - "anchor_base": 16, - "anchor_height": [2, 4, 7, 11, 16, 23, 33, 48, 68, 97, 139, 198, 283, 406], - "anchor_width": [16], - - # rpn - "rpn_in_channels": 256, - "rpn_feat_channels": 512, - "rpn_loss_cls_weight": 1.0, - "rpn_loss_reg_weight": 3.0, - "rpn_cls_out_channels": 2, - - # bbox_assign_sampler - "neg_iou_thr": 0.5, - "pos_iou_thr": 0.7, - "min_pos_iou": 0.001, - "num_gts": 256, - "num_expected_neg": 512, - "num_expected_pos": 256, - - #proposal - "activate_num_classes": 2, - "use_sigmoid_cls": False, - - # train proposal - "rpn_proposal_nms_across_levels": False, - "rpn_proposal_nms_pre": 2000, - "rpn_proposal_nms_post": 1000, - "rpn_proposal_max_num": 1000, - "rpn_proposal_nms_thr": 0.7, - "rpn_proposal_min_bbox_size": 8, - - # rnn structure - "input_size": 512, - "hidden_size": 128, - - # training - "warmup_mode": "linear", - # batch_size only support 1 - "batch_size": 1, - "momentum": 0.9, - "save_checkpoint": True, - "save_checkpoint_epochs": 10, - "keep_checkpoint_max": 5, - "save_checkpoint_path": "./", - "use_dropout": False, - "loss_scale": 1, - "weight_decay": 1e-4, - - # test proposal - "rpn_nms_pre": 2000, - "rpn_nms_post": 1000, - "rpn_max_num": 1000, - "rpn_nms_thr": 0.7, - "rpn_min_bbox_min_size": 8, - "test_iou_thr": 0.7, - "test_max_per_img": 100, - "test_batch_size": 1, - "use_python_proposal": False, - - # text proposal connection - "max_horizontal_gap": 60, - "text_proposals_min_scores": 0.7, - "text_proposals_nms_thresh": 0.2, - "min_v_overlaps": 0.7, - "min_size_sim": 0.7, - "min_ratio": 0.5, - "line_min_score": 0.9, - "text_proposals_width": 16, - "min_num_proposals": 2, - - # create dataset - "coco_root": "", - "coco_train_data_type": "", - "cocotext_json": "", - "icdar11_train_path": [], - "icdar13_train_path": [], - "icdar15_train_path": [], - "icdar13_test_path": [], - "flick_train_path": [], - "svt_train_path": [], - "pretrain_dataset_path": "", - "finetune_dataset_path": "", - "test_dataset_path": "", - - # training dataset - "pretraining_dataset_file": "", - "finetune_dataset_file": "" -}) - -config_add = { - "feature_shapes": (config_default["img_height"] // 16, config_default["img_width"] // 16), - "num_bboxes": (config_default["img_height"] // 16) * \ - (config_default["img_width"] // 16) *config_default["num_anchors"], - "num_step": config_default["img_width"] // 16, - "rnn_batch_size": config_default["img_height"] // 16 -} -config = EasyDict({**config_default, **config_add}) diff --git a/model_zoo/official/cv/ctpn/src/dataset.py b/model_zoo/official/cv/ctpn/src/dataset.py index a7246e14392..f1f91e668ca 100644 --- a/model_zoo/official/cv/ctpn/src/dataset.py +++ b/model_zoo/official/cv/ctpn/src/dataset.py @@ -22,7 +22,8 @@ import mindspore.dataset as de import mindspore.dataset.vision.c_transforms as C import mindspore.dataset.transforms.c_transforms as CC import mindspore.common.dtype as mstype -from src.config import config +from src.model_utils.config import config + class PhotoMetricDistortion: """Photo Metric Distortion""" diff --git a/model_zoo/official/cv/ctpn/src/eval_utils.py b/model_zoo/official/cv/ctpn/src/eval_utils.py index 56f64969b4d..cbd0693d6e8 100644 --- a/model_zoo/official/cv/ctpn/src/eval_utils.py +++ b/model_zoo/official/cv/ctpn/src/eval_utils.py @@ -16,9 +16,10 @@ import os import subprocess import numpy as np -from src.config import config +from src.model_utils.config import config from src.text_connector.detector import detect + def exec_shell_cmd(cmd): sub = subprocess.Popen(args="{}".format(cmd), shell=True, stdin=subprocess.PIPE, \ stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True) @@ -34,11 +35,15 @@ def get_eval_result(): hmean = exec_shell_cmd(get_eval_output) return float(hmean) + def eval_for_ctpn(network, dataset, eval_image_path): network.set_train(False) eval_iter = 0 img_basenames = [] - output_dir = os.path.join(os.getcwd(), "submit") + local_path = os.getcwd() + if config.enable_modelarts: + local_path = os.path.join(config.modelarts_home, config.object_name) + output_dir = os.path.join(local_path, "submit") if not os.path.exists(output_dir): os.mkdir(output_dir) for file in os.listdir(eval_image_path): diff --git a/model_zoo/official/cv/ctpn/src/model_utils/__init__.py b/model_zoo/official/cv/ctpn/src/model_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/model_zoo/official/cv/ctpn/src/model_utils/config.py b/model_zoo/official/cv/ctpn/src/model_utils/config.py new file mode 100644 index 00000000000..efc856cf0cf --- /dev/null +++ b/model_zoo/official/cv/ctpn/src/model_utils/config.py @@ -0,0 +1,130 @@ +# 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 WARRANT IES OR CONITTONS OF ANY KIND， either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ==================================================================================== + +"""Parse arguments""" +import os +import ast +import argparse +from pprint import pprint, pformat +import yaml + + +_config_path = '../../default_config.yaml' + + +class Config: + """ + Configuration namespace. Convert dictionary to members + """ + def __init__(self, cfg_dict): + for k, v in cfg_dict.items(): + if isinstance(v, (list, tuple)): + setattr(self, k, [Config(x) if isinstance(x, dict) else x for x in v]) + else: + setattr(self, k, Config(v) if isinstance(v, dict) else v) + + def __str__(self): + return pformat(self.__dict__) + + def __repr__(self): + return self.__str__() + + +def parse_cli_to_yaml(parser, cfg, helper=None, choices=None, cfg_path='default_config.yaml'): + """ + Parse command line arguments to the configuration according to the default yaml + + Args: + parser: Parent parser + cfg: Base configuration + helper: Helper description + cfg_path: Path to the default yaml config + """ + parser = argparse.ArgumentParser(description='[REPLACE THIS at config.py]', + parents=[parser]) + helper = {} if helper is None else helper + choices = {} if choices is None else choices + for item in cfg: + if not isinstance(cfg[item], list) and not isinstance(cfg[item], dict): + help_description = helper[item] if item in helper else 'Please reference to {}'.format(cfg_path) + choice = choices[item] if item in choices else None + if isinstance(cfg[item], bool): + parser.add_argument('--' + item, type=ast.literal_eval, default=cfg[item], choices=choice, + help=help_description) + else: + parser.add_argument('--' + item, type=type(cfg[item]), default=cfg[item], choices=choice, + help=help_description) + args = parser.parse_args() + return args + + +def parse_yaml(yaml_path): + """ + Parse the yaml config file + + Args: + yaml_path: Path to the yaml config + """ + with open(yaml_path, 'r') as fin: + try: + cfgs = yaml.load_all(fin.read(), Loader=yaml.FullLoader) + cfgs = [x for x in cfgs] + if len(cfgs) == 1: + cfg_helper = {} + cfg = cfgs[0] + cfg_choices = {} + elif len(cfgs) == 2: + cfg, cfg_helper = cfgs + cfg_choices = {} + elif len(cfgs) == 3: + cfg, cfg_helper, cfg_choices = cfgs + else: + raise ValueError('At most 3 docs (config description for help, choices) are supported in config yaml') + print(cfg_helper) + except: + raise ValueError('Failed to parse yaml') + return cfg, cfg_helper, cfg_choices + + +def merge(args, cfg): + """ + Merge the base config from yaml file and command line arguments + + Args: + args: command line arguments + cfg: Base configuration + """ + args_var = vars(args) + for item in args_var: + cfg[item] = args_var[item] + return cfg + + +def get_config(): + """ + Get Config according to the yaml file and cli arguments + """ + parser = argparse.ArgumentParser(description='default name', add_help=False) + current_dir = os.path.dirname(os.path.abspath(__file__)) + parser.add_argument('--config_path', type=str, default=os.path.join(current_dir, _config_path), + help='Config file path') + path_args, _ = parser.parse_known_args() + default, helper, choices = parse_yaml(path_args.config_path) + pprint(default) + args = parse_cli_to_yaml(parser=parser, cfg=default, helper=helper, choices=choices, cfg_path=path_args.config_path) + final_config = merge(args, default) + return Config(final_config) + +config = get_config() diff --git a/model_zoo/official/cv/ctpn/src/model_utils/device_adapter.py b/model_zoo/official/cv/ctpn/src/model_utils/device_adapter.py new file mode 100644 index 00000000000..ad8415af0f6 --- /dev/null +++ b/model_zoo/official/cv/ctpn/src/model_utils/device_adapter.py @@ -0,0 +1,26 @@ +# 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 WARRANT IES OR CONITTONS OF ANY KIND， either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ==================================================================================== + +"""Device adapter for ModelArts""" + +from .config import config +if config.enable_modelarts: + from .moxing_adapter import get_device_id, get_device_num, get_rank_id, get_job_id +else: + from .local_adapter import get_device_id, get_device_num, get_rank_id, get_job_id + +__all__ = [ + 'get_device_id', 'get_device_num', 'get_job_id', 'get_rank_id' +] diff --git a/model_zoo/official/cv/ctpn/src/model_utils/local_adapter.py b/model_zoo/official/cv/ctpn/src/model_utils/local_adapter.py new file mode 100644 index 00000000000..4ff88c4fba5 --- /dev/null +++ b/model_zoo/official/cv/ctpn/src/model_utils/local_adapter.py @@ -0,0 +1,36 @@ +# 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 WARRANT IES OR CONITTONS OF ANY KIND， either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ==================================================================================== + +"""Local adapter""" + +import os + +def get_device_id(): + device_id = os.getenv('DEVICE_ID', '0') + return int(device_id) + + +def get_device_num(): + device_num = os.getenv('RANK_SIZE', '1') + return int(device_num) + + +def get_rank_id(): + global_rank_id = os.getenv('RANK_ID', '0') + return int(global_rank_id) + + +def get_job_id(): + return 'Local Job' diff --git a/model_zoo/official/cv/ctpn/src/model_utils/moxing_adapter.py b/model_zoo/official/cv/ctpn/src/model_utils/moxing_adapter.py new file mode 100644 index 00000000000..c2d2282402b --- /dev/null +++ b/model_zoo/official/cv/ctpn/src/model_utils/moxing_adapter.py @@ -0,0 +1,124 @@ +# 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 WARRANT IES OR CONITTONS OF ANY KIND， either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ==================================================================================== + +"""Moxing adapter for ModelArts""" + +import os +import functools +from mindspore import context +from .config import config + + +_global_syn_count = 0 + + +def get_device_id(): + device_id = os.getenv('DEVICE_ID', '0') + return int(device_id) + + +def get_device_num(): + device_num = os.getenv('RANK_SIZE', '1') + return int(device_num) + + +def get_rank_id(): + global_rank_id = os.getenv('RANK_ID', '0') + return int(global_rank_id) + + +def get_job_id(): + job_id = os.getenv('JOB_ID') + job_id = job_id if job_id != "" else "default" + return job_id + + +def sync_data(from_path, to_path): + """ + Download data from remote obs to local directory if the first url is remote url and the second one is local + Uploca data from local directory to remote obs in contrast + """ + import moxing as mox + import time + global _global_syn_count + sync_lock = '/tmp/copy_sync.lock' + str(_global_syn_count) + _global_syn_count += 1 + + # Each server contains 8 devices as most + if get_device_id() % min(get_device_num(), 8) == 0 and not os.path.exists(sync_lock): + print('from path: ', from_path) + print('to path: ', to_path) + mox.file.copy_parallel(from_path, to_path) + print('===finished data synchronization===') + try: + os.mknod(sync_lock) + except IOError: + pass + print('===save flag===') + + while True: + if os.path.exists(sync_lock): + break + time.sleep(1) + print('Finish sync data from {} to {}'.format(from_path, to_path)) + + +def moxing_wrapper(pre_process=None, post_process=None): + """ + Moxing wrapper to download dataset and upload outputs + """ + def wrapper(run_func): + @functools.wraps(run_func) + def wrapped_func(*args, **kwargs): + # Download data from data_url + if config.enable_modelarts: + if config.data_url: + sync_data(config.data_url, config.data_path) + print('Dataset downloaded: ', os.listdir(config.data_path)) + if config.checkpoint_url: + if not os.path.exists(config.load_path): + # os.makedirs(config.load_path) + print('=' * 20 + 'makedirs') + if os.path.isdir(config.load_path): + print('=' * 20 + 'makedirs success') + else: + print('=' * 20 + 'makedirs fail') + sync_data(config.checkpoint_url, config.load_path) + print('Preload downloaded: ', os.listdir(config.load_path)) + if config.train_url: + sync_data(config.train_url, config.output_path) + print('Workspace downloaded: ', os.listdir(config.output_path)) + + context.set_context(save_graphs_path=os.path.join(config.output_path, str(get_rank_id()))) + config.device_num = get_device_num() + config.device_id = get_device_id() + if not os.path.exists(config.output_path): + os.makedirs(config.output_path) + + if pre_process: + pre_process() + + run_func(*args, **kwargs) + + # Upload data to train_url + if config.enable_modelarts: + if post_process: + post_process() + + if config.train_url: + print('Start to copy output directory') + sync_data(config.output_path, config.train_url) + return wrapped_func + return wrapper diff --git a/model_zoo/official/cv/ctpn/src/text_connector/detector.py b/model_zoo/official/cv/ctpn/src/text_connector/detector.py index 6518d920faf..7e5d724d440 100644 --- a/model_zoo/official/cv/ctpn/src/text_connector/detector.py +++ b/model_zoo/official/cv/ctpn/src/text_connector/detector.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================ import numpy as np -from src.config import config +from src.model_utils.config import config from src.text_connector.utils import nms from src.text_connector.connect_text_lines import connect_text_lines diff --git a/model_zoo/official/cv/ctpn/src/text_connector/get_successions.py b/model_zoo/official/cv/ctpn/src/text_connector/get_successions.py index 007a1b63257..c79e3e4d669 100644 --- a/model_zoo/official/cv/ctpn/src/text_connector/get_successions.py +++ b/model_zoo/official/cv/ctpn/src/text_connector/get_successions.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================ import numpy as np -from src.config import config +from src.model_utils.config import config from src.text_connector.utils import overlaps_v, size_similarity def get_successions(text_proposals, scores, im_size): diff --git a/model_zoo/official/cv/ctpn/train.py b/model_zoo/official/cv/ctpn/train.py index f42d147752e..37612171d8a 100644 --- a/model_zoo/official/cv/ctpn/train.py +++ b/model_zoo/official/cv/ctpn/train.py @@ -15,9 +15,8 @@ """train CTPN and get checkpoint files.""" import os -import time -import argparse import ast +import operator import mindspore.common.dtype as mstype from mindspore import context, Tensor from mindspore.communication.management import init @@ -28,38 +27,47 @@ from mindspore.train.serialization import load_checkpoint, load_param_into_net from mindspore.nn import Momentum from mindspore.common import set_seed from src.ctpn import CTPN -from src.config import config, pretrain_config, finetune_config from src.dataset import create_ctpn_dataset from src.lr_schedule import dynamic_lr from src.network_define import LossCallBack, LossNet, WithLossCell, TrainOneStepCell from src.eval_utils import eval_for_ctpn, get_eval_result from src.eval_callback import EvalCallBack +from src.model_utils.config import config +from src.model_utils.moxing_adapter import moxing_wrapper +from src.model_utils.device_adapter import get_device_num, get_device_id, get_rank_id + set_seed(1) -parser = argparse.ArgumentParser(description="CTPN training") -parser.add_argument("--run_distribute", type=ast.literal_eval, default=False, help="Run distribute, default: false.") -parser.add_argument("--pre_trained", type=str, default="", help="Pretrained file path.") -parser.add_argument("--device_id", type=int, default=0, help="Device id, default: 0.") -parser.add_argument("--device_num", type=int, default=1, help="Use device nums, default: 1.") -parser.add_argument("--rank_id", type=int, default=0, help="Rank id, default: 0.") -parser.add_argument("--task_type", type=str, default="Pretraining",\ - choices=['Pretraining', 'Finetune'], help="task type, default:Pretraining") -parser.add_argument("--run_eval", type=ast.literal_eval, default=False, \ - help="Run evaluation when training, default is False.") -parser.add_argument("--save_best_ckpt", type=ast.literal_eval, default=True, \ - help="Save best checkpoint when run_eval is True, default is True.") -parser.add_argument("--eval_image_path", type=str, default="", \ - help="eval image path, when run_eval is True, eval_image_path should be set.") -parser.add_argument("--eval_dataset_path", type=str, default="", \ - help="eval dataset path, when run_eval is True, eval_dataset_path should be set.") -parser.add_argument("--eval_start_epoch", type=int, default=10, \ - help="Evaluation start epoch when run_eval is True, default is 10.") -parser.add_argument("--eval_interval", type=int, default=10, \ - help="Evaluation interval when run_eval is True, default is 10.") -args_opt = parser.parse_args() -context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id, save_graphs=True) +context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=get_device_id(), save_graphs=True) + + +binOps = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Mod: operator.mod +} + + +def arithmeticeval(s): + node = ast.parse(s, mode='eval') + + def _eval(node): + if isinstance(node, ast.BinOp): + return binOps[type(node.op)](_eval(node.left), _eval(node.right)) + + if isinstance(node, ast.Num): + return node.n + + if isinstance(node, ast.Expression): + return _eval(node.body) + + raise Exception('unsupported type{}'.format(node)) + return _eval(node.body) + def apply_eval(eval_param): network = eval_param["eval_network"] @@ -69,43 +77,57 @@ def apply_eval(eval_param): hmean = get_eval_result() return hmean -if __name__ == '__main__': - if args_opt.run_distribute: - rank = args_opt.rank_id - device_num = args_opt.device_num + +def modelarts_pre_process(): + pass + + +@moxing_wrapper(pre_process=modelarts_pre_process) +def train(): + config.feature_shapes = [config.img_height // 16, config.img_width // 16] + config.num_bboxes = (config.img_height // 16) * (config.img_width // 16) * config.num_anchors + config.num_step = config.img_width // 16 + config.rnn_batch_size = config.img_height // 16 + config.weight_decay = arithmeticeval(config.weight_decay) + + if config.run_distribute: + rank = get_rank_id() + device_num = get_device_num() context.set_auto_parallel_context(device_num=device_num, parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True) init() else: rank = 0 device_num = 1 - if args_opt.task_type == "Pretraining": + if config.task_type == "Pretraining": print("Start to do pretraining") mindrecord_file = config.pretraining_dataset_file - training_cfg = pretrain_config + config.base_lr = config.pre_base_lr + config.warmup_step = config.pre_warmup_step + config.warmup_ratio = arithmeticeval(config.pre_warmup_ratio) + config.total_epoch = config.pre_total_epoch else: print("Start to do finetune") mindrecord_file = config.finetune_dataset_file - training_cfg = finetune_config - - print("CHECKING MINDRECORD FILES ...") - while not os.path.exists(mindrecord_file + ".db"): - time.sleep(5) + config.base_lr = config.fine_base_lr + config.warmup_step = config.fine_warmup_step + config.warmup_ratio = arithmeticeval(config.fine_warmup_ratio) + config.total_epoch = config.fine_total_epoch print("CHECKING MINDRECORD FILES DONE!") - loss_scale = float(config.loss_scale) + # loss_scale = float(config.loss_scale) # When create MindDataset, using the fitst mindrecord file, such as ctpn_pretrain.mindrecord0. - dataset = create_ctpn_dataset(mindrecord_file, repeat_num=1,\ + dataset = create_ctpn_dataset(mindrecord_file, repeat_num=1, \ batch_size=config.batch_size, device_num=device_num, rank_id=rank) dataset_size = dataset.get_dataset_size() net = CTPN(config=config, batch_size=config.batch_size) net = net.set_train() - load_path = args_opt.pre_trained - if args_opt.task_type == "Pretraining": - print("load backbone vgg16 ckpt {}".format(args_opt.pre_trained)) + load_path = config.pre_trained + if config.task_type == "Pretraining": + print("load backbone vgg16 ckpt {}".format(config.pre_trained)) param_dict = load_checkpoint(load_path) for item in list(param_dict.keys()): if not item.startswith('vgg16_feature_extractor'): @@ -113,15 +135,15 @@ if __name__ == '__main__': load_param_into_net(net, param_dict) else: if load_path != "": - print("load pretrain ckpt {}".format(args_opt.pre_trained)) + print("load pretrain ckpt {}".format(config.pre_trained)) param_dict = load_checkpoint(load_path) load_param_into_net(net, param_dict) loss = LossNet() - lr = Tensor(dynamic_lr(training_cfg, dataset_size), mstype.float32) + lr = Tensor(dynamic_lr(config, dataset_size), mstype.float32) opt = Momentum(params=net.trainable_params(), learning_rate=lr, momentum=config.momentum,\ weight_decay=config.weight_decay, loss_scale=config.loss_scale) net_with_loss = WithLossCell(net, loss) - if args_opt.run_distribute: + if config.run_distribute: net_with_grads = TrainOneStepCell(net_with_loss, opt, sens=config.loss_scale, reduce_flag=True, \ mean=True, degree=device_num) else: @@ -136,20 +158,24 @@ if __name__ == '__main__': keep_checkpoint_max=config.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix='ctpn', directory=save_checkpoint_path, config=ckptconfig) cb += [ckpoint_cb] - if args_opt.run_eval: - if args_opt.eval_dataset_path is None or (not os.path.isfile(args_opt.eval_dataset_path)): - raise ValueError("{} is not a existing path.".format(args_opt.eval_dataset_path)) - if args_opt.eval_image_path is None or (not os.path.isdir(args_opt.eval_image_path)): - raise ValueError("{} is not a existing path.".format(args_opt.eval_image_path)) - eval_dataset = create_ctpn_dataset(args_opt.eval_dataset_path, \ + if config.run_eval: + if config.eval_dataset_path is None or (not os.path.isfile(config.eval_dataset_path)): + raise ValueError("{} is not a existing path.".format(config.eval_dataset_path)) + if config.eval_image_path is None or (not os.path.isdir(config.eval_image_path)): + raise ValueError("{} is not a existing path.".format(config.eval_image_path)) + eval_dataset = create_ctpn_dataset(config.eval_dataset_path, \ batch_size=config.batch_size, repeat_num=1, is_training=False) eval_net = net eval_param_dict = {"eval_network": eval_net, "eval_dataset": eval_dataset, \ - "eval_image_path": args_opt.eval_image_path} - eval_cb = EvalCallBack(apply_eval, eval_param_dict, interval=args_opt.eval_interval, - eval_start_epoch=args_opt.eval_start_epoch, save_best_ckpt=True, + "eval_image_path": config.eval_image_path} + eval_cb = EvalCallBack(apply_eval, eval_param_dict, interval=config.eval_interval, + eval_start_epoch=config.eval_start_epoch, save_best_ckpt=True, ckpt_directory=save_checkpoint_path, besk_ckpt_name="best_acc.ckpt", metrics_name="hmean") cb += [eval_cb] model = Model(net_with_grads) - model.train(training_cfg.total_epoch, dataset, callbacks=cb, dataset_sink_mode=True) + model.train(config.total_epoch, dataset, callbacks=cb, dataset_sink_mode=True) + + +if __name__ == '__main__': + train()