save_load.py 6.1 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
def load_dygraph_params(config, model, logger, optimizer):
    ckp = config['Global']['checkpoints']
B
fix bug  
bingooo 已提交
94
    if ckp and os.path.exists(ckp + ".pdparams"):
D
Double_V 已提交
95 96 97 98 99 100
        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 {}
B
fix bug  
bingooo 已提交
101
        if not os.path.exists(pm) and not os.path.exists(pm + ".pdparams"):
D
Double_V 已提交
102 103 104 105 106 107 108 109 110
            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]
111 112 113 114
            else:
                logger.info(
                    f"The shape of model params {k1} {state_dict[k1].shape} not matched with loaded params {k2} {params[k2].shape} !"
                )
D
Double_V 已提交
115 116 117 118
        model.set_state_dict(new_state_dict)
        logger.info(f"loaded pretrained_model successful from {pm}")
        return {}

119

L
fix bug  
LDOUBLEV 已提交
120 121 122 123 124 125 126 127 128 129 130 131 132 133
def load_pretrained_params(model, path):
    if path is None:
        return False
    if not os.path.exists(path) and not os.path.exists(path + ".pdparams"):
        print(f"The pretrained_model {path} does not exists!")
        return False

    path = path if path.endswith('.pdparams') else path + '.pdparams'
    params = paddle.load(path)
    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]
L
LDOUBLEV 已提交
134 135 136 137
        else:
            print(
                f"The shape of model params {k1} {state_dict[k1].shape} not matched with loaded params {k2} {params[k2].shape} !"
            )
L
fix bug  
LDOUBLEV 已提交
138
    model.set_state_dict(new_state_dict)
L
LDOUBLEV 已提交
139
    print(f"load pretrain successful from {path}")
L
LDOUBLEV 已提交
140
    return model
D
Double_V 已提交
141

142

143
def save_model(model,
W
WenmuZhou 已提交
144 145 146 147 148 149
               optimizer,
               model_path,
               logger,
               is_best=False,
               prefix='ppocr',
               **kwargs):
L
LDOUBLEV 已提交
150 151 152
    """
    save model to the target path
    """
W
WenmuZhou 已提交
153 154
    _mkdir_if_not_exist(model_path, logger)
    model_prefix = os.path.join(model_path, prefix)
155
    paddle.save(model.state_dict(), model_prefix + '.pdparams')
W
WenmuZhou 已提交
156
    paddle.save(optimizer.state_dict(), model_prefix + '.pdopt')
W
WenmuZhou 已提交
157 158 159 160 161 162 163 164

    # 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))