merge uapi paddleseg (#3718)

Author: zhangyubo0722

paddleseg/core/train.py

@@ -345,15 +345,29 @@ def train(model,
  "iter_{}".format(iter))
  if not os.path.isdir(current_save_dir):
  os.makedirs(current_save_dir)
+ states_dict = {
+ 'mIoU': mean_iou,
+ 'Acc': acc,
+ 'iter': iter
+ }
  paddle.save(model.state_dict(),
  os.path.join(current_save_dir, 'model.pdparams'))
  paddle.save(optimizer.state_dict(),
  os.path.join(current_save_dir, 'model.pdopt'))
+ paddle.save(states_dict,
+ os.path.join(current_save_dir, 'model.pdstates'))
 
  if use_ema:
+ ema_states_dict = {
+ 'mIoU': ema_mean_iou,
+ 'Acc': ema_acc,
+ 'iter': iter
+ }
  paddle.save(
  ema_model.state_dict(),
  os.path.join(current_save_dir, 'ema_model.pdparams'))
+ paddle.save(ema_states_dict,
+ os.path.join(current_save_dir, 'ema_model.pdstates'))
 
  save_models.append(current_save_dir)
  if len(save_models) > keep_checkpoint_max > 0:
@@ -369,6 +383,8 @@ def train(model,
  paddle.save(
  model.state_dict(),
  os.path.join(best_model_dir, 'model.pdparams'))
+ paddle.save(states_dict,
+ os.path.join(best_model_dir, 'model.pdstates'))
  elif mean_iou < best_mean_iou:
  stop_count += 1
 
@@ -391,6 +407,8 @@ def train(model,
  paddle.save(ema_model.state_dict(),
  os.path.join(best_ema_model_dir,
  'ema_model.pdparams'))
+ paddle.save(ema_states_dict,
+ os.path.join(best_ema_model_dir, 'ema_model.pdstates'))
  logger.info(
  '[EVAL] The EMA model with the best validation mIoU ({:.4f}) was saved at iter {}.'
  .format(best_ema_mean_iou, best_ema_model_iter))