train.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  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 argparse
  18. import os
  19. import sys
  20. __dir__ = os.path.dirname(os.path.abspath(__file__))
  21. sys.path.append(__dir__)
  22. sys.path.append(os.path.abspath(os.path.join(__dir__, '../../')))
  23. import paddle
  24. from paddle.distributed import fleet
  25. from visualdl import LogWriter
  26. from paddlex.ppcls.data import build_dataloader
  27. from paddlex.ppcls.utils.config import get_config, print_config
  28. from paddlex.ppcls.utils import logger
  29. from paddlex.ppcls.utils.logger import init_logger
  30. from paddlex.ppcls.static.save_load import init_model, save_model
  31. from paddlex.ppcls.static import program
  32. def parse_args():
  33. parser = argparse.ArgumentParser("PaddleClas train script")
  34. parser.add_argument(
  35. '-c',
  36. '--config',
  37. type=str,
  38. default='configs/ResNet/ResNet50.yaml',
  39. help='config file path')
  40. parser.add_argument(
  41. '-p',
  42. '--profiler_options',
  43. type=str,
  44. default=None,
  45. help='The option of profiler, which should be in format \"key1=value1;key2=value2;key3=value3\".'
  46. )
  47. parser.add_argument(
  48. '-o',
  49. '--override',
  50. action='append',
  51. default=[],
  52. help='config options to be overridden')
  53. args = parser.parse_args()
  54. return args
  55. def main(args):
  56. """
  57. all the config of training paradigm should be in config["Global"]
  58. """
  59. config = get_config(args.config, overrides=args.override, show=False)
  60. global_config = config["Global"]
  61. mode = "train"
  62. log_file = os.path.join(global_config['output_dir'],
  63. config["Arch"]["name"], f"{mode}.log")
  64. init_logger(name='root', log_file=log_file)
  65. print_config(config)
  66. if global_config.get("is_distributed", True):
  67. fleet.init(is_collective=True)
  68. # assign the device
  69. use_gpu = global_config.get("use_gpu", True)
  70. # amp related config
  71. if 'AMP' in config:
  72. AMP_RELATED_FLAGS_SETTING = {
  73. 'FLAGS_cudnn_exhaustive_search': "1",
  74. 'FLAGS_conv_workspace_size_limit': "1500",
  75. 'FLAGS_cudnn_batchnorm_spatial_persistent': "1",
  76. 'FLAGS_max_indevice_grad_add': "8",
  77. "FLAGS_cudnn_batchnorm_spatial_persistent": "1",
  78. }
  79. for k in AMP_RELATED_FLAGS_SETTING:
  80. os.environ[k] = AMP_RELATED_FLAGS_SETTING[k]
  81. use_xpu = global_config.get("use_xpu", False)
  82. use_npu = global_config.get("use_npu", False)
  83. assert (
  84. use_gpu and use_xpu and use_npu
  85. ) is not True, "gpu, xpu and npu can not be true in the same time in static mode!"
  86. if use_gpu:
  87. device = paddle.set_device('gpu')
  88. elif use_xpu:
  89. device = paddle.set_device('xpu')
  90. elif use_npu:
  91. device = paddle.set_device('npu')
  92. else:
  93. device = paddle.set_device('cpu')
  94. # visualDL
  95. vdl_writer = None
  96. if global_config["use_visualdl"]:
  97. vdl_dir = os.path.join(global_config["output_dir"], "vdl")
  98. vdl_writer = LogWriter(vdl_dir)
  99. # build dataloader
  100. eval_dataloader = None
  101. use_dali = global_config.get('use_dali', False)
  102. class_num = config["Arch"].get("class_num", None)
  103. config["DataLoader"].update({"class_num": class_num})
  104. train_dataloader = build_dataloader(
  105. config["DataLoader"], "Train", device=device, use_dali=use_dali)
  106. if global_config["eval_during_train"]:
  107. eval_dataloader = build_dataloader(
  108. config["DataLoader"], "Eval", device=device, use_dali=use_dali)
  109. step_each_epoch = len(train_dataloader)
  110. # startup_prog is used to do some parameter init work,
  111. # and train prog is used to hold the network
  112. startup_prog = paddle.static.Program()
  113. train_prog = paddle.static.Program()
  114. best_top1_acc = 0.0 # best top1 acc record
  115. train_fetchs, lr_scheduler, train_feeds, optimizer = program.build(
  116. config,
  117. train_prog,
  118. startup_prog,
  119. class_num,
  120. step_each_epoch=step_each_epoch,
  121. is_train=True,
  122. is_distributed=global_config.get("is_distributed", True))
  123. if global_config["eval_during_train"]:
  124. eval_prog = paddle.static.Program()
  125. eval_fetchs, _, eval_feeds, _ = program.build(
  126. config,
  127. eval_prog,
  128. startup_prog,
  129. is_train=False,
  130. is_distributed=global_config.get("is_distributed", True))
  131. # clone to prune some content which is irrelevant in eval_prog
  132. eval_prog = eval_prog.clone(for_test=True)
  133. # create the "Executor" with the statement of which device
  134. exe = paddle.static.Executor(device)
  135. # Parameter initialization
  136. exe.run(startup_prog)
  137. # load pretrained models or checkpoints
  138. init_model(global_config, train_prog, exe)
  139. if 'AMP' in config and config.AMP.get("use_pure_fp16", False):
  140. optimizer.amp_init(
  141. device,
  142. scope=paddle.static.global_scope(),
  143. test_program=eval_prog
  144. if global_config["eval_during_train"] else None)
  145. if not global_config.get("is_distributed", True):
  146. compiled_train_prog = program.compile(
  147. config, train_prog, loss_name=train_fetchs["loss"][0].name)
  148. else:
  149. compiled_train_prog = train_prog
  150. if eval_dataloader is not None:
  151. compiled_eval_prog = program.compile(config, eval_prog)
  152. for epoch_id in range(global_config["epochs"]):
  153. # 1. train with train dataset
  154. program.run(train_dataloader, exe, compiled_train_prog, train_feeds,
  155. train_fetchs, epoch_id, 'train', config, vdl_writer,
  156. lr_scheduler, args.profiler_options)
  157. # 2. evaate with eval dataset
  158. if global_config["eval_during_train"] and epoch_id % global_config[
  159. "eval_interval"] == 0:
  160. top1_acc = program.run(eval_dataloader, exe, compiled_eval_prog,
  161. eval_feeds, eval_fetchs, epoch_id, "eval",
  162. config)
  163. if top1_acc > best_top1_acc:
  164. best_top1_acc = top1_acc
  165. message = "The best top1 acc {:.5f}, in epoch: {:d}".format(
  166. best_top1_acc, epoch_id)
  167. logger.info(message)
  168. if epoch_id % global_config["save_interval"] == 0:
  169. model_path = os.path.join(global_config["output_dir"],
  170. config["Arch"]["name"])
  171. save_model(train_prog, model_path, "best_model")
  172. # 3. save the persistable model
  173. if epoch_id % global_config["save_interval"] == 0:
  174. model_path = os.path.join(global_config["output_dir"],
  175. config["Arch"]["name"])
  176. save_model(train_prog, model_path, epoch_id)
  177. if __name__ == '__main__':
  178. paddle.enable_static()
  179. args = parse_args()
  180. main(args)