diff --git a/model_zoo/research/cv/centernet/README.md b/model_zoo/research/cv/centernet/README.md index 404afbc571..163cf71031 100644 --- a/model_zoo/research/cv/centernet/README.md +++ b/model_zoo/research/cv/centernet/README.md @@ -117,18 +117,21 @@ Dataset used: [COCO2017]() After installing MindSpore via the official website, you can start training and evaluation as follows: -Note: 1.the first run will generate the mindrecord file, which will take a long time. - 2.VALIDATION_JSON_FILE is ground truth label file. CHECKPOINT_PATH is a checkpoint file after training. +Note: 1.the first run of training will generate the mindrecord file, which will take a long time. + 2.MINDRECORD_DATASET_PATH is the mindrecord dataset directory. ```shell +# create dataset in mindrecord format +bash scripts/convert_dataset_to_mindrecord.sh + # standalone training -bash run_standalone_train_ascend.sh [DEVICE_ID] [EPOCH_SIZE] +bash scripts/run_standalone_train_ascend.sh [DEVICE_ID] [EPOCH_SIZE] # distributed training -bash run_distributed_train_ascend.sh [COCO_DATASET_PATH] [MINDRECORD_DATASET_PATH] [RANK_TABLE_FILE] +bash scripts/run_distributed_train_ascend.sh [MINDRECORD_DATASET_PATH] [RANK_TABLE_FILE] # eval -bash run_standalone_eval_ascend.sh [DEVICE_ID] +bash scripts/run_standalone_eval_ascend.sh [DEVICE_ID] ``` # [Script Description](#contents) @@ -149,6 +152,7 @@ bash run_standalone_eval_ascend.sh [DEVICE_ID] │ │ ├──hyper_parameter_config.ini // hyper parameter for distributed pretraining │ │ ├──get_distribute_pretrain_cmd.py // script for distributed pretraining │ │ ├──README.md + │ ├──convert_dataset_to_mindrecord.sh // shell script for converting coco type dataset to mindrecord │ ├──run_standalone_train_ascend.sh // shell script for standalone pretrain on ascend │ ├──run_distributed_train_ascend.sh // shell script for distributed pretrain on ascend │ ├──run_standalone_eval_ascend.sh // shell script for standalone evaluation on ascend @@ -168,6 +172,19 @@ bash run_standalone_eval_ascend.sh [DEVICE_ID] ## [Script Parameters](#contents) +### Create MindRecord type dataset + +```text +usage: dataset.py [--coco_data_dir COCO_DATA_DIR] + [--mindrecord_dir MINDRECORD_DIR] + [--mindrecord_prefix MINDRECORD_PREFIX] + +options: + --coco_data_dir path to coco dataset directory: PATH, default is "" + --mindrecord_dir path to mindrecord dataset directory: PATH, default is "" + --mindrecord_prefix prefix of MindRecord dataset filename: STR, default is "coco_hp.train.mind" +``` + ### Training ```text @@ -180,7 +197,8 @@ usage: train.py [--device_target DEVICE_TARGET] [--distribute DISTRIBUTE] [--save_checkpoint_path SAVE_CHECKPOINT_PATH] [--load_checkpoint_path LOAD_CHECKPOINT_PATH] [--save_checkpoint_steps N] [--save_checkpoint_num N] - [--data_dir DATA_DIR] [--mindrecord_dir MINDRECORD_DIR] + [--mindrecord_dir MINDRECORD_DIR] + [--mindrecord_prefix MINDRECORD_PREFIX] [--visual_image VISUAL_IMAGE] [--save_result_dir SAVE_RESULT_DIR] options: @@ -201,8 +219,8 @@ options: --load_checkpoint_path path to load checkpoint files: PATH, default is "" --save_checkpoint_steps steps for saving checkpoint files: N, default is 1000 --save_checkpoint_num number for saving checkpoint files: N, default is 1 - --data_dir path to original dataset directory: PATH, default is "" --mindrecord_dir path to mindrecord dataset directory: PATH, default is "" + --mindrecord_prefix prefix of MindRecord dataset filename: STR, default is "coco_hp.train.mind" --visual_image whether visualize the image and annotation info: "true" | "false", default is "false" --save_result_dir path to save the visualization results: PATH, default is "" ``` @@ -214,7 +232,7 @@ usage: eval.py [--device_target DEVICE_TARGET] [--device_id N] [--load_checkpoint_path LOAD_CHECKPOINT_PATH] [--data_dir DATA_DIR] [--run_mode RUN_MODE] [--visual_image VISUAL_IMAGE] - [enable_eval ENABLE_EVAL] [--save_result_dir SAVE_RESULT_DIR] + [--enable_eval ENABLE_EVAL] [--save_result_dir SAVE_RESULT_DIR] options: --device_target device where the code will be implemented: "Ascend" | "CPU", default is "Ascend" --device_id device id to run task, default is 0 @@ -324,6 +342,14 @@ Parameters for optimizer and learning rate: ## [Training Process](#contents) +Before your first training, convert coco type dataset to mindrecord files is needed to improve performance on host. + +```bash +bash scripts/convert_dataset_to_mindrecord.sh +``` + +The command above will run in the background, after converting mindrecord files will be located in path specified by yourself. + ### Training #### Running on Ascend diff --git a/model_zoo/research/cv/centernet/eval.py b/model_zoo/research/cv/centernet/eval.py index eaff6152ba..369ba3a524 100644 --- a/model_zoo/research/cv/centernet/eval.py +++ b/model_zoo/research/cv/centernet/eval.py @@ -55,9 +55,9 @@ def predict(): context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id) logger.info("Begin creating {} dataset".format(args_opt.run_mode)) - coco = COCOHP(args_opt.data_dir, dataset_config, net_config, run_mode=args_opt.run_mode) - coco.init(enable_visual_image=(args_opt.visual_image == "true"), save_path=args_opt.save_result_dir, - keep_res=eval_config.keep_res, flip_test=eval_config.flip_test) + coco = COCOHP(dataset_config, run_mode=args_opt.run_mode, net_opt=net_config, + enable_visual_image=(args_opt.visual_image == "true"), save_path=args_opt.save_result_dir,) + coco.init(args_opt.data_dir, keep_res=eval_config.keep_res, flip_test=eval_config.flip_test) dataset = coco.create_eval_dataset() net_for_eval = CenterNetMultiPoseEval(net_config, eval_config.flip_test, eval_config.K) diff --git a/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/get_distribute_train_cmd.py b/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/get_distribute_train_cmd.py index 24972645f5..df66c98728 100644 --- a/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/get_distribute_train_cmd.py +++ b/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/get_distribute_train_cmd.py @@ -38,12 +38,9 @@ def parse_args(): help="Run script path, it is better to use absolute path") parser.add_argument("--hyper_parameter_config_dir", type=str, default="", help="Hyper Parameter config path, it is better to use absolute path") - parser.add_argument("--mindrecord_dir", type=str, default="./Mindrecord_train", - help="Mindrecord directory. If the mindrecord_dir is empty, it wil generate mindrecord file by " - "data_dir and anno_path. Note if mindrecord_dir isn't empty, it will use mindrecord_dir " - "rather than data_dir and anno_path. Default is ./Mindrecord_train") - parser.add_argument("--data_dir", type=str, default="", - help="Data path, it is better to use absolute path") + parser.add_argument("--mindrecord_dir", type=str, default="", help="Mindrecord dataset directory") + parser.add_argument("--mindrecord_prefix", type=str, default="coco_hp.train.mind", + help="Prefix of MindRecord dataset filename.") parser.add_argument("--hccl_config_dir", type=str, default="", help="Hccl config path, it is better to use absolute path") parser.add_argument("--cmd_file", type=str, default="distributed_cmd.sh", @@ -74,8 +71,8 @@ def distribute_train(): args = parse_args() run_script = args.run_script_dir - data_dir = args.data_dir mindrecord_dir = args.mindrecord_dir + mindrecord_prefix = args.mindrecord_prefix cf = configparser.ConfigParser() cf.read(args.hyper_parameter_config_dir) cfg = dict(cf.items("config")) @@ -142,7 +139,6 @@ def distribute_train(): print("core_nums:", cmdopt) print("epoch_size:", str(cfg['epoch_size'])) - print("data_dir:", data_dir) print("mindrecord_dir:", mindrecord_dir) print("log_file_dir: " + cur_dir + "/LOG" + str(device_id) + "/training_log.txt") @@ -150,12 +146,12 @@ def distribute_train(): run_cmd = 'taskset -c ' + cmdopt + ' nohup python ' + run_script + " " opt = " ".join(["--" + key + "=" + str(cfg[key]) for key in cfg.keys()]) - if ('device_id' in opt) or ('device_num' in opt) or ('data_dir' in opt): + if ('device_id' in opt) or ('device_num' in opt) or ('mindrecord_dir' in opt): raise ValueError("hyper_parameter_config.ini can not setting 'device_id'," - " 'device_num' or 'data_dir'! ") + " 'device_num' or 'mindrecord_dir'! ") run_cmd += opt - run_cmd += " --data_dir=" + data_dir run_cmd += " --mindrecord_dir=" + mindrecord_dir + run_cmd += " --mindrecord_prefix=" + mindrecord_prefix run_cmd += ' --device_id=' + str(device_id) + ' --device_num=' \ + str(rank_size) + ' >./training_log.txt 2>&1 &' diff --git a/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/hyper_parameter_config.ini b/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/hyper_parameter_config.ini index 1d845d3cae..7f0fa64131 100644 --- a/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/hyper_parameter_config.ini +++ b/model_zoo/research/cv/centernet/scripts/ascend_distributed_launcher/hyper_parameter_config.ini @@ -9,6 +9,7 @@ load_checkpoint_path="" save_checkpoint_path=./ save_checkpoint_steps=3000 save_checkpoint_num=1 +mindrecord_prefix="coco_hp.train.mind" need_profiler=false profiler_path=./profiler visual_image=false \ No newline at end of file diff --git a/model_zoo/research/cv/centernet/scripts/convert_dataset_to_mindrecord.sh b/model_zoo/research/cv/centernet/scripts/convert_dataset_to_mindrecord.sh new file mode 100644 index 0000000000..1f34b827f2 --- /dev/null +++ b/model_zoo/research/cv/centernet/scripts/convert_dataset_to_mindrecord.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# 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. +# ============================================================================ + +echo "==============================================================================================================" +echo "Please run the scipt as: " +echo "bash convert_dataset_to_mindrecord.sh" +echo "==============================================================================================================" + +export GLOG_v=1 +PROJECT_DIR=$(cd "$(dirname "$0")" || exit; pwd) + +python ${PROJECT_DIR}/../src/dataset.py \ + --coco_data_dir="" \ + --mindrecord_dir="" \ + --mindrecord_prefix="coco_hp.train.mind" > create_dataset.log 2>&1 & \ No newline at end of file diff --git a/model_zoo/research/cv/centernet/scripts/run_distributed_train_ascend.sh b/model_zoo/research/cv/centernet/scripts/run_distributed_train_ascend.sh index 7b5d8fdbab..7e4242f1f1 100644 --- a/model_zoo/research/cv/centernet/scripts/run_distributed_train_ascend.sh +++ b/model_zoo/research/cv/centernet/scripts/run_distributed_train_ascend.sh @@ -27,9 +27,8 @@ CUR_DIR=`pwd` python ${CUR_DIR}/scripts/ascend_distributed_launcher/get_distribute_train_cmd.py \ --run_script_dir=${CUR_DIR}/train.py \ --hyper_parameter_config_dir=${CUR_DIR}/scripts/ascend_distributed_launcher/hyper_parameter_config.ini \ - --data_dir=$1 \ - --mindrecord_dir=$2 \ - --hccl_config_dir=$3 \ + --mindrecord_dir=$1 \ + --hccl_config_dir=$2 \ --hccl_time_out=1200 \ --cmd_file=distributed_cmd.sh diff --git a/model_zoo/research/cv/centernet/scripts/run_standalone_eval_ascend.sh b/model_zoo/research/cv/centernet/scripts/run_standalone_eval_ascend.sh index a6ddd09c0c..ad2c59399d 100644 --- a/model_zoo/research/cv/centernet/scripts/run_standalone_eval_ascend.sh +++ b/model_zoo/research/cv/centernet/scripts/run_standalone_eval_ascend.sh @@ -48,4 +48,4 @@ python ${PROJECT_DIR}/../eval.py \ --visual_image=true \ --enable_eval=true \ --save_result_dir="" \ - --run_mode=val > log.txt 2>&1 & \ No newline at end of file + --run_mode=val > eval_log.txt 2>&1 & \ No newline at end of file diff --git a/model_zoo/research/cv/centernet/scripts/run_standalone_train_ascend.sh b/model_zoo/research/cv/centernet/scripts/run_standalone_train_ascend.sh index 32d564610c..ff2f7e7aef 100644 --- a/model_zoo/research/cv/centernet/scripts/run_standalone_train_ascend.sh +++ b/model_zoo/research/cv/centernet/scripts/run_standalone_train_ascend.sh @@ -42,7 +42,7 @@ python ${PROJECT_DIR}/../train.py \ --load_checkpoint_path="" \ --save_checkpoint_steps=10000 \ --save_checkpoint_num=1 \ - --data_dir="" \ --mindrecord_dir="" \ + --mindrecord_prefix="coco_hp.train.mind" \ --visual_image=false \ - --save_result_dir=""> log.txt 2>&1 & \ No newline at end of file + --save_result_dir="" > training_log.txt 2>&1 & \ No newline at end of file diff --git a/model_zoo/research/cv/centernet/src/dataset.py b/model_zoo/research/cv/centernet/src/dataset.py index a4cfdcec26..463fd9b5c9 100644 --- a/model_zoo/research/cv/centernet/src/dataset.py +++ b/model_zoo/research/cv/centernet/src/dataset.py @@ -19,6 +19,7 @@ Data operations, will be used in train.py import os import copy import math +import argparse import cv2 import numpy as np import pycocotools.coco as coco @@ -26,10 +27,9 @@ import pycocotools.coco as coco import mindspore.dataset.engine.datasets as de from mindspore import log as logger from mindspore.mindrecord import FileWriter -from .image import color_aug -from .image import get_affine_transform, affine_transform -from .image import gaussian_radius, draw_umich_gaussian, draw_msra_gaussian, draw_dense_reg -from .visual import visual_image +from src.image import color_aug, get_affine_transform, affine_transform +from src.image import gaussian_radius, draw_umich_gaussian, draw_msra_gaussian, draw_dense_reg +from src.visual import visual_image _current_dir = os.path.dirname(os.path.realpath(__file__)) @@ -47,40 +47,39 @@ class COCOHP(de.Dataset): Returns: Prepocessed training or testing dataset for CenterNet network. """ - def __init__(self, data_dir, data_opt, net_opt, run_mode): + def __init__(self, data_opt, run_mode="train", net_opt=None, enable_visual_image=False, save_path=None): super(COCOHP, self).__init__() - if not os.path.isdir(data_dir): - raise RuntimeError("Invalid dataset path") + self._data_rng = np.random.RandomState(123) + self.data_opt = data_opt + self.data_opt.mean = self.data_opt.mean.reshape(1, 1, 3) + self.data_opt.std = self.data_opt.std.reshape(1, 1, 3) assert run_mode in ["train", "test", "val"], "only train/test/val mode are supported" self.run_mode = run_mode + if net_opt is not None: + self.net_opt = net_opt + self.enable_visual_image = enable_visual_image + if self.enable_visual_image: + self.save_path = os.path.join(save_path, self.run_mode, "input_image") + if not os.path.exists(self.save_path): + os.makedirs(self.save_path) + + + def init(self, data_dir, keep_res=False, flip_test=False): + """initailize additional info""" + logger.info('Initializing coco 2017 {} data.'.format(self.run_mode)) + if not os.path.isdir(data_dir): + raise RuntimeError("Invalid dataset path") if self.run_mode != "test": self.annot_path = os.path.join(data_dir, 'annotations', 'person_keypoints_{}2017.json').format(self.run_mode) else: self.annot_path = os.path.join(data_dir, 'annotations', 'image_info_test-dev2017.json') self.image_path = os.path.join(data_dir, '{}2017').format(self.run_mode) - - self._data_rng = np.random.RandomState(123) - self.data_opt = data_opt - self.data_opt.mean = self.data_opt.mean.reshape(1, 1, 3) - self.data_opt.std = self.data_opt.std.reshape(1, 1, 3) - self.net_opt = net_opt - self.coco = coco.COCO(self.annot_path) - - - def init(self, enable_visual_image=False, save_path=None, keep_res=False, flip_test=False): - """initailize additional info""" - logger.info('Initializing coco 2017 {} data.'.format(self.run_mode)) logger.info('Image path: {}'.format(self.image_path)) logger.info('Annotations: {}'.format(self.annot_path)) - self.enable_visual_image = enable_visual_image - if self.enable_visual_image: - self.save_path = os.path.join(save_path, self.run_mode, "input_image") - if not os.path.exists(self.save_path): - os.makedirs(self.save_path) - + self.coco = coco.COCO(self.annot_path) image_ids = self.coco.getImgIds() if self.run_mode != "test": self.images = [] @@ -102,8 +101,15 @@ class COCOHP(de.Dataset): def __len__(self): return self.num_samples - def transfer_coco_to_mindrecord(self, mindrecord_dir, file_name, shard_num=1): + def transfer_coco_to_mindrecord(self, mindrecord_dir, file_name="coco_hp.train.mind", shard_num=1): """Create MindRecord file by image_dir and anno_path.""" + if not os.path.isdir(mindrecord_dir): + os.makedirs(mindrecord_dir) + if os.path.isdir(self.image_path) and os.path.exists(self.annot_path): + logger.info("Create MindRecord based on COCO_HP dataset") + else: + raise ValueError('data_dir {} or anno_path {} does not exist'.format(self.image_path, self.annot_path)) + mindrecord_path = os.path.join(mindrecord_dir, file_name) writer = FileWriter(mindrecord_path, shard_num) centernet_json = { @@ -139,6 +145,7 @@ class COCOHP(de.Dataset): "category_id": np.array(category_id, np.int32)} writer.write_raw_data([row]) writer.commit() + logger.info("Create Mindrecord Done, at {}".format(mindrecord_dir)) def _coco_box_to_bbox(self, box): @@ -393,19 +400,11 @@ class COCOHP(de.Dataset): return ret - def create_train_dataset(self, mindrecord_dir, prefix, batch_size=1, + def create_train_dataset(self, mindrecord_dir, prefix="coco_hp.train.mind", batch_size=1, device_num=1, rank=0, num_parallel_workers=1, do_shuffle=True): """create train dataset based on mindrecord file""" if not os.path.isdir(mindrecord_dir): - os.makedirs(mindrecord_dir) - if os.path.isdir(self.image_path) and os.path.exists(self.annot_path): - logger.info("Create MindRecord based on COCO_HP dataset") - self.transfer_coco_to_mindrecord(mindrecord_dir, prefix, shard_num=8) - logger.info("Create Mindrecord Done, at {}".format(mindrecord_dir)) - else: - raise ValueError('data_dir {} or anno_path {} does not exist'.format(self.image_path, self.annot_path)) - else: - logger.info("MindRecord dataset already exists, dir: {}".format(mindrecord_dir)) + raise ValueError('MindRecord data_dir {} does not exist'.format(mindrecord_dir)) files = os.listdir(mindrecord_dir) data_files = [] @@ -415,7 +414,6 @@ class COCOHP(de.Dataset): if not data_files: raise ValueError('data_dir {} have no data files'.format(mindrecord_dir)) - columns = ["image", "num_objects", "keypoints", "bbox", "category_id"] ds = de.MindDataset(data_files, columns_list=columns, @@ -447,3 +445,17 @@ class COCOHP(de.Dataset): ds = de.GeneratorDataset(generator, column, num_parallel_workers=num_parallel_workers) ds = ds.batch(batch_size, drop_remainder=True, num_parallel_workers=8) return ds + + +if __name__ == '__main__': + # Convert coco2017 dataset to mindrecord to improve performance on host + from src.config import dataset_config + parser = argparse.ArgumentParser(description='CenterNet MindRecord dataset') + parser.add_argument("--coco_data_dir", type=str, default="", help="Coco dataset directory.") + parser.add_argument("--mindrecord_dir", type=str, default="", help="MindRecord dataset dir.") + parser.add_argument("--mindrecord_prefix", type=str, default="coco_hp.train.mind", + help="Prefix of MindRecord dataset filename.") + args_opt = parser.parse_args() + dsc = COCOHP(dataset_config, run_mode="train") + dsc.init(args_opt.coco_data_dir) + dsc.transfer_coco_to_mindrecord(args_opt.mindrecord_dir, args_opt.mindrecord_prefix, shard_num=8) diff --git a/model_zoo/research/cv/centernet/src/utils.py b/model_zoo/research/cv/centernet/src/utils.py index 27a5a291db..307003d2d0 100644 --- a/model_zoo/research/cv/centernet/src/utils.py +++ b/model_zoo/research/cv/centernet/src/utils.py @@ -20,7 +20,6 @@ import math import numpy as np import mindspore.nn as nn import mindspore.ops as ops -from mindspore import log as logger from mindspore.common.tensor import Tensor from mindspore.common.parameter import Parameter from mindspore.common.initializer import initializer @@ -494,11 +493,11 @@ class LossCallBack(Callback): if percent == 0: percent = 1 epoch_num -= 1 - logger.info("epoch: {}, current epoch percent: {}, step: {}, outputs are {}" - .format(int(epoch_num), "%.3f" % percent, cb_params.cur_step_num, str(cb_params.net_outputs))) + print("epoch: {}, current epoch percent: {}, step: {}, outputs are {}" + .format(int(epoch_num), "%.3f" % percent, cb_params.cur_step_num, str(cb_params.net_outputs))) else: - logger.info("epoch: {}, step: {}, outputs are {}".format(cb_params.cur_epoch_num, cb_params.cur_step_num, - str(cb_params.net_outputs))) + print("epoch: {}, step: {}, outputs are {}".format(cb_params.cur_epoch_num, cb_params.cur_step_num, + str(cb_params.net_outputs))) class CenterNetPolynomialDecayLR(LearningRateSchedule): diff --git a/model_zoo/research/cv/centernet/train.py b/model_zoo/research/cv/centernet/train.py index 30fabff6a1..46784be702 100644 --- a/model_zoo/research/cv/centernet/train.py +++ b/model_zoo/research/cv/centernet/train.py @@ -58,13 +58,9 @@ parser.add_argument("--save_checkpoint_path", type=str, default="", help="Save c parser.add_argument("--load_checkpoint_path", type=str, default="", help="Load checkpoint file path") parser.add_argument("--save_checkpoint_steps", type=int, default=1000, help="Save checkpoint steps, default is 1000.") parser.add_argument("--save_checkpoint_num", type=int, default=1, help="Save checkpoint numbers, default is 1.") -parser.add_argument("--mindrecord_dir", type=str, default="", - help="Mindrecord files directory. If is empty, mindrecord format files will be generated" - "based on the original dataset and annotation information. If mindrecord_dir isn't empty," - "mindrecord_dir will be used inplace of data_dir and anno_path.") -parser.add_argument("--data_dir", type=str, default="", help="Dataset directory, " - "the absolute image path is joined by the data_dir " - "and the relative path in anno_path") +parser.add_argument("--mindrecord_dir", type=str, default="", help="Mindrecord dataset files directory") +parser.add_argument("--mindrecord_prefix", type=str, default="coco_hp.train.mind", + help="Prefix of MindRecord dataset filename.") parser.add_argument("--visual_image", type=str, default="false", help="Visulize the ground truth and predicted image") parser.add_argument("--save_result_dir", type=str, default="", help="The path to save the predict results") @@ -148,16 +144,15 @@ def train(): else: rank = 0 device_num = 1 - num_workers = device_num * 8 + num_workers = 8 # Start create dataset! # mindrecord files will be generated at args_opt.mindrecord_dir such as centernet.mindrecord0, 1, ... file_num. logger.info("Begin creating dataset for CenterNet") - prefix = "coco_hp.train.mind" - coco = COCOHP(args_opt.data_dir, dataset_config, net_config, run_mode="train") - coco.init(enable_visual_image=(args_opt.visual_image == "true"), save_path=args_opt.save_result_dir) - dataset = coco.create_train_dataset(args_opt.mindrecord_dir, prefix, batch_size=train_config.batch_size, - device_num=device_num, rank=rank, num_parallel_workers=num_workers, - do_shuffle=args_opt.do_shuffle == 'true') + coco = COCOHP(dataset_config, run_mode="train", net_opt=net_config, + enable_visual_image=(args_opt.visual_image == "true"), save_path=args_opt.save_result_dir) + dataset = coco.create_train_dataset(args_opt.mindrecord_dir, args_opt.mindrecord_prefix, + batch_size=train_config.batch_size, device_num=device_num, rank=rank, + num_parallel_workers=num_workers, do_shuffle=args_opt.do_shuffle == 'true') dataset_size = dataset.get_dataset_size() logger.info("Create dataset done!")