run_task.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. #!/usr/bin/env python
  2. import os
  3. import paddle
  4. import paddlers
  5. from paddlers import transforms as T
  6. import custom_model
  7. import custom_trainer
  8. from config_utils import parse_args, build_objects, CfgNode
  9. def format_cfg(cfg, indent=0):
  10. s = ''
  11. if isinstance(cfg, dict):
  12. for i, (k, v) in enumerate(sorted(cfg.items())):
  13. s += ' ' * indent + str(k) + ': '
  14. if isinstance(v, (dict, list, CfgNode)):
  15. s += '\n' + format_cfg(v, indent=indent + 1)
  16. else:
  17. s += str(v)
  18. if i != len(cfg) - 1:
  19. s += '\n'
  20. elif isinstance(cfg, list):
  21. for i, v in enumerate(cfg):
  22. s += ' ' * indent + '- '
  23. if isinstance(v, (dict, list, CfgNode)):
  24. s += '\n' + format_cfg(v, indent=indent + 1)
  25. else:
  26. s += str(v)
  27. if i != len(cfg) - 1:
  28. s += '\n'
  29. elif isinstance(cfg, CfgNode):
  30. s += ' ' * indent + f"type: {cfg.type}" + '\n'
  31. s += ' ' * indent + f"module: {cfg.module}" + '\n'
  32. s += ' ' * indent + 'args: \n' + format_cfg(cfg.args, indent + 1)
  33. return s
  34. if __name__ == '__main__':
  35. CfgNode.set_context(globals())
  36. cfg = parse_args()
  37. print(format_cfg(cfg))
  38. # Automatically download data
  39. if cfg['download_on']:
  40. paddlers.utils.download_and_decompress(
  41. cfg['download_url'], path=cfg['download_path'])
  42. if cfg['cmd'] == 'train':
  43. if not isinstance(cfg['datasets']['train'].args, dict):
  44. raise ValueError("args of train dataset must be a dict!")
  45. if cfg['datasets']['train'].args.get('transforms', None) is not None:
  46. raise ValueError(
  47. "Found key 'transforms' in args of train dataset and the value is not None."
  48. )
  49. train_transforms = T.Compose(
  50. build_objects(
  51. cfg['transforms']['train'], mod=T))
  52. # Inplace modification
  53. cfg['datasets']['train'].args['transforms'] = train_transforms
  54. train_dataset = build_objects(
  55. cfg['datasets']['train'], mod=paddlers.datasets)
  56. if not isinstance(cfg['datasets']['eval'].args, dict):
  57. raise ValueError("args of eval dataset must be a dict!")
  58. if cfg['datasets']['eval'].args.get('transforms', None) is not None:
  59. raise ValueError(
  60. "Found key 'transforms' in args of eval dataset and the value is not None."
  61. )
  62. eval_transforms = T.Compose(build_objects(cfg['transforms']['eval'], mod=T))
  63. # Inplace modification
  64. cfg['datasets']['eval'].args['transforms'] = eval_transforms
  65. eval_dataset = build_objects(cfg['datasets']['eval'], mod=paddlers.datasets)
  66. model = build_objects(
  67. cfg['model'], mod=getattr(paddlers.tasks, cfg['task']))
  68. if cfg['cmd'] == 'train':
  69. if cfg['optimizer']:
  70. if len(cfg['optimizer'].args) == 0:
  71. cfg['optimizer'].args = {}
  72. if not isinstance(cfg['optimizer'].args, dict):
  73. raise TypeError("args of optimizer must be a dict!")
  74. if cfg['optimizer'].args.get('parameters', None) is not None:
  75. raise ValueError(
  76. "Found key 'parameters' in args of optimizer and the value is not None."
  77. )
  78. cfg['optimizer'].args['parameters'] = model.net.parameters()
  79. optimizer = build_objects(cfg['optimizer'], mod=paddle.optimizer)
  80. else:
  81. optimizer = None
  82. model.train(
  83. num_epochs=cfg['num_epochs'],
  84. train_dataset=train_dataset,
  85. train_batch_size=cfg['train_batch_size'],
  86. eval_dataset=eval_dataset,
  87. optimizer=optimizer,
  88. save_interval_epochs=cfg['save_interval_epochs'],
  89. log_interval_steps=cfg['log_interval_steps'],
  90. save_dir=cfg['save_dir'],
  91. learning_rate=cfg['learning_rate'],
  92. early_stop=cfg['early_stop'],
  93. early_stop_patience=cfg['early_stop_patience'],
  94. use_vdl=cfg['use_vdl'],
  95. resume_checkpoint=cfg['resume_checkpoint'] or None,
  96. **cfg['train'])
  97. elif cfg['cmd'] == 'eval':
  98. state_dict = paddle.load(
  99. os.path.join(cfg['resume_checkpoint'], 'model.pdparams'))
  100. model.net.set_state_dict(state_dict)
  101. res = model.evaluate(eval_dataset)
  102. print(res)