# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function import errno import os import pickle import six import paddle from ppocr.utils.logging import get_logger from ppocr.utils.network import maybe_download_params __all__ = ['load_model'] def _mkdir_if_not_exist(path, logger): """ mkdir if not exists, ignore the exception when multiprocess mkdir together """ if not os.path.exists(path): try: os.makedirs(path) except OSError as e: if e.errno == errno.EEXIST and os.path.isdir(path): logger.warning( 'be happy if some process has already created {}'.format( path)) else: raise OSError('Failed to mkdir {}'.format(path)) def load_model(config, model, optimizer=None, model_type='det'): """ load model from checkpoint or pretrained_model """ logger = get_logger() global_config = config['Global'] checkpoints = global_config.get('checkpoints') pretrained_model = global_config.get('pretrained_model') best_model_dict = {} is_float16 = False is_nlp_model = model_type == 'kie' and config["Architecture"][ "algorithm"] not in ["SDMGR"] if is_nlp_model is True: # NOTE: for kie model dsitillation, resume training is not supported now if config["Architecture"]["algorithm"] in ["Distillation"]: return best_model_dict checkpoints = config['Architecture']['Backbone']['checkpoints'] # load kie method metric if checkpoints: if os.path.exists(os.path.join(checkpoints, 'metric.states')): with open(os.path.join(checkpoints, 'metric.states'), 'rb') as f: states_dict = pickle.load(f) if six.PY2 else pickle.load( f, encoding='latin1') best_model_dict = states_dict.get('best_model_dict', {}) if 'epoch' in states_dict: best_model_dict['start_epoch'] = states_dict['epoch'] + 1 logger.info("resume from {}".format(checkpoints)) if optimizer is not None: if checkpoints[-1] in ['/', '\\']: checkpoints = checkpoints[:-1] if os.path.exists(checkpoints + '.pdopt'): optim_dict = paddle.load(checkpoints + '.pdopt') optimizer.set_state_dict(optim_dict) else: logger.warning( "{}.pdopt is not exists, params of optimizer is not loaded". format(checkpoints)) return best_model_dict if checkpoints: if checkpoints.endswith('.pdparams'): checkpoints = checkpoints.replace('.pdparams', '') assert os.path.exists(checkpoints + ".pdparams"), \ "The {}.pdparams does not exists!".format(checkpoints) # load params from trained model params = paddle.load(checkpoints + '.pdparams') state_dict = model.state_dict() new_state_dict = {} for key, value in state_dict.items(): if key not in params: logger.warning("{} not in loaded params {} !".format( key, params.keys())) continue pre_value = params[key] if pre_value.dtype == paddle.float16: is_float16 = True if pre_value.dtype != value.dtype: pre_value = pre_value.astype(value.dtype) if list(value.shape) == list(pre_value.shape): new_state_dict[key] = pre_value else: logger.warning( "The shape of model params {} {} not matched with loaded params shape {} !". format(key, value.shape, pre_value.shape)) model.set_state_dict(new_state_dict) if is_float16: logger.info( "The parameter type is float16, which is converted to float32 when loading" ) if optimizer is not None: if os.path.exists(checkpoints + '.pdopt'): optim_dict = paddle.load(checkpoints + '.pdopt') optimizer.set_state_dict(optim_dict) else: logger.warning( "{}.pdopt is not exists, params of optimizer is not loaded". format(checkpoints)) if os.path.exists(checkpoints + '.states'): with open(checkpoints + '.states', 'rb') as f: states_dict = pickle.load(f) if six.PY2 else pickle.load( f, encoding='latin1') best_model_dict = states_dict.get('best_model_dict', {}) if 'epoch' in states_dict: best_model_dict['start_epoch'] = states_dict['epoch'] + 1 logger.info("resume from {}".format(checkpoints)) elif pretrained_model: is_float16 = load_pretrained_params(model, pretrained_model) else: logger.info('train from scratch') best_model_dict['is_float16'] = is_float16 return best_model_dict def load_pretrained_params(model, path): logger = get_logger() path = maybe_download_params(path) if path.endswith('.pdparams'): path = path.replace('.pdparams', '') assert os.path.exists(path + ".pdparams"), \ "The {}.pdparams does not exists!".format(path) params = paddle.load(path + '.pdparams') state_dict = model.state_dict() new_state_dict = {} is_float16 = False for k1 in params.keys(): if k1 not in state_dict.keys(): logger.warning("The pretrained params {} not in model".format(k1)) else: if params[k1].dtype == paddle.float16: is_float16 = True if params[k1].dtype != state_dict[k1].dtype: params[k1] = params[k1].astype(state_dict[k1].dtype) if list(state_dict[k1].shape) == list(params[k1].shape): new_state_dict[k1] = params[k1] else: logger.warning( "The shape of model params {} {} not matched with loaded params {} {} !". format(k1, state_dict[k1].shape, k1, params[k1].shape)) model.set_state_dict(new_state_dict) if is_float16: logger.info( "The parameter type is float16, which is converted to float32 when loading" ) logger.info("load pretrain successful from {}".format(path)) return is_float16 def save_model(model, optimizer, model_path, logger, config, is_best=False, prefix='ppocr', **kwargs): """ save model to the target path """ _mkdir_if_not_exist(model_path, logger) model_prefix = os.path.join(model_path, prefix) if prefix == 'best_accuracy': best_model_path = os.path.join(model_path, 'best_model') _mkdir_if_not_exist(best_model_path, logger) paddle.save(optimizer.state_dict(), model_prefix + '.pdopt') if prefix == 'best_accuracy': paddle.save(optimizer.state_dict(), os.path.join(best_model_path, 'model.pdopt')) is_nlp_model = config['Architecture']["model_type"] == 'kie' and config[ "Architecture"]["algorithm"] not in ["SDMGR"] if is_nlp_model is not True: paddle.save(model.state_dict(), model_prefix + '.pdparams') metric_prefix = model_prefix if prefix == 'best_accuracy': paddle.save(model.state_dict(), os.path.join(best_model_path, 'model.pdparams')) else: # for kie system, we follow the save/load rules in NLP if config['Global']['distributed']: arch = model._layers else: arch = model if config["Architecture"]["algorithm"] in ["Distillation"]: arch = arch.Student arch.backbone.model.save_pretrained(model_prefix) metric_prefix = os.path.join(model_prefix, 'metric') if prefix == 'best_accuracy': arch.backbone.model.save_pretrained(best_model_path) # save metric and config with open(metric_prefix + '.states', 'wb') as f: pickle.dump(kwargs, f, protocol=2) if is_best: logger.info('save best model is to {}'.format(model_prefix)) else: logger.info("save model in {}".format(model_prefix))