save_load.py 5.2 KB
Newer Older
L
LDOUBLEV 已提交
1 2
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
W
WenmuZhou 已提交
3 4 5
# 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
L
LDOUBLEV 已提交
6 7 8
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
W
WenmuZhou 已提交
9 10 11 12 13
# 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.
L
LDOUBLEV 已提交
14 15 16 17 18 19 20

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import errno
import os
W
WenmuZhou 已提交
21 22
import pickle
import six
L
LDOUBLEV 已提交
23

W
WenmuZhou 已提交
24
import paddle
L
LDOUBLEV 已提交
25

littletomatodonkey's avatar
littletomatodonkey 已提交
26 27
from ppocr.utils.logging import get_logger

D
Double_V 已提交
28
__all__ = ['init_model', 'save_model', 'load_dygraph_params']
L
LDOUBLEV 已提交
29 30


W
WenmuZhou 已提交
31
def _mkdir_if_not_exist(path, logger):
L
LDOUBLEV 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
    """
    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))


littletomatodonkey's avatar
littletomatodonkey 已提交
47
def init_model(config, model, optimizer=None, lr_scheduler=None):
L
LDOUBLEV 已提交
48 49 50
    """
    load model from checkpoint or pretrained_model
    """
littletomatodonkey's avatar
littletomatodonkey 已提交
51
    logger = get_logger()
Y
YukSing 已提交
52 53 54
    global_config = config['Global']
    checkpoints = global_config.get('checkpoints')
    pretrained_model = global_config.get('pretrained_model')
W
WenmuZhou 已提交
55
    best_model_dict = {}
L
LDOUBLEV 已提交
56
    if checkpoints:
W
WenmuZhou 已提交
57 58 59 60
        assert os.path.exists(checkpoints + ".pdparams"), \
            "Given dir {}.pdparams not exist.".format(checkpoints)
        assert os.path.exists(checkpoints + ".pdopt"), \
            "Given dir {}.pdopt not exist.".format(checkpoints)
W
WenmuZhou 已提交
61 62
        para_dict = paddle.load(checkpoints + '.pdparams')
        opti_dict = paddle.load(checkpoints + '.pdopt')
W
WenmuZhou 已提交
63
        model.set_state_dict(para_dict)
W
WenmuZhou 已提交
64 65 66 67 68 69 70 71 72 73 74 75
        if optimizer is not None:
            optimizer.set_state_dict(opti_dict)

        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:
D
dyning 已提交
76 77
        if not isinstance(pretrained_model, list):
            pretrained_model = [pretrained_model]
78
        for pretrained in pretrained_model:
littletomatodonkey's avatar
littletomatodonkey 已提交
79 80 81 82 83 84
            if not (os.path.isdir(pretrained) or
                    os.path.exists(pretrained + '.pdparams')):
                raise ValueError("Model pretrain path {} does not "
                                 "exists.".format(pretrained))
            param_state_dict = paddle.load(pretrained + '.pdparams')
            model.set_state_dict(param_state_dict)
D
dyning 已提交
85 86
            logger.info("load pretrained model from {}".format(
                pretrained_model))
87
    else:
W
WenmuZhou 已提交
88 89
        logger.info('train from scratch')
    return best_model_dict
L
LDOUBLEV 已提交
90 91


D
Double_V 已提交
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119
def load_dygraph_params(config, model, logger, optimizer):
    ckp = config['Global']['checkpoints']
    if ckp and os.path.exists(ckp):
        pre_best_model_dict = init_model(config, model, optimizer)
        return pre_best_model_dict
    else:
        pm = config['Global']['pretrained_model']
        if pm is None:
            return {}
        if not os.path.exists(pm) or not os.path.exists(pm + ".pdparams"):
            logger.info(f"The pretrained_model {pm} does not exists!")
            return {}
        pm = pm if pm.endswith('.pdparams') else pm + '.pdparams'
        params = paddle.load(pm)
        state_dict = model.state_dict()
        new_state_dict = {}
        for k1, k2 in zip(state_dict.keys(), params.keys()):
            if list(state_dict[k1].shape) == list(params[k2].shape):
                new_state_dict[k1] = params[k2]
        else:
            logger.info(
                f"The shape of model params {k1} {state_dict[k1].shape} not matched with loaded params {k2} {params[k2].shape} !"
            )
        model.set_state_dict(new_state_dict)
        logger.info(f"loaded pretrained_model successful from {pm}")
        return {}


120
def save_model(model,
W
WenmuZhou 已提交
121 122 123 124 125 126
               optimizer,
               model_path,
               logger,
               is_best=False,
               prefix='ppocr',
               **kwargs):
L
LDOUBLEV 已提交
127 128 129
    """
    save model to the target path
    """
W
WenmuZhou 已提交
130 131
    _mkdir_if_not_exist(model_path, logger)
    model_prefix = os.path.join(model_path, prefix)
132
    paddle.save(model.state_dict(), model_prefix + '.pdparams')
W
WenmuZhou 已提交
133
    paddle.save(optimizer.state_dict(), model_prefix + '.pdopt')
W
WenmuZhou 已提交
134 135 136 137 138 139 140 141

    # save metric and config
    with open(model_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))