train.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. # Copyright (c) 2021 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, division, print_function
  15. import time
  16. import paddle
  17. from paddlex.ppcls.engine.train.utils import update_loss, update_metric, log_info
  18. from paddlex.ppcls.utils import profiler
  19. def train_epoch(engine, epoch_id, print_batch_step):
  20. tic = time.time()
  21. for iter_id, batch in enumerate(engine.train_dataloader):
  22. if iter_id >= engine.max_iter:
  23. break
  24. profiler.add_profiler_step(engine.config["profiler_options"])
  25. if iter_id == 5:
  26. for key in engine.time_info:
  27. engine.time_info[key].reset()
  28. engine.time_info["reader_cost"].update(time.time() - tic)
  29. if engine.use_dali:
  30. batch = [
  31. paddle.to_tensor(batch[0]['data']),
  32. paddle.to_tensor(batch[0]['label'])
  33. ]
  34. batch_size = batch[0].shape[0]
  35. if not engine.config["Global"].get("use_multilabel", False):
  36. batch[1] = batch[1].reshape([batch_size, -1])
  37. engine.global_step += 1
  38. # image input
  39. if engine.amp:
  40. with paddle.amp.auto_cast(custom_black_list={
  41. "flatten_contiguous_range", "greater_than"
  42. }):
  43. out = forward(engine, batch)
  44. else:
  45. out = forward(engine, batch)
  46. loss_dict = engine.train_loss_func(out, batch[1])
  47. # step opt and lr
  48. if engine.amp:
  49. scaled = engine.scaler.scale(loss_dict["loss"])
  50. scaled.backward()
  51. engine.scaler.minimize(engine.optimizer, scaled)
  52. else:
  53. loss_dict["loss"].backward()
  54. engine.optimizer.step()
  55. engine.optimizer.clear_grad()
  56. engine.lr_sch.step()
  57. # below code just for logging
  58. # update metric_for_logger
  59. update_metric(engine, out, batch, batch_size)
  60. # update_loss_for_logger
  61. update_loss(engine, loss_dict, batch_size)
  62. engine.time_info["batch_cost"].update(time.time() - tic)
  63. if iter_id % print_batch_step == 0:
  64. log_info(engine, batch_size, epoch_id, iter_id)
  65. tic = time.time()
  66. def forward(engine, batch):
  67. if not engine.is_rec:
  68. return engine.model(batch[0])
  69. else:
  70. return engine.model(batch[0], batch[1])