Skip to content

GitLab

Commit aad3093a authored by WenmuZhou's avatar WenmuZhou
Browse files

dygraph first commit

parent 10f7e519
Hide whitespace changes
Inline Side-by-side
Showing with 455 additions and 718 deletions
tools/infer/predict_det.py
View file @ aad3093a
...@@ -13,71 +13,63 @@ ...@@ -13,71 +13,63 @@
# limitations under the License. # limitations under the License.
import os import os
import sys import sys
__dir__ = os.path.dirname(os.path.abspath(__file__)) __dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__) sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
import cv2 import cv2
import copy
import numpy as np import numpy as np
import math
import time import time
import sys import sys
import paddle.fluid as fluid import paddle
import tools.infer.utility as utility import tools.infer.utility as utility
from ppocr.utils.utility import initial_logger from ppocr.utils.logging import get_logger
logger = initial_logger()
from ppocr.utils.utility import get_image_file_list, check_and_read_gif from ppocr.utils.utility import get_image_file_list, check_and_read_gif
from ppocr.data.det.sast_process import SASTProcessTest from ppocr.data import create_operators, transform
from ppocr.data.det.east_process import EASTProcessTest from ppocr.postprocess import build_post_process
from ppocr.data.det.db_process import DBProcessTest
from ppocr.postprocess.db_postprocess import DBPostProcess
from ppocr.postprocess.east_postprocess import EASTPostPocess
from ppocr.postprocess.sast_postprocess import SASTPostProcess
class TextDetector(object): class TextDetector(object):
def __init__(self, args): def __init__(self, args):
max_side_len = args.det_max_side_len
self.det_algorithm = args.det_algorithm self.det_algorithm = args.det_algorithm
self.use_zero_copy_run = args.use_zero_copy_run self.use_zero_copy_run = args.use_zero_copy_run
preprocess_params = {'max_side_len': max_side_len}
postprocess_params = {} postprocess_params = {}
if self.det_algorithm == "DB": if self.det_algorithm == "DB":
self.preprocess_op = DBProcessTest(preprocess_params) pre_process_list = [{
'ResizeForTest': {
'limit_side_len': args.det_limit_side_len,
'limit_type': args.det_limit_type
}
}, {
'NormalizeImage': {
'std': [0.229, 0.224, 0.225],
'mean': [0.485, 0.456, 0.406],
'scale': '1./255.',
'order': 'hwc'
}
}, {
'ToCHWImage': None
}, {
'keepKeys': {
'keep_keys': ['image', 'shape']
}
}]
postprocess_params['name'] = 'DBPostProcess'
postprocess_params["thresh"] = args.det_db_thresh postprocess_params["thresh"] = args.det_db_thresh
postprocess_params["box_thresh"] = args.det_db_box_thresh postprocess_params["box_thresh"] = args.det_db_box_thresh
postprocess_params["max_candidates"] = 1000 postprocess_params["max_candidates"] = 1000
postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio
self.postprocess_op = DBPostProcess(postprocess_params)
elif self.det_algorithm == "EAST":
self.preprocess_op = EASTProcessTest(preprocess_params)
postprocess_params["score_thresh"] = args.det_east_score_thresh
postprocess_params["cover_thresh"] = args.det_east_cover_thresh
postprocess_params["nms_thresh"] = args.det_east_nms_thresh
self.postprocess_op = EASTPostPocess(postprocess_params)
elif self.det_algorithm == "SAST":
self.preprocess_op = SASTProcessTest(preprocess_params)
postprocess_params["score_thresh"] = args.det_sast_score_thresh
postprocess_params["nms_thresh"] = args.det_sast_nms_thresh
self.det_sast_polygon = args.det_sast_polygon
if self.det_sast_polygon:
postprocess_params["sample_pts_num"] = 6
postprocess_params["expand_scale"] = 1.2
postprocess_params["shrink_ratio_of_width"] = 0.2
else:
postprocess_params["sample_pts_num"] = 2
postprocess_params["expand_scale"] = 1.0
postprocess_params["shrink_ratio_of_width"] = 0.3
self.postprocess_op = SASTPostProcess(postprocess_params)
else: else:
logger.info("unknown det_algorithm:{}".format(self.det_algorithm)) logger.info("unknown det_algorithm:{}".format(self.det_algorithm))
sys.exit(0) sys.exit(0)
self.predictor, self.input_tensor, self.output_tensors =\ self.preprocess_op = create_operators(pre_process_list)
utility.create_predictor(args, mode="det") self.postprocess_op = build_post_process(postprocess_params)
self.predictor = paddle.jit.load(args.det_model_dir)
self.predictor.eval()
def order_points_clockwise(self, pts): def order_points_clockwise(self, pts):
""" """
...@@ -134,46 +126,31 @@ class TextDetector(object): ...@@ -134,46 +126,31 @@ class TextDetector(object):
def __call__(self, img): def __call__(self, img):
ori_im = img.copy() ori_im = img.copy()
im, ratio_list = self.preprocess_op(img) data = {'image': img}
if im is None: data = transform(data, self.preprocess_op)
img, shape_list = data
if img is None:
return None, 0 return None, 0
im = im.copy() img = np.expand_dims(img, axis=0)
shape_list = np.expand_dims(shape_list, axis=0)
starttime = time.time() starttime = time.time()
if self.use_zero_copy_run:
self.input_tensor.copy_from_cpu(im)
self.predictor.zero_copy_run()
else:
im = fluid.core.PaddleTensor(im)
self.predictor.run([im])
outputs = []
for output_tensor in self.output_tensors:
output = output_tensor.copy_to_cpu()
outputs.append(output)
outs_dict = {}
if self.det_algorithm == "EAST":
outs_dict['f_geo'] = outputs[0]
outs_dict['f_score'] = outputs[1]
elif self.det_algorithm == 'SAST':
outs_dict['f_border'] = outputs[0]
outs_dict['f_score'] = outputs[1]
outs_dict['f_tco'] = outputs[2]
outs_dict['f_tvo'] = outputs[3]
else:
outs_dict['maps'] = outputs[0]
dt_boxes_list = self.postprocess_op(outs_dict, [ratio_list]) preds = self.predictor(img)
dt_boxes = dt_boxes_list[0] post_result = self.postprocess_op(preds, shape_list)
if self.det_algorithm == "SAST" and self.det_sast_polygon:
dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape) dt_boxes = post_result[0]['points']
else: dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
elapse = time.time() - starttime elapse = time.time() - starttime
return dt_boxes, elapse return dt_boxes, elapse
if __name__ == "__main__": if __name__ == "__main__":
args = utility.parse_args() args = utility.parse_args()
place = paddle.CPUPlace()
paddle.disable_static(place)
image_file_list = get_image_file_list(args.image_dir) image_file_list = get_image_file_list(args.image_dir)
logger = get_logger()
text_detector = TextDetector(args) text_detector = TextDetector(args)
count = 0 count = 0
total_time = 0 total_time = 0
...@@ -187,6 +164,7 @@ if __name__ == "__main__": ...@@ -187,6 +164,7 @@ if __name__ == "__main__":
if img is None: if img is None:
logger.info("error in loading image:{}".format(image_file)) logger.info("error in loading image:{}".format(image_file))
continue continue
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
dt_boxes, elapse = text_detector(img) dt_boxes, elapse = text_detector(img)
if count > 0: if count > 0:
total_time += elapse total_time += elapse
......
tools/infer/utility.py
View file @ aad3093a
...@@ -13,12 +13,7 @@ ...@@ -13,12 +13,7 @@
# limitations under the License. # limitations under the License.
import argparse import argparse
import os, sys import os
from ppocr.utils.utility import initial_logger
logger = initial_logger()
from paddle.fluid.core import PaddleTensor
from paddle.fluid.core import AnalysisConfig
from paddle.fluid.core import create_paddle_predictor
import cv2 import cv2
import numpy as np import numpy as np
import json import json
...@@ -41,7 +36,8 @@ def parse_args(): ...@@ -41,7 +36,8 @@ def parse_args():
parser.add_argument("--image_dir", type=str) parser.add_argument("--image_dir", type=str)
parser.add_argument("--det_algorithm", type=str, default='DB') parser.add_argument("--det_algorithm", type=str, default='DB')
parser.add_argument("--det_model_dir", type=str) parser.add_argument("--det_model_dir", type=str)
parser.add_argument("--det_max_side_len", type=float, default=960) parser.add_argument("--det_limit_side_len", type=float, default=960)
parser.add_argument("--det_limit_type", type=str, default='max')
#DB parmas #DB parmas
parser.add_argument("--det_db_thresh", type=float, default=0.3) parser.add_argument("--det_db_thresh", type=float, default=0.3)
...@@ -75,54 +71,6 @@ def parse_args(): ...@@ -75,54 +71,6 @@ def parse_args():
return parser.parse_args() return parser.parse_args()
def create_predictor(args, mode):
if mode == "det":
model_dir = args.det_model_dir
else:
model_dir = args.rec_model_dir
if model_dir is None:
logger.info("not find {} model file path {}".format(mode, model_dir))
sys.exit(0)
model_file_path = model_dir + "/model"
params_file_path = model_dir + "/params"
if not os.path.exists(model_file_path):
logger.info("not find model file path {}".format(model_file_path))
sys.exit(0)
if not os.path.exists(params_file_path):
logger.info("not find params file path {}".format(params_file_path))
sys.exit(0)
config = AnalysisConfig(model_file_path, params_file_path)
if args.use_gpu:
config.enable_use_gpu(args.gpu_mem, 0)
else:
config.disable_gpu()
config.set_cpu_math_library_num_threads(6)
if args.enable_mkldnn:
config.enable_mkldnn()
#config.enable_memory_optim()
config.disable_glog_info()
if args.use_zero_copy_run:
config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
config.switch_use_feed_fetch_ops(False)
else:
config.switch_use_feed_fetch_ops(True)
predictor = create_paddle_predictor(config)
input_names = predictor.get_input_names()
input_tensor = predictor.get_input_tensor(input_names[0])
output_names = predictor.get_output_names()
output_tensors = []
for output_name in output_names:
output_tensor = predictor.get_output_tensor(output_name)
output_tensors.append(output_tensor)
return predictor, input_tensor, output_tensors
def draw_text_det_res(dt_boxes, img_path): def draw_text_det_res(dt_boxes, img_path):
src_im = cv2.imread(img_path) src_im = cv2.imread(img_path)
for box in dt_boxes: for box in dt_boxes:
...@@ -139,8 +87,8 @@ def resize_img(img, input_size=600): ...@@ -139,8 +87,8 @@ def resize_img(img, input_size=600):
im_shape = img.shape im_shape = img.shape
im_size_max = np.max(im_shape[0:2]) im_size_max = np.max(im_shape[0:2])
im_scale = float(input_size) / float(im_size_max) im_scale = float(input_size) / float(im_size_max)
im = cv2.resize(img, None, None, fx=im_scale, fy=im_scale) img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
return im return img
def draw_ocr(image, def draw_ocr(image,
......
tools/infer_det.py
View file @ aad3093a
...@@ -17,38 +17,25 @@ from __future__ import division ...@@ -17,38 +17,25 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np import numpy as np
from copy import deepcopy
import json
import os import os
import sys import sys
__dir__ = os.path.dirname(os.path.abspath(__file__)) __dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__) sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
def set_paddle_flags(**kwargs):
for key, value in kwargs.items():
if os.environ.get(key, None) is None:
os.environ[key] = str(value)
# NOTE(paddle-dev): All of these flags should be
# set before `import paddle`. Otherwise, it would
# not take any effect.
set_paddle_flags(
FLAGS_eager_delete_tensor_gb=0, # enable GC to save memory
)
from paddle import fluid
from ppocr.utils.utility import create_module, get_image_file_list
import program
from ppocr.utils.save_load import init_model
from ppocr.data.reader_main import reader_main
import cv2 import cv2
import json
import paddle
from ppocr.utils.utility import initial_logger from ppocr.utils.logging import get_logger
logger = initial_logger() from ppocr.data import create_operators, transform
from ppocr.modeling import build_model
from ppocr.postprocess import build_post_process
from ppocr.utils.save_load import init_model
from ppocr.utils.utility import print_dict, get_image_file_list
import tools.program as program
def draw_det_res(dt_boxes, config, img, img_name): def draw_det_res(dt_boxes, config, img, img_name):
...@@ -68,94 +55,68 @@ def draw_det_res(dt_boxes, config, img, img_name): ...@@ -68,94 +55,68 @@ def draw_det_res(dt_boxes, config, img, img_name):
def main(): def main():
config = program.load_config(FLAGS.config) global_config = config['Global']
program.merge_config(FLAGS.opt)
logger.info(config) # build model
model = build_model(config['Architecture'])
# check if set use_gpu=True in paddlepaddle cpu version
use_gpu = config['Global']['use_gpu'] init_model(config, model, logger)
program.check_gpu(use_gpu)
# build post process
place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() post_process_class = build_post_process(config['PostProcess'])
exe = fluid.Executor(place)
# create data ops
det_model = create_module(config['Architecture']['function'])(params=config) transforms = []
for op in config['EVAL']['dataset']['transforms']:
startup_prog = fluid.Program() op_name = list(op)[0]
eval_prog = fluid.Program() if 'Label' in op_name:
with fluid.program_guard(eval_prog, startup_prog): continue
with fluid.unique_name.guard(): elif op_name == 'keepKeys':
_, eval_outputs = det_model(mode="test") op[op_name]['keep_keys'] = ['image', 'shape']
fetch_name_list = list(eval_outputs.keys()) transforms.append(op)
eval_fetch_list = [eval_outputs[v].name for v in fetch_name_list]
ops = create_operators(transforms, global_config)
eval_prog = eval_prog.clone(for_test=True)
exe.run(startup_prog)
# load checkpoints
checkpoints = config['Global'].get('checkpoints')
if checkpoints:
path = checkpoints
fluid.load(eval_prog, path, exe)
logger.info("Finish initing model from {}".format(path))
else:
raise Exception("{} not exists!".format(checkpoints))
save_res_path = config['Global']['save_res_path'] save_res_path = config['Global']['save_res_path']
if not os.path.exists(os.path.dirname(save_res_path)): if not os.path.exists(os.path.dirname(save_res_path)):
os.makedirs(os.path.dirname(save_res_path)) os.makedirs(os.path.dirname(save_res_path))
with open(save_res_path, "wb") as fout:
test_reader = reader_main(config=config, mode='test') model.eval()
tackling_num = 0 with open(save_res_path, "wb") as fout:
for data in test_reader(): for file in get_image_file_list(config['Global']['infer_img']):
img_num = len(data) logger.info("infer_img: {}".format(file))
tackling_num = tackling_num + img_num with open(file, 'rb') as f:
logger.info("tackling_num:%d", tackling_num) img = f.read()
img_list = [] data = {'image': img}
ratio_list = [] batch = transform(data, ops)
img_name_list = []
for ino in range(img_num): images = np.expand_dims(batch[0], axis=0)
img_list.append(data[ino][0]) shape_list = np.expand_dims(batch[1], axis=0)
ratio_list.append(data[ino][1]) images = paddle.to_variable(images)
img_name_list.append(data[ino][2]) print(images.shape)
preds = model(images)
img_list = np.concatenate(img_list, axis=0) post_result = post_process_class(preds, shape_list)
outs = exe.run(eval_prog,\ boxes = post_result[0]['points']
feed={'image': img_list},\ # write resule
fetch_list=eval_fetch_list) dt_boxes_json = []
for box in boxes:
global_params = config['Global'] tmp_json = {"transcription": ""}
postprocess_params = deepcopy(config["PostProcess"]) tmp_json['points'] = box.tolist()
postprocess_params.update(global_params) dt_boxes_json.append(tmp_json)
postprocess = create_module(postprocess_params['function'])\ otstr = file + "\t" + json.dumps(dt_boxes_json) + "\n"
(params=postprocess_params) fout.write(otstr.encode())
if config['Global']['algorithm'] == 'EAST': src_img = cv2.imread(file)
dic = {'f_score': outs[0], 'f_geo': outs[1]} draw_det_res(boxes, config, src_img, file)
elif config['Global']['algorithm'] == 'DB':
dic = {'maps': outs[0]}
elif config['Global']['algorithm'] == 'SAST':
dic = {'f_score': outs[0], 'f_border': outs[1], 'f_tvo': outs[2], 'f_tco': outs[3]}
else:
raise Exception("only support algorithm: ['EAST', 'DB', 'SAST']")
dt_boxes_list = postprocess(dic, ratio_list)
for ino in range(img_num):
dt_boxes = dt_boxes_list[ino]
img_name = img_name_list[ino]
dt_boxes_json = []
for box in dt_boxes:
tmp_json = {"transcription": ""}
tmp_json['points'] = box.tolist()
dt_boxes_json.append(tmp_json)
otstr = img_name + "\t" + json.dumps(dt_boxes_json) + "\n"
fout.write(otstr.encode())
src_img = cv2.imread(img_name)
draw_det_res(dt_boxes, config, src_img, img_name)
logger.info("success!") logger.info("success!")
# save inference model
# paddle.jit.save(model, 'output/model')
if __name__ == '__main__': if __name__ == '__main__':
parser = program.ArgsParser() place, config = program.preprocess()
FLAGS = parser.parse_args() paddle.disable_static(place)
logger = get_logger()
print_dict(config, logger)
main() main()
tools/infer_rec.py
View file @ aad3093a
...@@ -17,160 +17,80 @@ from __future__ import division ...@@ -17,160 +17,80 @@ from __future__ import division
from __future__ import print_function from __future__ import print_function
import numpy as np import numpy as np
import os import os
import sys import sys
__dir__ = os.path.dirname(os.path.abspath(__file__)) __dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__) sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
import paddle
def set_paddle_flags(**kwargs): from ppocr.utils.logging import get_logger
for key, value in kwargs.items(): from ppocr.data import create_operators, transform
if os.environ.get(key, None) is None: from ppocr.modeling import build_model
os.environ[key] = str(value) from ppocr.postprocess import build_post_process
# NOTE(paddle-dev): All of these flags should be
# set before `import paddle`. Otherwise, it would
# not take any effect.
set_paddle_flags(
FLAGS_eager_delete_tensor_gb=0, # enable GC to save memory
)
import tools.program as program
from paddle import fluid
from ppocr.utils.utility import initial_logger
logger = initial_logger()
from ppocr.data.reader_main import reader_main
from ppocr.utils.save_load import init_model from ppocr.utils.save_load import init_model
from ppocr.utils.character import CharacterOps from ppocr.utils.utility import print_dict, get_image_file_list
from ppocr.utils.utility import create_module import tools.program as program
from ppocr.utils.utility import get_image_file_list
def main(): def main():
config = program.load_config(FLAGS.config) global_config = config['Global']
program.merge_config(FLAGS.opt)
logger.info(config) # build post process
char_ops = CharacterOps(config['Global']) post_process_class = build_post_process(config['PostProcess'],
loss_type = config['Global']['loss_type'] global_config)
config['Global']['char_ops'] = char_ops
# build model
# check if set use_gpu=True in paddlepaddle cpu version if hasattr(post_process_class, 'character'):
use_gpu = config['Global']['use_gpu'] config['Architecture']["Head"]['out_channels'] = len(
# check_gpu(use_gpu) getattr(post_process_class, 'character'))
place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() model = build_model(config['Architecture'])
exe = fluid.Executor(place)
init_model(config, model, logger)
rec_model = create_module(config['Architecture']['function'])(params=config)
startup_prog = fluid.Program() # create data ops
eval_prog = fluid.Program() transforms = []
with fluid.program_guard(eval_prog, startup_prog): for op in config['EVAL']['dataset']['transforms']:
with fluid.unique_name.guard(): op_name = list(op)[0]
_, outputs = rec_model(mode="test") if 'Label' in op_name:
fetch_name_list = list(outputs.keys()) continue
fetch_varname_list = [outputs[v].name for v in fetch_name_list] elif op_name in ['RecResizeImg']:
eval_prog = eval_prog.clone(for_test=True) op[op_name]['infer_mode'] = True
exe.run(startup_prog) elif op_name == 'keepKeys':
op[op_name]['keep_keys'] = ['image']
init_model(config, eval_prog, exe) transforms.append(op)
global_config['infer_mode'] = True
blobs = reader_main(config, 'test')() ops = create_operators(transforms, global_config)
infer_img = config['Global']['infer_img']
infer_list = get_image_file_list(infer_img) model.eval()
max_img_num = len(infer_list) for file in get_image_file_list(config['Global']['infer_img']):
if len(infer_list) == 0: logger.info("infer_img: {}".format(file))
logger.info("Can not find img in infer_img dir.") with open(file, 'rb') as f:
for i in range(max_img_num): img = f.read()
logger.info("infer_img:%s" % infer_list[i]) data = {'image': img}
img = next(blobs) batch = transform(data, ops)
if loss_type != "srn":
predict = exe.run(program=eval_prog, images = np.expand_dims(batch[0], axis=0)
feed={"image": img}, images = paddle.to_variable(images)
fetch_list=fetch_varname_list, preds = model(images)
return_numpy=False) post_result = post_process_class(preds)
else: for rec_reuslt in post_result:
encoder_word_pos_list = [] logger.info('\t result: {}'.format(rec_reuslt))
gsrm_word_pos_list = [] logger.info("success!")
gsrm_slf_attn_bias1_list = []
gsrm_slf_attn_bias2_list = [] # save inference model
encoder_word_pos_list.append(img[1]) # currently, paddle.jit.to_static not support rnn
gsrm_word_pos_list.append(img[2]) # paddle.jit.save(model, 'output/rec/model')
gsrm_slf_attn_bias1_list.append(img[3])
gsrm_slf_attn_bias2_list.append(img[4])
encoder_word_pos_list = np.concatenate(
encoder_word_pos_list, axis=0).astype(np.int64)
gsrm_word_pos_list = np.concatenate(
gsrm_word_pos_list, axis=0).astype(np.int64)
gsrm_slf_attn_bias1_list = np.concatenate(
gsrm_slf_attn_bias1_list, axis=0).astype(np.float32)
gsrm_slf_attn_bias2_list = np.concatenate(
gsrm_slf_attn_bias2_list, axis=0).astype(np.float32)
predict = exe.run(program=eval_prog, \
feed={'image': img[0], 'encoder_word_pos': encoder_word_pos_list,
'gsrm_word_pos': gsrm_word_pos_list, 'gsrm_slf_attn_bias1': gsrm_slf_attn_bias1_list,
'gsrm_slf_attn_bias2': gsrm_slf_attn_bias2_list}, \
fetch_list=fetch_varname_list, \
return_numpy=False)
if loss_type == "ctc":
preds = np.array(predict[0])
preds = preds.reshape(-1)
preds_lod = predict[0].lod()[0]
preds_text = char_ops.decode(preds)
probs = np.array(predict[1])
ind = np.argmax(probs, axis=1)
blank = probs.shape[1]
valid_ind = np.where(ind != (blank - 1))[0]
if len(valid_ind) == 0:
continue
score = np.mean(probs[valid_ind, ind[valid_ind]])
elif loss_type == "attention":
preds = np.array(predict[0])
probs = np.array(predict[1])
end_pos = np.where(preds[0, :] == 1)[0]
if len(end_pos) <= 1:
preds = preds[0, 1:]
score = np.mean(probs[0, 1:])
else:
preds = preds[0, 1:end_pos[1]]
score = np.mean(probs[0, 1:end_pos[1]])
preds = preds.reshape(-1)
preds_text = char_ops.decode(preds)
elif loss_type == "srn":
char_num = char_ops.get_char_num()
preds = np.array(predict[0])
preds = preds.reshape(-1)
probs = np.array(predict[1])
ind = np.argmax(probs, axis=1)
valid_ind = np.where(preds != int(char_num-1))[0]
if len(valid_ind) == 0:
continue
score = np.mean(probs[valid_ind, ind[valid_ind]])
preds = preds[:valid_ind[-1] + 1]
preds_text = char_ops.decode(preds)
logger.info("\t index: {}".format(preds))
logger.info("\t word : {}".format(preds_text))
logger.info("\t score: {}".format(score))
# save for inference model
target_var = []
for key, values in outputs.items():
target_var.append(values)
fluid.io.save_inference_model(
"./output/",
feeded_var_names=['image'],
target_vars=target_var,
executor=exe,
main_program=eval_prog,
model_filename="model",
params_filename="params")
if __name__ == '__main__': if __name__ == '__main__':
parser = program.ArgsParser() place, config = program.preprocess()
FLAGS = parser.parse_args() paddle.disable_static(place)
logger = get_logger()
print_dict(config, logger)
main() main()
tools/program.py
View file @ aad3093a
...@@ -16,23 +16,18 @@ from __future__ import absolute_import ...@@ -16,23 +16,18 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from argparse import ArgumentParser, RawDescriptionHelpFormatter import os
import sys import sys
import yaml import yaml
import os
from ppocr.utils.utility import create_module
from ppocr.utils.utility import initial_logger
logger = initial_logger()
import paddle.fluid as fluid
import time import time
import shutil
import paddle
import paddle.distributed as dist
from tqdm import tqdm
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from ppocr.utils.stats import TrainingStats from ppocr.utils.stats import TrainingStats
from eval_utils.eval_det_utils import eval_det_run
from eval_utils.eval_rec_utils import eval_rec_run
from ppocr.utils.save_load import save_model from ppocr.utils.save_load import save_model
import numpy as np
from ppocr.utils.character import cal_predicts_accuracy, cal_predicts_accuracy_srn, CharacterOps
class ArgsParser(ArgumentParser): class ArgsParser(ArgumentParser):
...@@ -89,13 +84,7 @@ def load_config(file_path): ...@@ -89,13 +84,7 @@ def load_config(file_path):
merge_config(default_config) merge_config(default_config)
_, ext = os.path.splitext(file_path) _, ext = os.path.splitext(file_path)
assert ext in ['.yml', '.yaml'], "only support yaml files for now" assert ext in ['.yml', '.yaml'], "only support yaml files for now"
merge_config(yaml.load(open(file_path), Loader=yaml.Loader)) merge_config(yaml.load(open(file_path, 'rb'), Loader=yaml.Loader))
assert "reader_yml" in global_config['Global'],\
"absence reader_yml in global"
reader_file_path = global_config['Global']['reader_yml']
_, ext = os.path.splitext(reader_file_path)
assert ext in ['.yml', '.yaml'], "only support yaml files for reader"
merge_config(yaml.load(open(reader_file_path), Loader=yaml.Loader))
return global_config return global_config
...@@ -139,102 +128,34 @@ def check_gpu(use_gpu): ...@@ -139,102 +128,34 @@ def check_gpu(use_gpu):
"model on CPU" "model on CPU"
try: try:
if use_gpu and not fluid.is_compiled_with_cuda(): if use_gpu and not paddle.fluid.is_compiled_with_cuda():
logger.error(err) print(err)
sys.exit(1) sys.exit(1)
except Exception as e: except Exception as e:
pass pass
def build(config, main_prog, startup_prog, mode): def train(config,
""" model,
Build a program using a model and an optimizer loss_class,
1. create feeds optimizer,
2. create a dataloader lr_scheduler,
3. create a model train_dataloader,
4. create fetchs valid_dataloader,
5. create an optimizer post_process_class,
Args: eval_class,
config(dict): config pre_best_model_dict,
main_prog(): main program logger,
startup_prog(): startup program vdl_writer=None):
is_train(bool): train or valid global_step = 0
Returns:
dataloader(): a bridge between the model and the data cal_metric_during_train = config['Global'].get('cal_metric_during_train',
fetchs(dict): dict of model outputs(included loss and measures) False)
"""
with fluid.program_guard(main_prog, startup_prog):
with fluid.unique_name.guard():
func_infor = config['Architecture']['function']
model = create_module(func_infor)(params=config)
dataloader, outputs = model(mode=mode)
fetch_name_list = list(outputs.keys())
fetch_varname_list = [outputs[v].name for v in fetch_name_list]
opt_loss_name = None
model_average = None
img_loss_name = None
word_loss_name = None
if mode == "train":
opt_loss = outputs['total_loss']
# srn loss
#img_loss = outputs['img_loss']
#word_loss = outputs['word_loss']
#img_loss_name = img_loss.name
#word_loss_name = word_loss.name
opt_params = config['Optimizer']
optimizer = create_module(opt_params['function'])(opt_params)
optimizer.minimize(opt_loss)
opt_loss_name = opt_loss.name
global_lr = optimizer._global_learning_rate()
fetch_name_list.insert(0, "lr")
fetch_varname_list.insert(0, global_lr.name)
if "loss_type" in config["Global"]:
if config['Global']["loss_type"] == 'srn':
model_average = fluid.optimizer.ModelAverage(
config['Global']['average_window'],
min_average_window=config['Global'][
'min_average_window'],
max_average_window=config['Global'][
'max_average_window'])
return (dataloader, fetch_name_list, fetch_varname_list, opt_loss_name,
model_average)
def build_export(config, main_prog, startup_prog):
"""
"""
with fluid.program_guard(main_prog, startup_prog):
with fluid.unique_name.guard():
func_infor = config['Architecture']['function']
model = create_module(func_infor)(params=config)
image, outputs = model(mode='export')
fetches_var_name = sorted([name for name in outputs.keys()])
fetches_var = [outputs[name] for name in fetches_var_name]
feeded_var_names = [image.name]
target_vars = fetches_var
return feeded_var_names, target_vars, fetches_var_name
def create_multi_devices_program(program, loss_var_name):
build_strategy = fluid.BuildStrategy()
build_strategy.memory_optimize = False
build_strategy.enable_inplace = True
exec_strategy = fluid.ExecutionStrategy()
exec_strategy.num_iteration_per_drop_scope = 1
compile_program = fluid.CompiledProgram(program).with_data_parallel(
loss_name=loss_var_name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
return compile_program
def train_eval_det_run(config, exe, train_info_dict, eval_info_dict):
train_batch_id = 0
log_smooth_window = config['Global']['log_smooth_window'] log_smooth_window = config['Global']['log_smooth_window']
epoch_num = config['Global']['epoch_num'] epoch_num = config['Global']['epoch_num']
print_batch_step = config['Global']['print_batch_step'] print_batch_step = config['Global']['print_batch_step']
eval_batch_step = config['Global']['eval_batch_step'] eval_batch_step = config['Global']['eval_batch_step']
start_eval_step = 0 start_eval_step = 0
if type(eval_batch_step) == list and len(eval_batch_step) >= 2: if type(eval_batch_step) == list and len(eval_batch_step) >= 2:
start_eval_step = eval_batch_step[0] start_eval_step = eval_batch_step[0]
...@@ -246,180 +167,173 @@ def train_eval_det_run(config, exe, train_info_dict, eval_info_dict): ...@@ -246,180 +167,173 @@ def train_eval_det_run(config, exe, train_info_dict, eval_info_dict):
save_model_dir = config['Global']['save_model_dir'] save_model_dir = config['Global']['save_model_dir']
if not os.path.exists(save_model_dir): if not os.path.exists(save_model_dir):
os.makedirs(save_model_dir) os.makedirs(save_model_dir)
train_stats = TrainingStats(log_smooth_window, main_indicator = eval_class.main_indicator
train_info_dict['fetch_name_list']) best_model_dict = {main_indicator: 0}
best_eval_hmean = -1 best_model_dict.update(pre_best_model_dict)
best_batch_id = 0 train_stats = TrainingStats(log_smooth_window, ['lr'])
best_epoch = 0 model.train()
train_loader = train_info_dict['reader']
for epoch in range(epoch_num): if 'start_epoch' in best_model_dict:
train_loader.start() start_epoch = best_model_dict['start_epoch']
try: else:
while True: start_epoch = 0
t1 = time.time()
train_outs = exe.run( for epoch in range(start_epoch, epoch_num):
program=train_info_dict['compile_program'], for idx, batch in enumerate(train_dataloader):
fetch_list=train_info_dict['fetch_varname_list'], if idx >= len(train_dataloader):
return_numpy=False) break
stats = {} if not isinstance(lr_scheduler, float):
for tno in range(len(train_outs)): lr_scheduler.step()
fetch_name = train_info_dict['fetch_name_list'][tno] lr = optimizer.get_lr()
fetch_value = np.mean(np.array(train_outs[tno])) t1 = time.time()
stats[fetch_name] = fetch_value batch = [paddle.to_variable(x) for x in batch]
t2 = time.time() images = batch[0]
train_batch_elapse = t2 - t1 preds = model(images)
train_stats.update(stats) loss = loss_class(preds, batch)
if train_batch_id > 0 and train_batch_id \ avg_loss = loss['loss']
% print_batch_step == 0: if config['Global']['distributed']:
logs = train_stats.log() avg_loss = model.scale_loss(avg_loss)
strs = 'epoch: {}, iter: {}, {}, time: {:.3f}'.format( avg_loss.backward()
epoch, train_batch_id, logs, train_batch_elapse) model.apply_collective_grads()
logger.info(strs) else:
avg_loss.backward()
if train_batch_id > start_eval_step and\ optimizer.step()
(train_batch_id - start_eval_step) % eval_batch_step == 0: optimizer.clear_grad()
metrics = eval_det_run(exe, config, eval_info_dict, "eval")
hmean = metrics['hmean'] # logger and visualdl
if hmean >= best_eval_hmean: stats = {k: v.numpy().mean() for k, v in loss.items()}
best_eval_hmean = hmean stats['lr'] = lr
best_batch_id = train_batch_id train_stats.update(stats)
best_epoch = epoch
save_path = save_model_dir + "/best_accuracy" if cal_metric_during_train: # onlt rec and cls need
save_model(train_info_dict['train_program'], save_path) batch = [item.numpy() for item in batch]
strs = 'Test iter: {}, metrics:{}, best_hmean:{:.6f}, best_epoch:{}, best_batch_id:{}'.format( post_result = post_process_class(preds, batch[1])
train_batch_id, metrics, best_eval_hmean, best_epoch, eval_class(post_result, batch)
best_batch_id) metirc = eval_class.get_metric()
logger.info(strs) train_stats.update(metirc)
train_batch_id += 1
t2 = time.time()
except fluid.core.EOFException: train_batch_elapse = t2 - t1
train_loader.reset()
if epoch == 0 and save_epoch_step == 1: if vdl_writer is not None and dist.get_rank() == 0:
save_path = save_model_dir + "/iter_epoch_0" for k, v in train_stats.get().items():
save_model(train_info_dict['train_program'], save_path) vdl_writer.add_scalar('TRAIN/{}'.format(k), v, global_step)
if epoch > 0 and epoch % save_epoch_step == 0: vdl_writer.add_scalar('TRAIN/lr', lr, global_step)
save_path = save_model_dir + "/iter_epoch_%d" % (epoch)
save_model(train_info_dict['train_program'], save_path) if global_step > 0 and global_step % print_batch_step == 0:
logs = train_stats.log()
strs = 'epoch: [{}/{}], iter: {}, {}, time: {:.3f}'.format(
epoch, epoch_num, global_step, logs, train_batch_elapse)
logger.info(strs)
# eval
if global_step > start_eval_step and \
(global_step - start_eval_step) % eval_batch_step == 0 and dist.get_rank() == 0:
cur_metirc = eval(model, valid_dataloader, post_process_class,
eval_class)
cur_metirc_str = 'cur metirc, {}'.format(', '.join(
['{}: {}'.format(k, v) for k, v in cur_metirc.items()]))
logger.info(cur_metirc_str)
# logger metric
if vdl_writer is not None:
for k, v in cur_metirc.items():
if isinstance(v, (float, int)):
vdl_writer.add_scalar('EVAL/{}'.format(k),
cur_metirc[k], global_step)
if cur_metirc[main_indicator] >= best_model_dict[
main_indicator]:
best_model_dict.update(cur_metirc)
best_model_dict['best_epoch'] = epoch
save_model(
model,
optimizer,
save_model_dir,
logger,
is_best=True,
prefix='best_accuracy',
best_model_dict=best_model_dict,
epoch=epoch)
best_str = 'best metirc, {}'.format(', '.join([
'{}: {}'.format(k, v) for k, v in best_model_dict.items()
]))
logger.info(best_str)
# logger best metric
if vdl_writer is not None:
vdl_writer.add_scalar('EVAL/best_{}'.format(main_indicator),
best_model_dict[main_indicator],
global_step)
global_step += 1
if dist.get_rank() == 0:
save_model(
model,
optimizer,
save_model_dir,
logger,
is_best=False,
prefix='latest',
best_model_dict=best_model_dict,
epoch=epoch)
if dist.get_rank() == 0 and epoch > 0 and epoch % save_epoch_step == 0:
save_model(
model,
optimizer,
save_model_dir,
logger,
is_best=False,
prefix='iter_epoch_{}'.format(epoch),
best_model_dict=best_model_dict,
epoch=epoch)
best_str = 'best metirc, {}'.format(', '.join(
['{}: {}'.format(k, v) for k, v in best_model_dict.items()]))
logger.info(best_str)
if dist.get_rank() == 0 and vdl_writer is not None:
vdl_writer.close()
return return
def train_eval_rec_run(config, exe, train_info_dict, eval_info_dict): def eval(model, valid_dataloader, post_process_class, eval_class):
train_batch_id = 0 model.eval()
log_smooth_window = config['Global']['log_smooth_window'] with paddle.no_grad():
epoch_num = config['Global']['epoch_num'] total_frame = 0.0
print_batch_step = config['Global']['print_batch_step'] total_time = 0.0
eval_batch_step = config['Global']['eval_batch_step'] pbar = tqdm(total=len(valid_dataloader), desc='eval model: ')
start_eval_step = 0 for idx, batch in enumerate(valid_dataloader):
if type(eval_batch_step) == list and len(eval_batch_step) >= 2: if idx >= len(valid_dataloader):
start_eval_step = eval_batch_step[0] break
eval_batch_step = eval_batch_step[1] images = paddle.to_variable(batch[0])
logger.info( start = time.time()
"During the training process, after the {}th iteration, an evaluation is run every {} iterations". preds = model(images)
format(start_eval_step, eval_batch_step))
save_epoch_step = config['Global']['save_epoch_step'] batch = [item.numpy() for item in batch]
save_model_dir = config['Global']['save_model_dir'] # Obtain usable results from post-processing methods
if not os.path.exists(save_model_dir): post_result = post_process_class(preds, batch[1])
os.makedirs(save_model_dir) total_time += time.time() - start
train_stats = TrainingStats(log_smooth_window, ['loss', 'acc']) # Evaluate the results of the current batch
best_eval_acc = -1 eval_class(post_result, batch)
best_batch_id = 0 pbar.update(1)
best_epoch = 0 total_frame += len(images)
train_loader = train_info_dict['reader'] # Get final metirc，eg. acc or hmean
for epoch in range(epoch_num): metirc = eval_class.get_metric()
train_loader.start() pbar.close()
try: model.train()
while True: metirc['fps'] = total_frame / total_time
t1 = time.time() return metirc
train_outs = exe.run(
program=train_info_dict['compile_program'],
fetch_list=train_info_dict['fetch_varname_list'],
return_numpy=False)
fetch_map = dict(
zip(train_info_dict['fetch_name_list'],
range(len(train_outs))))
loss = np.mean(np.array(train_outs[fetch_map['total_loss']]))
lr = np.mean(np.array(train_outs[fetch_map['lr']]))
preds_idx = fetch_map['decoded_out']
preds = np.array(train_outs[preds_idx])
labels_idx = fetch_map['label']
labels = np.array(train_outs[labels_idx])
if config['Global']['loss_type'] != 'srn':
preds_lod = train_outs[preds_idx].lod()[0]
labels_lod = train_outs[labels_idx].lod()[0]
acc, acc_num, img_num = cal_predicts_accuracy(
config['Global']['char_ops'], preds, preds_lod, labels,
labels_lod)
else:
acc, acc_num, img_num = cal_predicts_accuracy_srn(
config['Global']['char_ops'], preds, labels,
config['Global']['max_text_length'])
t2 = time.time()
train_batch_elapse = t2 - t1
stats = {'loss': loss, 'acc': acc}
train_stats.update(stats)
if train_batch_id > start_eval_step and (train_batch_id - start_eval_step) \
% print_batch_step == 0:
logs = train_stats.log()
strs = 'epoch: {}, iter: {}, lr: {:.6f}, {}, time: {:.3f}'.format(
epoch, train_batch_id, lr, logs, train_batch_elapse)
logger.info(strs)
if train_batch_id > 0 and\
train_batch_id % eval_batch_step == 0:
model_average = train_info_dict['model_average']
if model_average != None:
model_average.apply(exe)
metrics = eval_rec_run(exe, config, eval_info_dict, "eval")
eval_acc = metrics['avg_acc']
eval_sample_num = metrics['total_sample_num']
if eval_acc > best_eval_acc:
best_eval_acc = eval_acc
best_batch_id = train_batch_id
best_epoch = epoch
save_path = save_model_dir + "/best_accuracy"
save_model(train_info_dict['train_program'], save_path)
strs = 'Test iter: {}, acc:{:.6f}, best_acc:{:.6f}, best_epoch:{}, best_batch_id:{}, eval_sample_num:{}'.format(
train_batch_id, eval_acc, best_eval_acc, best_epoch,
best_batch_id, eval_sample_num)
logger.info(strs)
train_batch_id += 1
except fluid.core.EOFException:
train_loader.reset()
if epoch == 0 and save_epoch_step == 1:
save_path = save_model_dir + "/iter_epoch_0"
save_model(train_info_dict['train_program'], save_path)
if epoch > 0 and epoch % save_epoch_step == 0:
save_path = save_model_dir + "/iter_epoch_%d" % (epoch)
save_model(train_info_dict['train_program'], save_path)
return
def preprocess(): def preprocess():
FLAGS = ArgsParser().parse_args() FLAGS = ArgsParser().parse_args()
config = load_config(FLAGS.config) config = load_config(FLAGS.config)
merge_config(FLAGS.opt) merge_config(FLAGS.opt)
logger.info(config)
# check if set use_gpu=True in paddlepaddle cpu version # check if set use_gpu=True in paddlepaddle cpu version
use_gpu = config['Global']['use_gpu'] use_gpu = config['Global']['use_gpu']
check_gpu(use_gpu) check_gpu(use_gpu)
alg = config['Global']['algorithm'] alg = config['Architecture']['algorithm']
assert alg in ['EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN'] assert alg in [
if alg in ['Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN']: 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN'
config['Global']['char_ops'] = CharacterOps(config['Global']) ]
place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
startup_program = fluid.Program()
train_program = fluid.Program()
if alg in ['EAST', 'DB', 'SAST']:
train_alg_type = 'det'
else:
train_alg_type = 'rec'
return startup_program, train_program, place, config, train_alg_type device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'
device = paddle.set_device(device)
return device, config
tools/train.py
View file @ aad3093a
...@@ -18,107 +18,122 @@ from __future__ import print_function ...@@ -18,107 +18,122 @@ from __future__ import print_function
import os import os
import sys import sys
__dir__ = os.path.dirname(os.path.abspath(__file__)) __dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__) sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..'))) sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
import yaml
import paddle
import paddle.distributed as dist
def set_paddle_flags(**kwargs): paddle.manual_seed(2)
for key, value in kwargs.items():
if os.environ.get(key, None) is None:
os.environ[key] = str(value)
from ppocr.utils.logging import get_logger
from ppocr.data import build_dataloader
from ppocr.modeling import build_model, build_loss
from ppocr.optimizer import build_optimizer
from ppocr.postprocess import build_post_process
from ppocr.metrics import build_metric
from ppocr.utils.save_load import init_model
from ppocr.utils.utility import print_dict
import tools.program as program
# NOTE(paddle-dev): All of these flags should be dist.get_world_size()
# set before `import paddle`. Otherwise, it would
# not take any effect.
set_paddle_flags(
FLAGS_eager_delete_tensor_gb=0, # enable GC to save memory
)
import tools.program as program
from paddle import fluid
from ppocr.utils.utility import initial_logger
logger = initial_logger()
from ppocr.data.reader_main import reader_main
from ppocr.utils.save_load import init_model
from paddle.fluid.contrib.model_stat import summary
def main():
train_build_outputs = program.build(
config, train_program, startup_program, mode='train')
train_loader = train_build_outputs[0]
train_fetch_name_list = train_build_outputs[1]
train_fetch_varname_list = train_build_outputs[2]
train_opt_loss_name = train_build_outputs[3]
model_average = train_build_outputs[-1]
eval_program = fluid.Program()
eval_build_outputs = program.build(
config, eval_program, startup_program, mode='eval')
eval_fetch_name_list = eval_build_outputs[1]
eval_fetch_varname_list = eval_build_outputs[2]
eval_program = eval_program.clone(for_test=True)
train_reader = reader_main(config=config, mode="train")
train_loader.set_sample_list_generator(train_reader, places=place)
eval_reader = reader_main(config=config, mode="eval")
exe = fluid.Executor(place)
exe.run(startup_program)
# compile program for multi-devices
train_compile_program = program.create_multi_devices_program(
train_program, train_opt_loss_name)
# dump mode structure
if config['Global']['debug']:
if train_alg_type == 'rec' and 'attention' in config['Global']['loss_type']:
logger.warning('Does not suport dump attention...')
else:
summary(train_program)
init_model(config, train_program, exe)
train_info_dict = {'compile_program':train_compile_program,\
'train_program':train_program,\
'reader':train_loader,\
'fetch_name_list':train_fetch_name_list,\
'fetch_varname_list':train_fetch_varname_list,\
'model_average': model_average}
eval_info_dict = {'program':eval_program,\
'reader':eval_reader,\
'fetch_name_list':eval_fetch_name_list,\
'fetch_varname_list':eval_fetch_varname_list}
if train_alg_type == 'det':
program.train_eval_det_run(config, exe, train_info_dict, eval_info_dict)
else:
program.train_eval_rec_run(config, exe, train_info_dict, eval_info_dict)
def main(config, device, logger, vdl_writer):
# init dist environment
if config['Global']['distributed']:
dist.init_parallel_env()
def test_reader(): global_config = config['Global']
logger.info(config) # build dataloader
train_reader = reader_main(config=config, mode="train") train_loader, train_info_dict = build_dataloader(
config['TRAIN'], device, global_config['distributed'], global_config)
if config['EVAL']:
eval_loader, _ = build_dataloader(config['EVAL'], device, False,
global_config)
else:
eval_loader = None
# build post process
post_process_class = build_post_process(config['PostProcess'],
global_config)
# build model
# for rec algorithm
if hasattr(post_process_class, 'character'):
config['Architecture']["Head"]['out_channels'] = len(
getattr(post_process_class, 'character'))
model = build_model(config['Architecture'])
if config['Global']['distributed']:
model = paddle.DataParallel(model)
# build optim
optimizer, lr_scheduler = build_optimizer(
config['Optimizer'],
epochs=config['Global']['epoch_num'],
step_each_epoch=len(train_loader),
parameters=model.parameters())
best_model_dict = init_model(config, model, logger, optimizer)
# build loss
loss_class = build_loss(config['Loss'])
# build metric
eval_class = build_metric(config['Metric'])
# start train
program.train(config, model, loss_class, optimizer, lr_scheduler,
train_loader, eval_loader, post_process_class, eval_class,
best_model_dict, logger, vdl_writer)
def test_reader(config, place, logger):
train_loader = build_dataloader(config['TRAIN'], place)
import time import time
starttime = time.time() starttime = time.time()
count = 0 count = 0
try: try:
for data in train_reader(): for data in train_loader():
count += 1 count += 1
if count % 1 == 0: if count % 1 == 0:
batch_time = time.time() - starttime batch_time = time.time() - starttime
starttime = time.time() starttime = time.time()
logger.info("reader:", count, len(data), batch_time) logger.info("reader: {}, {}, {}".format(count,
len(data), batch_time))
except Exception as e: except Exception as e:
logger.info(e) logger.info(e)
logger.info("finish reader: {}, Success!".format(count)) logger.info("finish reader: {}, Success!".format(count))
def dis_main():
device, config = program.preprocess()
config['Global']['distributed'] = dist.get_world_size() != 1
paddle.disable_static(device)
# save_config
os.makedirs(config['Global']['save_model_dir'], exist_ok=True)
with open(
os.path.join(config['Global']['save_model_dir'], 'config.yml'),
'w') as f:
yaml.dump(dict(config), f, default_flow_style=False, sort_keys=False)
logger = get_logger(
log_file='{}/train.log'.format(config['Global']['save_model_dir']))
if config['Global']['use_visualdl']:
from visualdl import LogWriter
vdl_writer = LogWriter(logdir=config['Global']['save_model_dir'])
else:
vdl_writer = None
print_dict(config, logger)
logger.info('train with paddle {} and device {}'.format(paddle.__version__,
device))
main(config, device, logger, vdl_writer)
# test_reader(config, place, logger)
if __name__ == '__main__': if __name__ == '__main__':
startup_program, train_program, place, config, train_alg_type = program.preprocess() # main()
main() # dist.spawn(dis_main, nprocs=2, selelcted_gpus='6,7')
# test_reader() dis_main()
train.sh 0 → 100644
View file @ aad3093a
python -m paddle.distributed.launch --selected_gpus '0,1,2,3,4,5,6,7' tools/train.py -c configs/det/det_mv3_db.yml
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment