trainer.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535
  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 time
  19. import random
  20. import datetime
  21. import numpy as np
  22. from PIL import Image
  23. import paddle
  24. import paddle.distributed as dist
  25. from paddle.distributed import fleet
  26. from paddle import amp
  27. from paddle.static import InputSpec
  28. from paddlex.ppdet.optimizer import ModelEMA
  29. from paddlex.ppdet.core.workspace import create
  30. from paddlex.ppdet.utils.checkpoint import load_weight, load_pretrain_weight
  31. from paddlex.ppdet.utils.visualizer import visualize_results, save_result
  32. from paddlex.ppdet.metrics import JDEDetMetric, JDEReIDMetric
  33. from paddlex.ppdet.metrics import Metric, COCOMetric, VOCMetric, WiderFaceMetric, get_infer_results, KeyPointTopDownCOCOEval
  34. from paddlex.ppdet.data.source.category import get_categories
  35. from paddlex.ppdet.utils import stats
  36. from .callbacks import Callback, ComposeCallback, LogPrinter, Checkpointer, WiferFaceEval, VisualDLWriter
  37. from .export_utils import _dump_infer_config
  38. from paddlex.ppdet.utils.logger import setup_logger
  39. logger = setup_logger('ppdet.engine')
  40. __all__ = ['Trainer']
  41. class Trainer(object):
  42. def __init__(self, cfg, mode='train'):
  43. self.cfg = cfg
  44. assert mode.lower() in ['train', 'eval', 'test'], \
  45. "mode should be 'train', 'eval' or 'test'"
  46. self.mode = mode.lower()
  47. self.optimizer = None
  48. self.is_loaded_weights = False
  49. # build data loader
  50. self.dataset = cfg['{}Dataset'.format(self.mode.capitalize())]
  51. if self.mode == 'train':
  52. self.loader = create('{}Reader'.format(self.mode.capitalize()))(
  53. self.dataset, cfg.worker_num)
  54. if cfg.architecture == 'JDE' and self.mode == 'train':
  55. cfg['JDEEmbeddingHead'][
  56. 'num_identifiers'] = self.dataset.total_identities
  57. if cfg.architecture == 'FairMOT' and self.mode == 'train':
  58. cfg['FairMOTEmbeddingHead'][
  59. 'num_identifiers'] = self.dataset.total_identities
  60. # build model
  61. if 'model' not in self.cfg:
  62. self.model = create(cfg.architecture)
  63. else:
  64. self.model = self.cfg.model
  65. self.is_loaded_weights = True
  66. self.use_ema = ('use_ema' in cfg and cfg['use_ema'])
  67. if self.use_ema:
  68. self.ema = ModelEMA(
  69. cfg['ema_decay'], self.model, use_thres_step=True)
  70. # EvalDataset build with BatchSampler to evaluate in single device
  71. # TODO: multi-device evaluate
  72. if self.mode == 'eval':
  73. self._eval_batch_sampler = paddle.io.BatchSampler(
  74. self.dataset, batch_size=self.cfg.EvalReader['batch_size'])
  75. self.loader = create('{}Reader'.format(self.mode.capitalize()))(
  76. self.dataset, cfg.worker_num, self._eval_batch_sampler)
  77. # TestDataset build after user set images, skip loader creation here
  78. # build optimizer in train mode
  79. if self.mode == 'train':
  80. steps_per_epoch = len(self.loader)
  81. self.lr = create('LearningRate')(steps_per_epoch)
  82. self.optimizer = create('OptimizerBuilder')(
  83. self.lr, self.model.parameters())
  84. self._nranks = dist.get_world_size()
  85. self._local_rank = dist.get_rank()
  86. self.status = {}
  87. self.start_epoch = 0
  88. self.end_epoch = cfg.epoch
  89. # initial default callbacks
  90. self._init_callbacks()
  91. # initial default metrics
  92. self._init_metrics()
  93. self._reset_metrics()
  94. def _init_callbacks(self):
  95. if self.mode == 'train':
  96. self._callbacks = [LogPrinter(self), Checkpointer(self)]
  97. if self.cfg.get('use_vdl', False):
  98. self._callbacks.append(VisualDLWriter(self))
  99. self._compose_callback = ComposeCallback(self._callbacks)
  100. elif self.mode == 'eval':
  101. self._callbacks = [LogPrinter(self)]
  102. if self.cfg.metric == 'WiderFace':
  103. self._callbacks.append(WiferFaceEval(self))
  104. self._compose_callback = ComposeCallback(self._callbacks)
  105. elif self.mode == 'test' and self.cfg.get('use_vdl', False):
  106. self._callbacks = [VisualDLWriter(self)]
  107. self._compose_callback = ComposeCallback(self._callbacks)
  108. else:
  109. self._callbacks = []
  110. self._compose_callback = None
  111. def _init_metrics(self, validate=False):
  112. if self.mode == 'test' or (self.mode == 'train' and not validate):
  113. self._metrics = []
  114. return
  115. classwise = self.cfg['classwise'] if 'classwise' in self.cfg else False
  116. if self.cfg.metric == 'COCO':
  117. # TODO: bias should be unified
  118. bias = self.cfg['bias'] if 'bias' in self.cfg else 0
  119. output_eval = self.cfg['output_eval'] \
  120. if 'output_eval' in self.cfg else None
  121. save_prediction_only = self.cfg.get('save_prediction_only', False)
  122. # pass clsid2catid info to metric instance to avoid multiple loading
  123. # annotation file
  124. clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()} \
  125. if self.mode == 'eval' else None
  126. # when do validation in train, annotation file should be get from
  127. # EvalReader instead of self.dataset(which is TrainReader)
  128. anno_file = self.dataset.get_anno()
  129. if self.mode == 'train' and validate:
  130. eval_dataset = self.cfg['EvalDataset']
  131. eval_dataset.check_or_download_dataset()
  132. anno_file = eval_dataset.get_anno()
  133. IouType = self.cfg['IouType'] if 'IouType' in self.cfg else 'bbox'
  134. self._metrics = [
  135. COCOMetric(
  136. anno_file=anno_file,
  137. clsid2catid=clsid2catid,
  138. classwise=classwise,
  139. output_eval=output_eval,
  140. bias=bias,
  141. IouType=IouType,
  142. save_prediction_only=save_prediction_only)
  143. ]
  144. elif self.cfg.metric == 'VOC':
  145. self._metrics = [
  146. VOCMetric(
  147. label_list=self.dataset.get_label_list(),
  148. class_num=self.cfg.num_classes,
  149. map_type=self.cfg.map_type,
  150. classwise=classwise)
  151. ]
  152. elif self.cfg.metric == 'WiderFace':
  153. multi_scale = self.cfg.multi_scale_eval if 'multi_scale_eval' in self.cfg else True
  154. self._metrics = [
  155. WiderFaceMetric(
  156. image_dir=os.path.join(self.dataset.dataset_dir,
  157. self.dataset.image_dir),
  158. anno_file=self.dataset.get_anno(),
  159. multi_scale=multi_scale)
  160. ]
  161. elif self.cfg.metric == 'KeyPointTopDownCOCOEval':
  162. eval_dataset = self.cfg['EvalDataset']
  163. eval_dataset.check_or_download_dataset()
  164. anno_file = eval_dataset.get_anno()
  165. self._metrics = [
  166. KeyPointTopDownCOCOEval(anno_file,
  167. len(eval_dataset), self.cfg.num_joints,
  168. self.cfg.save_dir)
  169. ]
  170. elif self.cfg.metric == 'MOTDet':
  171. self._metrics = [JDEDetMetric(), ]
  172. elif self.cfg.metric == 'ReID':
  173. self._metrics = [JDEReIDMetric(), ]
  174. else:
  175. logger.warn("Metric not support for metric type {}".format(
  176. self.cfg.metric))
  177. self._metrics = []
  178. def _reset_metrics(self):
  179. for metric in self._metrics:
  180. metric.reset()
  181. def register_callbacks(self, callbacks):
  182. callbacks = [c for c in list(callbacks) if c is not None]
  183. for c in callbacks:
  184. assert isinstance(c, Callback), \
  185. "metrics shoule be instances of subclass of Metric"
  186. self._callbacks.extend(callbacks)
  187. self._compose_callback = ComposeCallback(self._callbacks)
  188. def register_metrics(self, metrics):
  189. metrics = [m for m in list(metrics) if m is not None]
  190. for m in metrics:
  191. assert isinstance(m, Metric), \
  192. "metrics shoule be instances of subclass of Metric"
  193. self._metrics.extend(metrics)
  194. def load_weights(self, weights):
  195. if self.is_loaded_weights:
  196. return
  197. self.start_epoch = 0
  198. if hasattr(self.model, 'detector'):
  199. if self.model.__class__.__name__ == 'FairMOT':
  200. load_pretrain_weight(self.model, weights)
  201. else:
  202. load_pretrain_weight(self.model.detector, weights)
  203. else:
  204. load_pretrain_weight(self.model, weights)
  205. logger.debug("Load weights {} to start training".format(weights))
  206. def resume_weights(self, weights):
  207. # support Distill resume weights
  208. if hasattr(self.model, 'student_model'):
  209. self.start_epoch = load_weight(self.model.student_model, weights,
  210. self.optimizer)
  211. else:
  212. self.start_epoch = load_weight(self.model, weights, self.optimizer)
  213. logger.debug("Resume weights of epoch {}".format(self.start_epoch))
  214. def train(self, validate=False):
  215. assert self.mode == 'train', "Model not in 'train' mode"
  216. # if validation in training is enabled, metrics should be re-init
  217. if validate:
  218. self._init_metrics(validate=validate)
  219. self._reset_metrics()
  220. model = self.model
  221. if self.cfg.get('fleet', False):
  222. model = fleet.distributed_model(model)
  223. self.optimizer = fleet.distributed_optimizer(
  224. self.optimizer).user_defined_optimizer
  225. elif self._nranks > 1:
  226. find_unused_parameters = self.cfg[
  227. 'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
  228. model = paddle.DataParallel(
  229. self.model, find_unused_parameters=find_unused_parameters)
  230. # initial fp16
  231. if self.cfg.get('fp16', False):
  232. scaler = amp.GradScaler(
  233. enable=self.cfg.use_gpu, init_loss_scaling=1024)
  234. self.status.update({
  235. 'epoch_id': self.start_epoch,
  236. 'step_id': 0,
  237. 'steps_per_epoch': len(self.loader)
  238. })
  239. self.status['batch_time'] = stats.SmoothedValue(
  240. self.cfg.log_iter, fmt='{avg:.4f}')
  241. self.status['data_time'] = stats.SmoothedValue(
  242. self.cfg.log_iter, fmt='{avg:.4f}')
  243. self.status['training_staus'] = stats.TrainingStats(self.cfg.log_iter)
  244. for epoch_id in range(self.start_epoch, self.cfg.epoch):
  245. self.status['mode'] = 'train'
  246. self.status['epoch_id'] = epoch_id
  247. self._compose_callback.on_epoch_begin(self.status)
  248. self.loader.dataset.set_epoch(epoch_id)
  249. model.train()
  250. iter_tic = time.time()
  251. for step_id, data in enumerate(self.loader):
  252. self.status['data_time'].update(time.time() - iter_tic)
  253. self.status['step_id'] = step_id
  254. self._compose_callback.on_step_begin(self.status)
  255. if self.cfg.get('fp16', False):
  256. with amp.auto_cast(enable=self.cfg.use_gpu):
  257. # model forward
  258. outputs = model(data)
  259. loss = outputs['loss']
  260. # model backward
  261. scaled_loss = scaler.scale(loss)
  262. scaled_loss.backward()
  263. # in dygraph mode, optimizer.minimize is equal to optimizer.step
  264. scaler.minimize(self.optimizer, scaled_loss)
  265. else:
  266. # model forward
  267. outputs = model(data)
  268. loss = outputs['loss']
  269. # model backward
  270. loss.backward()
  271. self.optimizer.step()
  272. curr_lr = self.optimizer.get_lr()
  273. self.lr.step()
  274. self.optimizer.clear_grad()
  275. self.status['learning_rate'] = curr_lr
  276. if self._nranks < 2 or self._local_rank == 0:
  277. self.status['training_staus'].update(outputs)
  278. self.status['batch_time'].update(time.time() - iter_tic)
  279. self._compose_callback.on_step_end(self.status)
  280. if self.use_ema:
  281. self.ema.update(self.model)
  282. iter_tic = time.time()
  283. # apply ema weight on model
  284. if self.use_ema:
  285. weight = self.model.state_dict()
  286. self.model.set_dict(self.ema.apply())
  287. self._compose_callback.on_epoch_end(self.status)
  288. if validate and (self._nranks < 2 or self._local_rank == 0) \
  289. and ((epoch_id + 1) % self.cfg.snapshot_epoch == 0 \
  290. or epoch_id == self.end_epoch - 1):
  291. if not hasattr(self, '_eval_loader'):
  292. # build evaluation dataset and loader
  293. self._eval_dataset = self.cfg.EvalDataset
  294. self._eval_batch_sampler = \
  295. paddle.io.BatchSampler(
  296. self._eval_dataset,
  297. batch_size=self.cfg.EvalReader['batch_size'])
  298. self._eval_loader = create('EvalReader')(
  299. self._eval_dataset,
  300. self.cfg.worker_num,
  301. batch_sampler=self._eval_batch_sampler)
  302. with paddle.no_grad():
  303. self.status['save_best_model'] = True
  304. self._eval_with_loader(self._eval_loader)
  305. # restore origin weight on model
  306. if self.use_ema:
  307. self.model.set_dict(weight)
  308. def _eval_with_loader(self, loader):
  309. sample_num = 0
  310. tic = time.time()
  311. self._compose_callback.on_epoch_begin(self.status)
  312. self.status['mode'] = 'eval'
  313. self.model.eval()
  314. for step_id, data in enumerate(loader):
  315. self.status['step_id'] = step_id
  316. self._compose_callback.on_step_begin(self.status)
  317. # forward
  318. outs = self.model(data)
  319. # update metrics
  320. for metric in self._metrics:
  321. metric.update(data, outs)
  322. sample_num += data['im_id'].numpy().shape[0]
  323. self._compose_callback.on_step_end(self.status)
  324. self.status['sample_num'] = sample_num
  325. self.status['cost_time'] = time.time() - tic
  326. # accumulate metric to log out
  327. for metric in self._metrics:
  328. metric.accumulate()
  329. metric.log()
  330. self._compose_callback.on_epoch_end(self.status)
  331. # reset metric states for metric may performed multiple times
  332. self._reset_metrics()
  333. def evaluate(self):
  334. with paddle.no_grad():
  335. self._eval_with_loader(self.loader)
  336. def predict(self,
  337. images,
  338. draw_threshold=0.5,
  339. output_dir='output',
  340. save_txt=False):
  341. self.dataset.set_images(images)
  342. loader = create('TestReader')(self.dataset, 0)
  343. imid2path = self.dataset.get_imid2path()
  344. anno_file = self.dataset.get_anno()
  345. clsid2catid, catid2name = get_categories(
  346. self.cfg.metric, anno_file=anno_file)
  347. # Run Infer
  348. self.status['mode'] = 'test'
  349. self.model.eval()
  350. for step_id, data in enumerate(loader):
  351. self.status['step_id'] = step_id
  352. # forward
  353. outs = self.model(data)
  354. for key in ['im_shape', 'scale_factor', 'im_id']:
  355. outs[key] = data[key]
  356. for key, value in outs.items():
  357. if hasattr(value, 'numpy'):
  358. outs[key] = value.numpy()
  359. batch_res = get_infer_results(outs, clsid2catid)
  360. bbox_num = outs['bbox_num']
  361. start = 0
  362. for i, im_id in enumerate(outs['im_id']):
  363. image_path = imid2path[int(im_id)]
  364. image = Image.open(image_path).convert('RGB')
  365. self.status['original_image'] = np.array(image.copy())
  366. end = start + bbox_num[i]
  367. bbox_res = batch_res['bbox'][start:end] \
  368. if 'bbox' in batch_res else None
  369. mask_res = batch_res['mask'][start:end] \
  370. if 'mask' in batch_res else None
  371. segm_res = batch_res['segm'][start:end] \
  372. if 'segm' in batch_res else None
  373. keypoint_res = batch_res['keypoint'][start:end] \
  374. if 'keypoint' in batch_res else None
  375. image = visualize_results(
  376. image, bbox_res, mask_res, segm_res, keypoint_res,
  377. int(im_id), catid2name, draw_threshold)
  378. self.status['result_image'] = np.array(image.copy())
  379. if self._compose_callback:
  380. self._compose_callback.on_step_end(self.status)
  381. # save image with detection
  382. save_name = self._get_save_image_name(output_dir, image_path)
  383. logger.info("Detection bbox results save in {}".format(
  384. save_name))
  385. image.save(save_name, quality=95)
  386. if save_txt:
  387. save_path = os.path.splitext(save_name)[0] + '.txt'
  388. results = {}
  389. results["im_id"] = im_id
  390. if bbox_res:
  391. results["bbox_res"] = bbox_res
  392. if keypoint_res:
  393. results["keypoint_res"] = keypoint_res
  394. save_result(save_path, results, catid2name, draw_threshold)
  395. start = end
  396. def _get_save_image_name(self, output_dir, image_path):
  397. """
  398. Get save image name from source image path.
  399. """
  400. if not os.path.exists(output_dir):
  401. os.makedirs(output_dir)
  402. image_name = os.path.split(image_path)[-1]
  403. name, ext = os.path.splitext(image_name)
  404. return os.path.join(output_dir, "{}".format(name)) + ext
  405. def export(self, output_dir='output_inference'):
  406. self.model.eval()
  407. model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
  408. save_dir = os.path.join(output_dir, model_name)
  409. if not os.path.exists(save_dir):
  410. os.makedirs(save_dir)
  411. image_shape = None
  412. if 'inputs_def' in self.cfg['TestReader']:
  413. inputs_def = self.cfg['TestReader']['inputs_def']
  414. image_shape = inputs_def.get('image_shape', None)
  415. # set image_shape=[3, -1, -1] as default
  416. if image_shape is None:
  417. image_shape = [3, -1, -1]
  418. self.model.eval()
  419. if hasattr(self.model, 'deploy'): self.model.deploy = True
  420. # Save infer cfg
  421. _dump_infer_config(self.cfg,
  422. os.path.join(save_dir, 'infer_cfg.yml'),
  423. image_shape, self.model)
  424. input_spec = [{
  425. "image": InputSpec(
  426. shape=[None] + image_shape, name='image'),
  427. "im_shape": InputSpec(
  428. shape=[None, 2], name='im_shape'),
  429. "scale_factor": InputSpec(
  430. shape=[None, 2], name='scale_factor')
  431. }]
  432. # dy2st and save model
  433. if 'slim' not in self.cfg or self.cfg['slim_type'] != 'QAT':
  434. static_model = paddle.jit.to_static(
  435. self.model, input_spec=input_spec)
  436. # NOTE: dy2st do not pruned program, but jit.save will prune program
  437. # input spec, prune input spec here and save with pruned input spec
  438. pruned_input_spec = self._prune_input_spec(
  439. input_spec, static_model.forward.main_program,
  440. static_model.forward.outputs)
  441. paddle.jit.save(
  442. static_model,
  443. os.path.join(save_dir, 'model'),
  444. input_spec=pruned_input_spec)
  445. logger.info("Export model and saved in {}".format(save_dir))
  446. else:
  447. self.cfg.slim.save_quantized_model(
  448. self.model,
  449. os.path.join(save_dir, 'model'),
  450. input_spec=input_spec)
  451. def _prune_input_spec(self, input_spec, program, targets):
  452. # try to prune static program to figure out pruned input spec
  453. # so we perform following operations in static mode
  454. paddle.enable_static()
  455. pruned_input_spec = [{}]
  456. program = program.clone()
  457. program = program._prune(targets=targets)
  458. global_block = program.global_block()
  459. for name, spec in input_spec[0].items():
  460. try:
  461. v = global_block.var(name)
  462. pruned_input_spec[0][name] = spec
  463. except Exception:
  464. pass
  465. paddle.disable_static()
  466. return pruned_input_spec