You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

160 lines
7.7 KiB

  1. # Copyright 2021 Huawei Technologies Co., Ltd
  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. # ============================================================================
  15. """train script"""
  16. import os
  17. import time
  18. import argparse
  19. import ast
  20. import mindspore.common.dtype as mstype
  21. from mindspore.context import ParallelMode
  22. from mindspore import context
  23. from mindspore import log as logger
  24. from mindspore.communication.management import init
  25. from mindspore.train.callback import Callback, CheckpointConfig, ModelCheckpoint, TimeMonitor
  26. from mindspore.train import Model
  27. from mindspore.common import set_seed
  28. from mindspore.train.loss_scale_manager import DynamicLossScaleManager
  29. from mindspore.nn.optim import Adam
  30. from src.config import config
  31. from src.seq2seq import Seq2Seq
  32. from src.gru_for_train import GRUWithLossCell, GRUTrainOneStepWithLossScaleCell, GRUTrainOneStepCell
  33. from src.dataset import create_gru_dataset
  34. from src.lr_schedule import dynamic_lr
  35. set_seed(1)
  36. parser = argparse.ArgumentParser(description="GRU training")
  37. parser.add_argument("--device_target", type=str, default="Ascend",
  38. help="device where the code will be implemented, default is Ascend")
  39. parser.add_argument("--run_distribute", type=ast.literal_eval, default=False, help="Run distribute, default: false.")
  40. parser.add_argument("--dataset_path", type=str, default=None, help="Dataset path")
  41. parser.add_argument("--pre_trained", type=str, default=None, help="Pretrained file path.")
  42. parser.add_argument("--device_id", type=int, default=0, help="Device id, default: 0.")
  43. parser.add_argument("--device_num", type=int, default=1, help="Use device nums, default: 1.")
  44. parser.add_argument("--rank_id", type=int, default=0, help="Rank id, default: 0.")
  45. parser.add_argument('--ckpt_path', type=str, default='outputs/', help='Checkpoint save location. Default: outputs/')
  46. parser.add_argument('--outputs_dir', type=str, default='./', help='Checkpoint save location. Default: outputs/')
  47. args = parser.parse_args()
  48. def get_ms_timestamp():
  49. t = time.time()
  50. return int(round(t * 1000))
  51. time_stamp_init = False
  52. time_stamp_first = 0
  53. class LossCallBack(Callback):
  54. """
  55. Monitor the loss in training.
  56. If the loss is NAN or INF terminating training.
  57. Note:
  58. If per_print_times is 0 do not print loss.
  59. Args:
  60. per_print_times (int): Print loss every times. Default: 1.
  61. """
  62. def __init__(self, per_print_times=1, rank_id=0):
  63. super(LossCallBack, self).__init__()
  64. if not isinstance(per_print_times, int) or per_print_times < 0:
  65. raise ValueError("print_step must be int and >= 0.")
  66. self._per_print_times = per_print_times
  67. self.rank_id = rank_id
  68. global time_stamp_init, time_stamp_first
  69. if not time_stamp_init:
  70. time_stamp_first = get_ms_timestamp()
  71. time_stamp_init = True
  72. def step_end(self, run_context):
  73. """Monitor the loss in training."""
  74. global time_stamp_first
  75. time_stamp_current = get_ms_timestamp()
  76. cb_params = run_context.original_args()
  77. print("time: {}, epoch: {}, step: {}, outputs are {}".format(time_stamp_current - time_stamp_first,
  78. cb_params.cur_epoch_num,
  79. cb_params.cur_step_num,
  80. str(cb_params.net_outputs)))
  81. with open("./loss_{}.log".format(self.rank_id), "a+") as f:
  82. if context.get_context("device_target") == "Ascend":
  83. f.write("time: {}, epoch: {}, step: {}, loss: {}, overflow: {}, loss_scale: {}".format(
  84. time_stamp_current - time_stamp_first,
  85. cb_params.cur_epoch_num,
  86. cb_params.cur_step_num,
  87. str(cb_params.net_outputs[0].asnumpy()),
  88. str(cb_params.net_outputs[1].asnumpy()),
  89. str(cb_params.net_outputs[2].asnumpy())))
  90. else:
  91. f.write("time: {}, epoch: {}, step: {}, loss: {}".format(
  92. time_stamp_current - time_stamp_first,
  93. cb_params.cur_epoch_num,
  94. cb_params.cur_step_num,
  95. str(cb_params.net_outputs.asnumpy())))
  96. f.write('\n')
  97. if __name__ == '__main__':
  98. context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, \
  99. device_id=args.device_id, save_graphs=False)
  100. if args.device_target == "GPU":
  101. if config.compute_type != mstype.float32:
  102. logger.warning('GPU only support fp32 temporarily, run with fp32.')
  103. config.compute_type = mstype.float32
  104. if args.run_distribute:
  105. if args.device_target == "Ascend":
  106. rank = args.rank_id
  107. device_num = args.device_num
  108. context.set_auto_parallel_context(device_num=device_num,
  109. parallel_mode=ParallelMode.DATA_PARALLEL,
  110. gradients_mean=True)
  111. init()
  112. elif args.device_target == "GPU":
  113. init("nccl")
  114. context.set_auto_parallel_context(parallel_mode=ParallelMode.AUTO_PARALLEL,
  115. gradients_mean=True)
  116. else:
  117. raise ValueError(args.device_target)
  118. else:
  119. rank = 0
  120. device_num = 1
  121. mindrecord_file = args.dataset_path
  122. if not os.path.exists(mindrecord_file):
  123. print("dataset file {} not exists, please check!".format(mindrecord_file))
  124. raise ValueError(mindrecord_file)
  125. dataset = create_gru_dataset(epoch_count=config.num_epochs, batch_size=config.batch_size,
  126. dataset_path=mindrecord_file, rank_size=device_num, rank_id=rank)
  127. dataset_size = dataset.get_dataset_size()
  128. print("dataset size is {}".format(dataset_size))
  129. network = Seq2Seq(config)
  130. network = GRUWithLossCell(network)
  131. lr = dynamic_lr(config, dataset_size)
  132. opt = Adam(network.trainable_params(), learning_rate=lr)
  133. scale_manager = DynamicLossScaleManager(init_loss_scale=config.init_loss_scale_value,
  134. scale_factor=config.scale_factor,
  135. scale_window=config.scale_window)
  136. update_cell = scale_manager.get_update_cell()
  137. if args.device_target == "Ascend":
  138. netwithgrads = GRUTrainOneStepWithLossScaleCell(network, opt, update_cell)
  139. else:
  140. netwithgrads = GRUTrainOneStepCell(network, opt)
  141. time_cb = TimeMonitor(data_size=dataset_size)
  142. loss_cb = LossCallBack(rank_id=rank)
  143. cb = [time_cb, loss_cb]
  144. #Save Checkpoint
  145. if config.save_checkpoint:
  146. ckpt_config = CheckpointConfig(save_checkpoint_steps=config.ckpt_epoch * dataset_size,
  147. keep_checkpoint_max=config.keep_checkpoint_max)
  148. save_ckpt_path = os.path.join(args.outputs_dir, 'ckpt_'+str(args.rank_id)+'/')
  149. ckpt_cb = ModelCheckpoint(config=ckpt_config,
  150. directory=save_ckpt_path,
  151. prefix='{}'.format(args.rank_id))
  152. cb += [ckpt_cb]
  153. netwithgrads.set_train(True)
  154. model = Model(netwithgrads)
  155. model.train(config.num_epochs, dataset, callbacks=cb, dataset_sink_mode=True)