trainer.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import os
  18. import sys
  19. import copy
  20. import time
  21. import numpy as np
  22. import typing
  23. from PIL import Image, ImageOps
  24. import paddle
  25. import paddle.distributed as dist
  26. from paddle.distributed import fleet
  27. from paddle import amp
  28. from paddle.static import InputSpec
  29. from paddlex.ppdet.optimizer import ModelEMA
  30. from paddlex.ppdet.core.workspace import create
  31. from paddlex.ppdet.modeling.architectures.meta_arch import BaseArch
  32. from paddlex.ppdet.utils.checkpoint import load_weight, load_pretrain_weight
  33. from paddlex.ppdet.utils.visualizer import visualize_results, save_result
  34. from paddlex.ppdet.metrics import Metric, COCOMetric, VOCMetric, WiderFaceMetric, get_infer_results, KeyPointTopDownCOCOEval, KeyPointTopDownMPIIEval
  35. from paddlex.ppdet.metrics import RBoxMetric, JDEDetMetric, SNIPERCOCOMetric
  36. from paddlex.ppdet.data.source.sniper_coco import SniperCOCODataSet
  37. from paddlex.ppdet.data.source.category import get_categories
  38. from paddlex.ppdet.utils import stats
  39. from paddlex.ppdet.utils import profiler
  40. from .callbacks import Callback, ComposeCallback, LogPrinter, Checkpointer, WiferFaceEval, VisualDLWriter, SniperProposalsGenerator
  41. from .export_utils import _dump_infer_config, _prune_input_spec
  42. from paddlex.ppdet.utils.logger import setup_logger
  43. logger = setup_logger('ppdet.engine')
  44. __all__ = ['Trainer']
  45. MOT_ARCH = ['DeepSORT', 'JDE', 'FairMOT']
  46. class Trainer(object):
  47. def __init__(self, cfg, mode='train'):
  48. self.cfg = cfg
  49. assert mode.lower() in ['train', 'eval', 'test'], \
  50. "mode should be 'train', 'eval' or 'test'"
  51. self.mode = mode.lower()
  52. self.optimizer = None
  53. self.is_loaded_weights = False
  54. # build data loader
  55. if cfg.architecture in MOT_ARCH and self.mode in ['eval', 'test']:
  56. self.dataset = cfg['{}MOTDataset'.format(self.mode.capitalize())]
  57. else:
  58. self.dataset = cfg['{}Dataset'.format(self.mode.capitalize())]
  59. if cfg.architecture == 'DeepSORT' and self.mode == 'train':
  60. logger.error('DeepSORT has no need of training on mot dataset.')
  61. sys.exit(1)
  62. if self.mode == 'train':
  63. self.loader = create('{}Reader'.format(self.mode.capitalize()))(
  64. self.dataset, cfg.worker_num)
  65. if cfg.architecture == 'JDE' and self.mode == 'train':
  66. cfg['JDEEmbeddingHead'][
  67. 'num_identities'] = self.dataset.num_identities_dict[0]
  68. # JDE only support single class MOT now.
  69. if cfg.architecture == 'FairMOT' and self.mode == 'train':
  70. cfg['FairMOTEmbeddingHead'][
  71. 'num_identities_dict'] = self.dataset.num_identities_dict
  72. # FairMOT support single class and multi-class MOT now.
  73. # build model
  74. if 'model' not in self.cfg:
  75. self.model = create(cfg.architecture)
  76. else:
  77. self.model = self.cfg.model
  78. self.is_loaded_weights = True
  79. #normalize params for deploy
  80. self.model.load_meanstd(cfg['TestReader']['sample_transforms'])
  81. self.use_ema = ('use_ema' in cfg and cfg['use_ema'])
  82. if self.use_ema:
  83. ema_decay = self.cfg.get('ema_decay', 0.9998)
  84. cycle_epoch = self.cfg.get('cycle_epoch', -1)
  85. self.ema = ModelEMA(
  86. self.model,
  87. decay=ema_decay,
  88. use_thres_step=True,
  89. cycle_epoch=cycle_epoch)
  90. # EvalDataset build with BatchSampler to evaluate in single device
  91. # TODO: multi-device evaluate
  92. if self.mode == 'eval':
  93. self._eval_batch_sampler = paddle.io.BatchSampler(
  94. self.dataset, batch_size=self.cfg.EvalReader['batch_size'])
  95. reader_name = '{}Reader'.format(self.mode.capitalize())
  96. # If metric is VOC, need to be set collate_batch=False.
  97. if cfg.metric == 'VOC':
  98. cfg[reader_name]['collate_batch'] = False
  99. self.loader = create(reader_name)(self.dataset, cfg.worker_num,
  100. self._eval_batch_sampler)
  101. # TestDataset build after user set images, skip loader creation here
  102. # build optimizer in train mode
  103. if self.mode == 'train':
  104. steps_per_epoch = len(self.loader)
  105. self.lr = create('LearningRate')(steps_per_epoch)
  106. self.optimizer = create('OptimizerBuilder')(self.lr, self.model)
  107. if self.cfg.get('unstructured_prune'):
  108. self.pruner = create('UnstructuredPruner')(self.model,
  109. steps_per_epoch)
  110. self._nranks = dist.get_world_size()
  111. self._local_rank = dist.get_rank()
  112. self.status = {}
  113. self.start_epoch = 0
  114. self.end_epoch = 0 if 'epoch' not in cfg else cfg.epoch
  115. # initial default callbacks
  116. self._init_callbacks()
  117. # initial default metrics
  118. self._init_metrics()
  119. self._reset_metrics()
  120. def _init_callbacks(self):
  121. if self.mode == 'train':
  122. self._callbacks = [LogPrinter(self), Checkpointer(self)]
  123. if self.cfg.get('use_vdl', False):
  124. self._callbacks.append(VisualDLWriter(self))
  125. if self.cfg.get('save_proposals', False):
  126. self._callbacks.append(SniperProposalsGenerator(self))
  127. self._compose_callback = ComposeCallback(self._callbacks)
  128. elif self.mode == 'eval':
  129. self._callbacks = [LogPrinter(self)]
  130. if self.cfg.metric == 'WiderFace':
  131. self._callbacks.append(WiferFaceEval(self))
  132. self._compose_callback = ComposeCallback(self._callbacks)
  133. elif self.mode == 'test' and self.cfg.get('use_vdl', False):
  134. self._callbacks = [VisualDLWriter(self)]
  135. self._compose_callback = ComposeCallback(self._callbacks)
  136. else:
  137. self._callbacks = []
  138. self._compose_callback = None
  139. def _init_metrics(self, validate=False):
  140. if self.mode == 'test' or (self.mode == 'train' and not validate):
  141. self._metrics = []
  142. return
  143. classwise = self.cfg['classwise'] if 'classwise' in self.cfg else False
  144. if self.cfg.metric == 'COCO' or self.cfg.metric == "SNIPERCOCO":
  145. # TODO: bias should be unified
  146. bias = self.cfg['bias'] if 'bias' in self.cfg else 0
  147. output_eval = self.cfg['output_eval'] \
  148. if 'output_eval' in self.cfg else None
  149. save_prediction_only = self.cfg.get('save_prediction_only', False)
  150. # pass clsid2catid info to metric instance to avoid multiple loading
  151. # annotation file
  152. clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()} \
  153. if self.mode == 'eval' else None
  154. # when do validation in train, annotation file should be get from
  155. # EvalReader instead of self.dataset(which is TrainReader)
  156. anno_file = self.dataset.get_anno()
  157. dataset = self.dataset
  158. if self.mode == 'train' and validate:
  159. eval_dataset = self.cfg['EvalDataset']
  160. eval_dataset.check_or_download_dataset()
  161. anno_file = eval_dataset.get_anno()
  162. dataset = eval_dataset
  163. IouType = self.cfg['IouType'] if 'IouType' in self.cfg else 'bbox'
  164. if self.cfg.metric == "COCO":
  165. self._metrics = [
  166. COCOMetric(
  167. anno_file=anno_file,
  168. clsid2catid=clsid2catid,
  169. classwise=classwise,
  170. output_eval=output_eval,
  171. bias=bias,
  172. IouType=IouType,
  173. save_prediction_only=save_prediction_only)
  174. ]
  175. elif self.cfg.metric == "SNIPERCOCO": # sniper
  176. self._metrics = [
  177. SNIPERCOCOMetric(
  178. anno_file=anno_file,
  179. dataset=dataset,
  180. clsid2catid=clsid2catid,
  181. classwise=classwise,
  182. output_eval=output_eval,
  183. bias=bias,
  184. IouType=IouType,
  185. save_prediction_only=save_prediction_only)
  186. ]
  187. elif self.cfg.metric == 'RBOX':
  188. # TODO: bias should be unified
  189. bias = self.cfg['bias'] if 'bias' in self.cfg else 0
  190. output_eval = self.cfg['output_eval'] \
  191. if 'output_eval' in self.cfg else None
  192. save_prediction_only = self.cfg.get('save_prediction_only', False)
  193. # pass clsid2catid info to metric instance to avoid multiple loading
  194. # annotation file
  195. clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()} \
  196. if self.mode == 'eval' else None
  197. # when do validation in train, annotation file should be get from
  198. # EvalReader instead of self.dataset(which is TrainReader)
  199. anno_file = self.dataset.get_anno()
  200. if self.mode == 'train' and validate:
  201. eval_dataset = self.cfg['EvalDataset']
  202. eval_dataset.check_or_download_dataset()
  203. anno_file = eval_dataset.get_anno()
  204. self._metrics = [
  205. RBoxMetric(
  206. anno_file=anno_file,
  207. clsid2catid=clsid2catid,
  208. classwise=classwise,
  209. output_eval=output_eval,
  210. bias=bias,
  211. save_prediction_only=save_prediction_only)
  212. ]
  213. elif self.cfg.metric == 'VOC':
  214. self._metrics = [
  215. VOCMetric(
  216. label_list=self.dataset.get_label_list(),
  217. class_num=self.cfg.num_classes,
  218. map_type=self.cfg.map_type,
  219. classwise=classwise)
  220. ]
  221. elif self.cfg.metric == 'WiderFace':
  222. multi_scale = self.cfg.multi_scale_eval if 'multi_scale_eval' in self.cfg else True
  223. self._metrics = [
  224. WiderFaceMetric(
  225. image_dir=os.path.join(self.dataset.dataset_dir,
  226. self.dataset.image_dir),
  227. anno_file=self.dataset.get_anno(),
  228. multi_scale=multi_scale)
  229. ]
  230. elif self.cfg.metric == 'KeyPointTopDownCOCOEval':
  231. eval_dataset = self.cfg['EvalDataset']
  232. eval_dataset.check_or_download_dataset()
  233. anno_file = eval_dataset.get_anno()
  234. save_prediction_only = self.cfg.get('save_prediction_only', False)
  235. self._metrics = [
  236. KeyPointTopDownCOCOEval(
  237. anno_file,
  238. len(eval_dataset),
  239. self.cfg.num_joints,
  240. self.cfg.save_dir,
  241. save_prediction_only=save_prediction_only)
  242. ]
  243. elif self.cfg.metric == 'KeyPointTopDownMPIIEval':
  244. eval_dataset = self.cfg['EvalDataset']
  245. eval_dataset.check_or_download_dataset()
  246. anno_file = eval_dataset.get_anno()
  247. save_prediction_only = self.cfg.get('save_prediction_only', False)
  248. self._metrics = [
  249. KeyPointTopDownMPIIEval(
  250. anno_file,
  251. len(eval_dataset),
  252. self.cfg.num_joints,
  253. self.cfg.save_dir,
  254. save_prediction_only=save_prediction_only)
  255. ]
  256. elif self.cfg.metric == 'MOTDet':
  257. self._metrics = [JDEDetMetric(), ]
  258. else:
  259. logger.warning("Metric not support for metric type {}".format(
  260. self.cfg.metric))
  261. self._metrics = []
  262. def _reset_metrics(self):
  263. for metric in self._metrics:
  264. metric.reset()
  265. def register_callbacks(self, callbacks):
  266. callbacks = [c for c in list(callbacks) if c is not None]
  267. for c in callbacks:
  268. assert isinstance(c, Callback), \
  269. "metrics shoule be instances of subclass of Metric"
  270. self._callbacks.extend(callbacks)
  271. self._compose_callback = ComposeCallback(self._callbacks)
  272. def register_metrics(self, metrics):
  273. metrics = [m for m in list(metrics) if m is not None]
  274. for m in metrics:
  275. assert isinstance(m, Metric), \
  276. "metrics shoule be instances of subclass of Metric"
  277. self._metrics.extend(metrics)
  278. def load_weights(self, weights):
  279. if self.is_loaded_weights:
  280. return
  281. self.start_epoch = 0
  282. load_pretrain_weight(self.model, weights)
  283. logger.debug("Load weights {} to start training".format(weights))
  284. def load_weights_sde(self, det_weights, reid_weights):
  285. if self.model.detector:
  286. load_weight(self.model.detector, det_weights)
  287. load_weight(self.model.reid, reid_weights)
  288. else:
  289. load_weight(self.model.reid, reid_weights)
  290. def resume_weights(self, weights):
  291. # support Distill resume weights
  292. if hasattr(self.model, 'student_model'):
  293. self.start_epoch = load_weight(self.model.student_model, weights,
  294. self.optimizer)
  295. else:
  296. self.start_epoch = load_weight(self.model, weights, self.optimizer)
  297. logger.debug("Resume weights of epoch {}".format(self.start_epoch))
  298. def train(self, validate=False):
  299. assert self.mode == 'train', "Model not in 'train' mode"
  300. Init_mark = False
  301. sync_bn = (
  302. getattr(self.cfg, 'norm_type', None) in [None, 'sync_bn'] and
  303. self.cfg.use_gpu and self._nranks > 1)
  304. if sync_bn:
  305. self.model = BaseArch.convert_sync_batchnorm(self.model)
  306. model = self.model
  307. if self.cfg.get('fleet', False):
  308. model = fleet.distributed_model(model)
  309. self.optimizer = fleet.distributed_optimizer(self.optimizer)
  310. elif self._nranks > 1:
  311. find_unused_parameters = self.cfg[
  312. 'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
  313. model = paddle.DataParallel(
  314. self.model, find_unused_parameters=find_unused_parameters)
  315. # initial fp16
  316. if self.cfg.get('fp16', False):
  317. scaler = amp.GradScaler(
  318. enable=self.cfg.use_gpu, init_loss_scaling=1024)
  319. self.status.update({
  320. 'epoch_id': self.start_epoch,
  321. 'step_id': 0,
  322. 'steps_per_epoch': len(self.loader)
  323. })
  324. self.status['batch_time'] = stats.SmoothedValue(
  325. self.cfg.log_iter, fmt='{avg:.4f}')
  326. self.status['data_time'] = stats.SmoothedValue(
  327. self.cfg.log_iter, fmt='{avg:.4f}')
  328. self.status['training_staus'] = stats.TrainingStats(self.cfg.log_iter)
  329. if self.cfg.get('print_flops', False):
  330. flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
  331. self.dataset, self.cfg.worker_num)
  332. self._flops(flops_loader)
  333. profiler_options = self.cfg.get('profiler_options', None)
  334. self._compose_callback.on_train_begin(self.status)
  335. for epoch_id in range(self.start_epoch, self.cfg.epoch):
  336. self.status['mode'] = 'train'
  337. self.status['epoch_id'] = epoch_id
  338. self._compose_callback.on_epoch_begin(self.status)
  339. self.loader.dataset.set_epoch(epoch_id)
  340. model.train()
  341. iter_tic = time.time()
  342. for step_id, data in enumerate(self.loader):
  343. self.status['data_time'].update(time.time() - iter_tic)
  344. self.status['step_id'] = step_id
  345. profiler.add_profiler_step(profiler_options)
  346. self._compose_callback.on_step_begin(self.status)
  347. data['epoch_id'] = epoch_id
  348. if self.cfg.get('fp16', False):
  349. with amp.auto_cast(enable=self.cfg.use_gpu):
  350. # model forward
  351. outputs = model(data)
  352. loss = outputs['loss']
  353. # model backward
  354. scaled_loss = scaler.scale(loss)
  355. scaled_loss.backward()
  356. # in dygraph mode, optimizer.minimize is equal to optimizer.step
  357. scaler.minimize(self.optimizer, scaled_loss)
  358. else:
  359. # model forward
  360. outputs = model(data)
  361. loss = outputs['loss']
  362. # model backward
  363. loss.backward()
  364. self.optimizer.step()
  365. curr_lr = self.optimizer.get_lr()
  366. self.lr.step()
  367. if self.cfg.get('unstructured_prune'):
  368. self.pruner.step()
  369. self.optimizer.clear_grad()
  370. self.status['learning_rate'] = curr_lr
  371. if self._nranks < 2 or self._local_rank == 0:
  372. self.status['training_staus'].update(outputs)
  373. self.status['batch_time'].update(time.time() - iter_tic)
  374. self._compose_callback.on_step_end(self.status)
  375. if self.use_ema:
  376. self.ema.update(self.model)
  377. iter_tic = time.time()
  378. # apply ema weight on model
  379. if self.use_ema:
  380. weight = copy.deepcopy(self.model.state_dict())
  381. self.model.set_dict(self.ema.apply())
  382. if self.cfg.get('unstructured_prune'):
  383. self.pruner.update_params()
  384. self._compose_callback.on_epoch_end(self.status)
  385. if validate and (self._nranks < 2 or self._local_rank == 0) \
  386. and ((epoch_id + 1) % self.cfg.snapshot_epoch == 0 \
  387. or epoch_id == self.end_epoch - 1):
  388. if not hasattr(self, '_eval_loader'):
  389. # build evaluation dataset and loader
  390. self._eval_dataset = self.cfg.EvalDataset
  391. self._eval_batch_sampler = \
  392. paddle.io.BatchSampler(
  393. self._eval_dataset,
  394. batch_size=self.cfg.EvalReader['batch_size'])
  395. # If metric is VOC, need to be set collate_batch=False.
  396. if self.cfg.metric == 'VOC':
  397. self.cfg['EvalReader']['collate_batch'] = False
  398. self._eval_loader = create('EvalReader')(
  399. self._eval_dataset,
  400. self.cfg.worker_num,
  401. batch_sampler=self._eval_batch_sampler)
  402. # if validation in training is enabled, metrics should be re-init
  403. # Init_mark makes sure this code will only execute once
  404. if validate and Init_mark == False:
  405. Init_mark = True
  406. self._init_metrics(validate=validate)
  407. self._reset_metrics()
  408. with paddle.no_grad():
  409. self.status['save_best_model'] = True
  410. self._eval_with_loader(self._eval_loader)
  411. # restore origin weight on model
  412. if self.use_ema:
  413. self.model.set_dict(weight)
  414. self._compose_callback.on_train_end(self.status)
  415. def _eval_with_loader(self, loader):
  416. sample_num = 0
  417. tic = time.time()
  418. self._compose_callback.on_epoch_begin(self.status)
  419. self.status['mode'] = 'eval'
  420. self.model.eval()
  421. if self.cfg.get('print_flops', False):
  422. flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
  423. self.dataset, self.cfg.worker_num, self._eval_batch_sampler)
  424. self._flops(flops_loader)
  425. for step_id, data in enumerate(loader):
  426. self.status['step_id'] = step_id
  427. self._compose_callback.on_step_begin(self.status)
  428. # forward
  429. outs = self.model(data)
  430. # update metrics
  431. for metric in self._metrics:
  432. metric.update(data, outs)
  433. # multi-scale inputs: all inputs have same im_id
  434. if isinstance(data, typing.Sequence):
  435. sample_num += data[0]['im_id'].numpy().shape[0]
  436. else:
  437. sample_num += data['im_id'].numpy().shape[0]
  438. self._compose_callback.on_step_end(self.status)
  439. self.status['sample_num'] = sample_num
  440. self.status['cost_time'] = time.time() - tic
  441. # accumulate metric to log out
  442. for metric in self._metrics:
  443. metric.accumulate()
  444. metric.log()
  445. self._compose_callback.on_epoch_end(self.status)
  446. # reset metric states for metric may performed multiple times
  447. self._reset_metrics()
  448. def evaluate(self):
  449. with paddle.no_grad():
  450. self._eval_with_loader(self.loader)
  451. def predict(self,
  452. images,
  453. draw_threshold=0.5,
  454. output_dir='output',
  455. save_txt=False):
  456. self.dataset.set_images(images)
  457. loader = create('TestReader')(self.dataset, 0)
  458. imid2path = self.dataset.get_imid2path()
  459. anno_file = self.dataset.get_anno()
  460. clsid2catid, catid2name = get_categories(
  461. self.cfg.metric, anno_file=anno_file)
  462. # Run Infer
  463. self.status['mode'] = 'test'
  464. self.model.eval()
  465. if self.cfg.get('print_flops', False):
  466. flops_loader = create('TestReader')(self.dataset, 0)
  467. self._flops(flops_loader)
  468. results = []
  469. for step_id, data in enumerate(loader):
  470. self.status['step_id'] = step_id
  471. # forward
  472. outs = self.model(data)
  473. for key in ['im_shape', 'scale_factor', 'im_id']:
  474. if isinstance(data, typing.Sequence):
  475. outs[key] = data[0][key]
  476. else:
  477. outs[key] = data[key]
  478. for key, value in outs.items():
  479. if hasattr(value, 'numpy'):
  480. outs[key] = value.numpy()
  481. results.append(outs)
  482. # sniper
  483. if type(self.dataset) == SniperCOCODataSet:
  484. results = self.dataset.anno_cropper.aggregate_chips_detections(
  485. results)
  486. for outs in results:
  487. batch_res = get_infer_results(outs, clsid2catid)
  488. bbox_num = outs['bbox_num']
  489. start = 0
  490. for i, im_id in enumerate(outs['im_id']):
  491. image_path = imid2path[int(im_id)]
  492. image = Image.open(image_path).convert('RGB')
  493. image = ImageOps.exif_transpose(image)
  494. self.status['original_image'] = np.array(image.copy())
  495. end = start + bbox_num[i]
  496. bbox_res = batch_res['bbox'][start:end] \
  497. if 'bbox' in batch_res else None
  498. mask_res = batch_res['mask'][start:end] \
  499. if 'mask' in batch_res else None
  500. segm_res = batch_res['segm'][start:end] \
  501. if 'segm' in batch_res else None
  502. keypoint_res = batch_res['keypoint'][start:end] \
  503. if 'keypoint' in batch_res else None
  504. image = visualize_results(
  505. image, bbox_res, mask_res, segm_res, keypoint_res,
  506. int(im_id), catid2name, draw_threshold)
  507. self.status['result_image'] = np.array(image.copy())
  508. if self._compose_callback:
  509. self._compose_callback.on_step_end(self.status)
  510. # save image with detection
  511. save_name = self._get_save_image_name(output_dir, image_path)
  512. logger.info("Detection bbox results save in {}".format(
  513. save_name))
  514. image.save(save_name, quality=95)
  515. if save_txt:
  516. save_path = os.path.splitext(save_name)[0] + '.txt'
  517. results = {}
  518. results["im_id"] = im_id
  519. if bbox_res:
  520. results["bbox_res"] = bbox_res
  521. if keypoint_res:
  522. results["keypoint_res"] = keypoint_res
  523. save_result(save_path, results, catid2name, draw_threshold)
  524. start = end
  525. def _get_save_image_name(self, output_dir, image_path):
  526. """
  527. Get save image name from source image path.
  528. """
  529. if not os.path.exists(output_dir):
  530. os.makedirs(output_dir)
  531. image_name = os.path.split(image_path)[-1]
  532. name, ext = os.path.splitext(image_name)
  533. return os.path.join(output_dir, "{}".format(name)) + ext
  534. def _get_infer_cfg_and_input_spec(self, save_dir, prune_input=True):
  535. image_shape = None
  536. im_shape = [None, 2]
  537. scale_factor = [None, 2]
  538. if self.cfg.architecture in MOT_ARCH:
  539. test_reader_name = 'TestMOTReader'
  540. else:
  541. test_reader_name = 'TestReader'
  542. if 'inputs_def' in self.cfg[test_reader_name]:
  543. inputs_def = self.cfg[test_reader_name]['inputs_def']
  544. image_shape = inputs_def.get('image_shape', None)
  545. # set image_shape=[None, 3, -1, -1] as default
  546. if image_shape is None:
  547. image_shape = [None, 3, -1, -1]
  548. if len(image_shape) == 3:
  549. image_shape = [None] + image_shape
  550. else:
  551. im_shape = [image_shape[0], 2]
  552. scale_factor = [image_shape[0], 2]
  553. if hasattr(self.model, 'deploy'):
  554. self.model.deploy = True
  555. if hasattr(self.model, 'fuse_norm'):
  556. self.model.fuse_norm = self.cfg['TestReader'].get('fuse_normalize',
  557. False)
  558. # Save infer cfg
  559. _dump_infer_config(self.cfg,
  560. os.path.join(save_dir, 'infer_cfg.yml'),
  561. image_shape, self.model)
  562. input_spec = [{
  563. "image": InputSpec(
  564. shape=image_shape, name='image'),
  565. "im_shape": InputSpec(
  566. shape=im_shape, name='im_shape'),
  567. "scale_factor": InputSpec(
  568. shape=scale_factor, name='scale_factor')
  569. }]
  570. if self.cfg.architecture == 'DeepSORT':
  571. input_spec[0].update({
  572. "crops": InputSpec(
  573. shape=[None, 3, 192, 64], name='crops')
  574. })
  575. if prune_input:
  576. static_model = paddle.jit.to_static(
  577. self.model, input_spec=input_spec)
  578. # NOTE: dy2st do not pruned program, but jit.save will prune program
  579. # input spec, prune input spec here and save with pruned input spec
  580. pruned_input_spec = _prune_input_spec(
  581. input_spec, static_model.forward.main_program,
  582. static_model.forward.outputs)
  583. else:
  584. static_model = None
  585. pruned_input_spec = input_spec
  586. # TODO: Hard code, delete it when support prune input_spec.
  587. if self.cfg.architecture == 'PicoDet':
  588. pruned_input_spec = [{
  589. "image": InputSpec(
  590. shape=image_shape, name='image')
  591. }]
  592. return static_model, pruned_input_spec
  593. def export(self, output_dir='output_inference'):
  594. self.model.eval()
  595. model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
  596. save_dir = os.path.join(output_dir, model_name)
  597. if not os.path.exists(save_dir):
  598. os.makedirs(save_dir)
  599. static_model, pruned_input_spec = self._get_infer_cfg_and_input_spec(
  600. save_dir)
  601. # dy2st and save model
  602. if 'slim' not in self.cfg or self.cfg['slim_type'] != 'QAT':
  603. paddle.jit.save(
  604. static_model,
  605. os.path.join(save_dir, 'model'),
  606. input_spec=pruned_input_spec)
  607. else:
  608. self.cfg.slim.save_quantized_model(
  609. self.model,
  610. os.path.join(save_dir, 'model'),
  611. input_spec=pruned_input_spec)
  612. logger.info("Export model and saved in {}".format(save_dir))
  613. def post_quant(self, output_dir='output_inference'):
  614. model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
  615. save_dir = os.path.join(output_dir, model_name)
  616. if not os.path.exists(save_dir):
  617. os.makedirs(save_dir)
  618. for idx, data in enumerate(self.loader):
  619. self.model(data)
  620. if idx == int(self.cfg.get('quant_batch_num', 10)):
  621. break
  622. # TODO: support prune input_spec
  623. _, pruned_input_spec = self._get_infer_cfg_and_input_spec(
  624. save_dir, prune_input=False)
  625. self.cfg.slim.save_quantized_model(
  626. self.model,
  627. os.path.join(save_dir, 'model'),
  628. input_spec=pruned_input_spec)
  629. logger.info("Export Post-Quant model and saved in {}".format(save_dir))
  630. def _flops(self, loader):
  631. self.model.eval()
  632. try:
  633. import paddleslim
  634. except Exception as e:
  635. logger.warning(
  636. 'Unable to calculate flops, please install paddleslim, for example: `pip install paddleslim`'
  637. )
  638. return
  639. from paddleslim.analysis import dygraph_flops as flops
  640. input_data = None
  641. for data in loader:
  642. input_data = data
  643. break
  644. input_spec = [{
  645. "image": input_data['image'][0].unsqueeze(0),
  646. "im_shape": input_data['im_shape'][0].unsqueeze(0),
  647. "scale_factor": input_data['scale_factor'][0].unsqueeze(0)
  648. }]
  649. flops = flops(self.model, input_spec) / (1000**3)
  650. logger.info(" Model FLOPs : {:.6f}G. (image shape is {})".format(
  651. flops, input_data['image'][0].unsqueeze(0).shape))