engine.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436
  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. # http://www.apache.org/licenses/LICENSE-2.0
  7. #
  8. # Unless required by applicable law or agreed to in writing, software
  9. # distributed under the License is distributed on an "AS IS" BASIS,
  10. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11. # See the License for the specific language governing permissions and
  12. # limitations under the License.
  13. from __future__ import absolute_import
  14. from __future__ import division
  15. from __future__ import print_function
  16. import os
  17. import platform
  18. import paddle
  19. import paddle.distributed as dist
  20. from visualdl import LogWriter
  21. from paddle import nn
  22. import numpy as np
  23. import random
  24. from paddlex.ppcls.utils.check import check_gpu
  25. from paddlex.ppcls.utils.misc import AverageMeter
  26. from paddlex.ppcls.utils import logger
  27. from paddlex.ppcls.utils.logger import init_logger
  28. from paddlex.ppcls.utils.config import print_config
  29. from paddlex.ppcls.data import build_dataloader
  30. from paddlex.ppcls.arch import build_model, RecModel, DistillationModel
  31. from paddlex.ppcls.arch import apply_to_static
  32. from paddlex.ppcls.loss import build_loss
  33. from paddlex.ppcls.metric import build_metrics
  34. from paddlex.ppcls.optimizer import build_optimizer
  35. from paddlex.ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
  36. from paddlex.ppcls.utils.save_load import init_model
  37. from paddlex.ppcls.utils import save_load
  38. from paddlex.ppcls.data.utils.get_image_list import get_image_list
  39. from paddlex.ppcls.data.postprocess import build_postprocess
  40. from paddlex.ppcls.data import create_operators
  41. from paddlex.ppcls.engine.train import train_epoch
  42. from paddlex.ppcls.engine import evaluation
  43. from paddlex.ppcls.arch.gears.identity_head import IdentityHead
  44. from paddlex.ppcls.engine.slim import get_pruner, get_quaner
  45. class Engine(object):
  46. def __init__(self, config, mode="train"):
  47. assert mode in ["train", "eval", "infer", "export"]
  48. self.mode = mode
  49. self.config = config
  50. self.eval_mode = self.config["Global"].get("eval_mode",
  51. "classification")
  52. if "Head" in self.config["Arch"]:
  53. self.is_rec = True
  54. else:
  55. self.is_rec = False
  56. # set seed
  57. seed = self.config["Global"].get("seed", False)
  58. if seed or seed == 0:
  59. assert isinstance(seed, int), "The 'seed' must be a integer!"
  60. paddle.seed(seed)
  61. np.random.seed(seed)
  62. random.seed(seed)
  63. # init logger
  64. self.output_dir = self.config['Global']['output_dir']
  65. log_file = os.path.join(self.output_dir, self.config["Arch"]["name"],
  66. f"{mode}.log")
  67. init_logger(name='root', log_file=log_file)
  68. print_config(config)
  69. # init train_func and eval_func
  70. assert self.eval_mode in ["classification", "retrieval"], logger.error(
  71. "Invalid eval mode: {}".format(self.eval_mode))
  72. self.train_epoch_func = train_epoch
  73. self.eval_func = getattr(evaluation, self.eval_mode + "_eval")
  74. self.use_dali = self.config['Global'].get("use_dali", False)
  75. # for visualdl
  76. self.vdl_writer = None
  77. if self.config['Global']['use_visualdl'] and mode == "train":
  78. vdl_writer_path = os.path.join(self.output_dir, "vdl")
  79. if not os.path.exists(vdl_writer_path):
  80. os.makedirs(vdl_writer_path)
  81. self.vdl_writer = LogWriter(logdir=vdl_writer_path)
  82. # set device
  83. assert self.config["Global"]["device"] in ["cpu", "gpu", "xpu", "npu"]
  84. self.device = paddle.set_device(self.config["Global"]["device"])
  85. logger.info('train with paddle {} and device {}'.format(
  86. paddle.__version__, self.device))
  87. # AMP training
  88. self.amp = True if "AMP" in self.config else False
  89. if self.amp and self.config["AMP"] is not None:
  90. self.scale_loss = self.config["AMP"].get("scale_loss", 1.0)
  91. self.use_dynamic_loss_scaling = self.config["AMP"].get(
  92. "use_dynamic_loss_scaling", False)
  93. else:
  94. self.scale_loss = 1.0
  95. self.use_dynamic_loss_scaling = False
  96. if self.amp:
  97. AMP_RELATED_FLAGS_SETTING = {
  98. 'FLAGS_cudnn_batchnorm_spatial_persistent': 1,
  99. 'FLAGS_max_inplace_grad_add': 8,
  100. }
  101. paddle.fluid.set_flags(AMP_RELATED_FLAGS_SETTING)
  102. #TODO(gaotingquan): support rec
  103. class_num = config["Arch"].get("class_num", None)
  104. self.config["DataLoader"].update({"class_num": class_num})
  105. # build dataloader
  106. if self.mode == 'train':
  107. self.train_dataloader = build_dataloader(
  108. self.config["DataLoader"], "Train", self.device, self.use_dali)
  109. if self.mode == "eval" or (self.mode == "train" and
  110. self.config["Global"]["eval_during_train"]):
  111. if self.eval_mode == "classification":
  112. self.eval_dataloader = build_dataloader(
  113. self.config["DataLoader"], "Eval", self.device,
  114. self.use_dali)
  115. elif self.eval_mode == "retrieval":
  116. self.gallery_query_dataloader = None
  117. if len(self.config["DataLoader"]["Eval"].keys()) == 1:
  118. key = list(self.config["DataLoader"]["Eval"].keys())[0]
  119. self.gallery_query_dataloader = build_dataloader(
  120. self.config["DataLoader"]["Eval"], key, self.device,
  121. self.use_dali)
  122. else:
  123. self.gallery_dataloader = build_dataloader(
  124. self.config["DataLoader"]["Eval"], "Gallery",
  125. self.device, self.use_dali)
  126. self.query_dataloader = build_dataloader(
  127. self.config["DataLoader"]["Eval"], "Query",
  128. self.device, self.use_dali)
  129. # build loss
  130. if self.mode == "train":
  131. loss_info = self.config["Loss"]["Train"]
  132. self.train_loss_func = build_loss(loss_info)
  133. if self.mode == "eval" or (self.mode == "train" and
  134. self.config["Global"]["eval_during_train"]):
  135. loss_config = self.config.get("Loss", None)
  136. if loss_config is not None:
  137. loss_config = loss_config.get("Eval")
  138. if loss_config is not None:
  139. self.eval_loss_func = build_loss(loss_config)
  140. else:
  141. self.eval_loss_func = None
  142. else:
  143. self.eval_loss_func = None
  144. # build metric
  145. if self.mode == 'train':
  146. metric_config = self.config.get("Metric")
  147. if metric_config is not None:
  148. metric_config = metric_config.get("Train")
  149. if metric_config is not None:
  150. self.train_metric_func = build_metrics(metric_config)
  151. else:
  152. self.train_metric_func = None
  153. else:
  154. self.train_metric_func = None
  155. if self.mode == "eval" or (self.mode == "train" and
  156. self.config["Global"]["eval_during_train"]):
  157. metric_config = self.config.get("Metric")
  158. if self.eval_mode == "classification":
  159. if metric_config is not None:
  160. metric_config = metric_config.get("Eval")
  161. if metric_config is not None:
  162. self.eval_metric_func = build_metrics(metric_config)
  163. elif self.eval_mode == "retrieval":
  164. if metric_config is None:
  165. metric_config = [{"name": "Recallk", "topk": (1, 5)}]
  166. else:
  167. metric_config = metric_config["Eval"]
  168. self.eval_metric_func = build_metrics(metric_config)
  169. else:
  170. self.eval_metric_func = None
  171. # build model
  172. self.model = build_model(self.config["Arch"])
  173. # set @to_static for benchmark, skip this by default.
  174. apply_to_static(self.config, self.model)
  175. # for slim
  176. self.pruner = get_pruner(self.config, self.model)
  177. self.quanter = get_quaner(self.config, self.model)
  178. # load_pretrain
  179. if self.config["Global"]["pretrained_model"] is not None:
  180. if self.config["Global"]["pretrained_model"].startswith("http"):
  181. load_dygraph_pretrain_from_url(
  182. self.model, self.config["Global"]["pretrained_model"])
  183. else:
  184. load_dygraph_pretrain(
  185. self.model, self.config["Global"]["pretrained_model"])
  186. # build optimizer
  187. if self.mode == 'train':
  188. self.optimizer, self.lr_sch = build_optimizer(
  189. self.config["Optimizer"], self.config["Global"]["epochs"],
  190. len(self.train_dataloader), [self.model])
  191. # for distributed
  192. self.config["Global"][
  193. "distributed"] = paddle.distributed.get_world_size() != 1
  194. if self.config["Global"]["distributed"]:
  195. dist.init_parallel_env()
  196. if self.config["Global"]["distributed"]:
  197. self.model = paddle.DataParallel(self.model)
  198. # build postprocess for infer
  199. if self.mode == 'infer':
  200. self.preprocess_func = create_operators(self.config["Infer"][
  201. "transforms"])
  202. self.postprocess_func = build_postprocess(self.config["Infer"][
  203. "PostProcess"])
  204. def train(self):
  205. assert self.mode == "train"
  206. print_batch_step = self.config['Global']['print_batch_step']
  207. save_interval = self.config["Global"]["save_interval"]
  208. best_metric = {
  209. "metric": 0.0,
  210. "epoch": 0,
  211. }
  212. # key:
  213. # val: metrics list word
  214. self.output_info = dict()
  215. self.time_info = {
  216. "batch_cost": AverageMeter(
  217. "batch_cost", '.5f', postfix=" s,"),
  218. "reader_cost": AverageMeter(
  219. "reader_cost", ".5f", postfix=" s,"),
  220. }
  221. # global iter counter
  222. self.global_step = 0
  223. if self.config["Global"]["checkpoints"] is not None:
  224. metric_info = init_model(self.config["Global"], self.model,
  225. self.optimizer)
  226. if metric_info is not None:
  227. best_metric.update(metric_info)
  228. # for amp training
  229. if self.amp:
  230. self.scaler = paddle.amp.GradScaler(
  231. init_loss_scaling=self.scale_loss,
  232. use_dynamic_loss_scaling=self.use_dynamic_loss_scaling)
  233. self.max_iter = len(self.train_dataloader) - 1 if platform.system(
  234. ) == "Windows" else len(self.train_dataloader)
  235. for epoch_id in range(best_metric["epoch"] + 1,
  236. self.config["Global"]["epochs"] + 1):
  237. acc = 0.0
  238. # for one epoch train
  239. self.train_epoch_func(self, epoch_id, print_batch_step)
  240. if self.use_dali:
  241. self.train_dataloader.reset()
  242. metric_msg = ", ".join([
  243. "{}: {:.5f}".format(key, self.output_info[key].avg)
  244. for key in self.output_info
  245. ])
  246. logger.info("[Train][Epoch {}/{}][Avg]{}".format(
  247. epoch_id, self.config["Global"]["epochs"], metric_msg))
  248. self.output_info.clear()
  249. # eval model and save model if possible
  250. if self.config["Global"][
  251. "eval_during_train"] and epoch_id % self.config["Global"][
  252. "eval_interval"] == 0:
  253. acc = self.eval(epoch_id)
  254. if acc > best_metric["metric"]:
  255. best_metric["metric"] = acc
  256. best_metric["epoch"] = epoch_id
  257. save_load.save_model(
  258. self.model,
  259. self.optimizer,
  260. best_metric,
  261. self.output_dir,
  262. model_name=self.config["Arch"]["name"],
  263. prefix="best_model")
  264. logger.info("[Eval][Epoch {}][best metric: {}]".format(
  265. epoch_id, best_metric["metric"]))
  266. logger.scaler(
  267. name="eval_acc",
  268. value=acc,
  269. step=epoch_id,
  270. writer=self.vdl_writer)
  271. self.model.train()
  272. # save model
  273. if epoch_id % save_interval == 0:
  274. save_load.save_model(
  275. self.model,
  276. self.optimizer, {"metric": acc,
  277. "epoch": epoch_id},
  278. self.output_dir,
  279. model_name=self.config["Arch"]["name"],
  280. prefix="epoch_{}".format(epoch_id))
  281. # save the latest model
  282. save_load.save_model(
  283. self.model,
  284. self.optimizer, {"metric": acc,
  285. "epoch": epoch_id},
  286. self.output_dir,
  287. model_name=self.config["Arch"]["name"],
  288. prefix="latest")
  289. if self.vdl_writer is not None:
  290. self.vdl_writer.close()
  291. @paddle.no_grad()
  292. def eval(self, epoch_id=0):
  293. assert self.mode in ["train", "eval"]
  294. self.model.eval()
  295. eval_result = self.eval_func(self, epoch_id)
  296. self.model.train()
  297. return eval_result
  298. @paddle.no_grad()
  299. def infer(self):
  300. assert self.mode == "infer" and self.eval_mode == "classification"
  301. total_trainer = paddle.distributed.get_world_size()
  302. local_rank = paddle.distributed.get_rank()
  303. image_list = get_image_list(self.config["Infer"]["infer_imgs"])
  304. # data split
  305. image_list = image_list[local_rank::total_trainer]
  306. batch_size = self.config["Infer"]["batch_size"]
  307. self.model.eval()
  308. batch_data = []
  309. image_file_list = []
  310. for idx, image_file in enumerate(image_list):
  311. with open(image_file, 'rb') as f:
  312. x = f.read()
  313. for process in self.preprocess_func:
  314. x = process(x)
  315. batch_data.append(x)
  316. image_file_list.append(image_file)
  317. if len(batch_data) >= batch_size or idx == len(image_list) - 1:
  318. batch_tensor = paddle.to_tensor(batch_data)
  319. out = self.model(batch_tensor)
  320. if isinstance(out, list):
  321. out = out[0]
  322. if isinstance(out, dict):
  323. out = out["output"]
  324. result = self.postprocess_func(out, image_file_list)
  325. print(result)
  326. batch_data.clear()
  327. image_file_list.clear()
  328. def export(self):
  329. assert self.mode == "export"
  330. use_multilabel = self.config["Global"].get("use_multilabel", False)
  331. model = ExportModel(self.config["Arch"], self.model, use_multilabel)
  332. if self.config["Global"]["pretrained_model"] is not None:
  333. load_dygraph_pretrain(model.base_model,
  334. self.config["Global"]["pretrained_model"])
  335. model.eval()
  336. save_path = os.path.join(self.config["Global"]["save_inference_dir"],
  337. "inference")
  338. if self.quanter:
  339. self.quanter.save_quantized_model(
  340. model.base_model,
  341. save_path,
  342. input_spec=[
  343. paddle.static.InputSpec(
  344. shape=[None] + self.config["Global"]["image_shape"],
  345. dtype='float32')
  346. ])
  347. else:
  348. model = paddle.jit.to_static(
  349. model,
  350. input_spec=[
  351. paddle.static.InputSpec(
  352. shape=[None] + self.config["Global"]["image_shape"],
  353. dtype='float32')
  354. ])
  355. paddle.jit.save(model, save_path)
  356. class ExportModel(nn.Layer):
  357. """
  358. ExportModel: add softmax onto the model
  359. """
  360. def __init__(self, config, model, use_multilabel):
  361. super().__init__()
  362. self.base_model = model
  363. # we should choose a final model to export
  364. if isinstance(self.base_model, DistillationModel):
  365. self.infer_model_name = config["infer_model_name"]
  366. else:
  367. self.infer_model_name = None
  368. self.infer_output_key = config.get("infer_output_key", None)
  369. if self.infer_output_key == "features" and isinstance(self.base_model,
  370. RecModel):
  371. self.base_model.head = IdentityHead()
  372. if use_multilabel:
  373. self.out_act = nn.Sigmoid()
  374. else:
  375. if config.get("infer_add_softmax", True):
  376. self.out_act = nn.Softmax(axis=-1)
  377. else:
  378. self.out_act = None
  379. def eval(self):
  380. self.training = False
  381. for layer in self.sublayers():
  382. layer.training = False
  383. layer.eval()
  384. def forward(self, x):
  385. x = self.base_model(x)
  386. if isinstance(x, list):
  387. x = x[0]
  388. if self.infer_model_name is not None:
  389. x = x[self.infer_model_name]
  390. if self.infer_output_key is not None:
  391. x = x[self.infer_output_key]
  392. if self.out_act is not None:
  393. x = self.out_act(x)
  394. return x