checkpoint.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216
  1. # Copyright (c) 2020 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
  15. from __future__ import division
  16. from __future__ import print_function
  17. from __future__ import unicode_literals
  18. import errno
  19. import os
  20. import time
  21. import re
  22. import numpy as np
  23. import paddle
  24. import paddle.nn as nn
  25. from .download import get_weights_path
  26. from .logger import setup_logger
  27. logger = setup_logger(__name__)
  28. def is_url(path):
  29. """
  30. Whether path is URL.
  31. Args:
  32. path (string): URL string or not.
  33. """
  34. return path.startswith('http://') \
  35. or path.startswith('https://') \
  36. or path.startswith('ppdet://')
  37. def _get_unique_endpoints(trainer_endpoints):
  38. # Sorting is to avoid different environmental variables for each card
  39. trainer_endpoints.sort()
  40. ips = set()
  41. unique_endpoints = set()
  42. for endpoint in trainer_endpoints:
  43. ip = endpoint.split(":")[0]
  44. if ip in ips:
  45. continue
  46. ips.add(ip)
  47. unique_endpoints.add(endpoint)
  48. logger.info("unique_endpoints {}".format(unique_endpoints))
  49. return unique_endpoints
  50. def get_weights_path_dist(path):
  51. env = os.environ
  52. if 'PADDLE_TRAINERS_NUM' in env and 'PADDLE_TRAINER_ID' in env:
  53. trainer_id = int(env['PADDLE_TRAINER_ID'])
  54. num_trainers = int(env['PADDLE_TRAINERS_NUM'])
  55. if num_trainers <= 1:
  56. path = get_weights_path(path)
  57. else:
  58. from paddlex.ppdet.utils.download import map_path, WEIGHTS_HOME
  59. weight_path = map_path(path, WEIGHTS_HOME)
  60. lock_path = weight_path + '.lock'
  61. if not os.path.exists(weight_path):
  62. from paddle.distributed import ParallelEnv
  63. unique_endpoints = _get_unique_endpoints(ParallelEnv()
  64. .trainer_endpoints[:])
  65. try:
  66. os.makedirs(os.path.dirname(weight_path))
  67. except OSError as e:
  68. if e.errno != errno.EEXIST:
  69. raise
  70. with open(lock_path, 'w'): # touch
  71. os.utime(lock_path, None)
  72. if ParallelEnv().current_endpoint in unique_endpoints:
  73. get_weights_path(path)
  74. os.remove(lock_path)
  75. else:
  76. while os.path.exists(lock_path):
  77. time.sleep(1)
  78. path = weight_path
  79. else:
  80. path = get_weights_path(path)
  81. return path
  82. def _strip_postfix(path):
  83. path, ext = os.path.splitext(path)
  84. assert ext in ['', '.pdparams', '.pdopt', '.pdmodel'], \
  85. "Unknown postfix {} from weights".format(ext)
  86. return path
  87. def load_weight(model, weight, optimizer=None):
  88. if is_url(weight):
  89. weight = get_weights_path_dist(weight)
  90. path = _strip_postfix(weight)
  91. pdparam_path = path + '.pdparams'
  92. if not os.path.exists(pdparam_path):
  93. raise ValueError("Model pretrain path {} does not "
  94. "exists.".format(pdparam_path))
  95. param_state_dict = paddle.load(pdparam_path)
  96. model_dict = model.state_dict()
  97. model_weight = {}
  98. incorrect_keys = 0
  99. for key in model_dict.keys():
  100. if key in param_state_dict.keys():
  101. model_weight[key] = param_state_dict[key]
  102. else:
  103. logger.info('Unmatched key: {}'.format(key))
  104. incorrect_keys += 1
  105. assert incorrect_keys == 0, "Load weight {} incorrectly, \
  106. {} keys unmatched, please check again.".format(weight,
  107. incorrect_keys)
  108. logger.info('Finish resuming model weights: {}'.format(pdparam_path))
  109. model.set_dict(model_weight)
  110. last_epoch = 0
  111. if optimizer is not None and os.path.exists(path + '.pdopt'):
  112. optim_state_dict = paddle.load(path + '.pdopt')
  113. # to solve resume bug, will it be fixed in paddle 2.0
  114. for key in optimizer.state_dict().keys():
  115. if not key in optim_state_dict.keys():
  116. optim_state_dict[key] = optimizer.state_dict()[key]
  117. if 'last_epoch' in optim_state_dict:
  118. last_epoch = optim_state_dict.pop('last_epoch')
  119. optimizer.set_state_dict(optim_state_dict)
  120. return last_epoch
  121. def load_pretrain_weight(model, pretrain_weight):
  122. if is_url(pretrain_weight):
  123. pretrain_weight = get_weights_path_dist(pretrain_weight)
  124. path = _strip_postfix(pretrain_weight)
  125. if not (os.path.isdir(path) or os.path.isfile(path) or
  126. os.path.exists(path + '.pdparams')):
  127. raise ValueError("Model pretrain path `{}` does not exists. "
  128. "If you don't want to load pretrain model, "
  129. "please delete `pretrain_weights` field in "
  130. "config file.".format(path))
  131. model_dict = model.state_dict()
  132. weights_path = path + '.pdparams'
  133. param_state_dict = paddle.load(weights_path)
  134. ignore_weights = set()
  135. # hack: fit for faster rcnn. Pretrain weights contain prefix of 'backbone'
  136. # while res5 module is located in bbox_head.head. Replace the prefix of
  137. # res5 with 'bbox_head.head' to load pretrain weights correctly.
  138. for k in param_state_dict.keys():
  139. if 'backbone.res5' in k:
  140. new_k = k.replace('backbone', 'bbox_head.head')
  141. if new_k in model_dict.keys():
  142. value = param_state_dict.pop(k)
  143. param_state_dict[new_k] = value
  144. for name, weight in param_state_dict.items():
  145. if name in model_dict.keys():
  146. if list(weight.shape) != list(model_dict[name].shape):
  147. logger.info(
  148. '{} not used, shape {} unmatched with {} in model.'.format(
  149. name, weight.shape, list(model_dict[name].shape)))
  150. ignore_weights.add(name)
  151. else:
  152. logger.info('Redundant weight {} and ignore it.'.format(name))
  153. ignore_weights.add(name)
  154. for weight in ignore_weights:
  155. param_state_dict.pop(weight, None)
  156. model.set_dict(param_state_dict)
  157. logger.info('Finish loading model weights: {}'.format(weights_path))
  158. def save_model(model, optimizer, save_dir, save_name, last_epoch):
  159. """
  160. save model into disk.
  161. Args:
  162. model (paddle.nn.Layer): the Layer instalce to save parameters.
  163. optimizer (paddle.optimizer.Optimizer): the Optimizer instance to
  164. save optimizer states.
  165. save_dir (str): the directory to be saved.
  166. save_name (str): the path to be saved.
  167. last_epoch (int): the epoch index.
  168. """
  169. if paddle.distributed.get_rank() != 0:
  170. return
  171. if not os.path.exists(save_dir):
  172. os.makedirs(save_dir)
  173. save_path = os.path.join(save_dir, save_name)
  174. if isinstance(model, nn.Layer):
  175. paddle.save(model.state_dict(), save_path + ".pdparams")
  176. else:
  177. assert isinstance(model,
  178. dict), 'model is not a instance of nn.layer or dict'
  179. paddle.save(model, save_path + ".pdparams")
  180. state_dict = optimizer.state_dict()
  181. state_dict['last_epoch'] = last_epoch
  182. paddle.save(state_dict, save_path + ".pdopt")
  183. logger.info("Save checkpoint: {}".format(save_dir))