add_rec_sar, test=dygraph

This commit is contained in:
andyjpaddle
2021-08-24 03:49:26 +00:00
parent ffa94415c3
commit 8a95b3352d
15 changed files with 207 additions and 13 deletions
+8 -5
View File
@@ -186,6 +186,7 @@ def train(config,
model.train()
use_srn = config['Architecture']['algorithm'] == "SRN"
use_sar = config['Architecture']['algorithm'] == 'SAR'
try:
model_type = config['Architecture']['model_type']
except:
@@ -213,7 +214,7 @@ def train(config,
images = batch[0]
if use_srn:
model_average = True
if use_srn or model_type == 'table':
if use_srn or model_type == 'table' or use_sar:
preds = model(images, data=batch[1:])
else:
preds = model(images)
@@ -277,7 +278,8 @@ def train(config,
post_process_class,
eval_class,
model_type,
use_srn=use_srn)
use_srn=use_srn,
use_sar=use_sar)
cur_metric_str = 'cur metric, {}'.format(', '.join(
['{}: {}'.format(k, v) for k, v in cur_metric.items()]))
logger.info(cur_metric_str)
@@ -349,7 +351,8 @@ def eval(model,
post_process_class,
eval_class,
model_type,
use_srn=False):
use_srn=False,
use_sar=False):
model.eval()
with paddle.no_grad():
total_frame = 0.0
@@ -362,7 +365,7 @@ def eval(model,
break
images = batch[0]
start = time.time()
if use_srn or model_type == 'table':
if use_srn or model_type == 'table' or use_sar:
preds = model(images, data=batch[1:])
else:
preds = model(images)
@@ -398,7 +401,7 @@ def preprocess(is_train=False):
alg = config['Architecture']['algorithm']
assert alg in [
'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
'CLS', 'PGNet', 'Distillation', 'TableAttn'
'CLS', 'PGNet', 'Distillation', 'TableAttn', 'SAR'
]
device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'