save_load.py 4.95 KB
Newer Older
LDOUBLEV's avatar
LDOUBLEV committed
1
2
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
WenmuZhou's avatar
WenmuZhou committed
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
LDOUBLEV's avatar
LDOUBLEV committed
6
7
8
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
WenmuZhou's avatar
WenmuZhou committed
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.
LDOUBLEV's avatar
LDOUBLEV committed
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
WenmuZhou's avatar
WenmuZhou committed
21
22
import pickle
import six
LDOUBLEV's avatar
LDOUBLEV committed
23

WenmuZhou's avatar
WenmuZhou committed
24
import paddle
LDOUBLEV's avatar
LDOUBLEV committed
25

littletomatodonkey's avatar
littletomatodonkey committed
26
27
from ppocr.utils.logging import get_logger

28
__all__ = ['load_model']
LDOUBLEV's avatar
LDOUBLEV committed
29
30


WenmuZhou's avatar
WenmuZhou committed
31
def _mkdir_if_not_exist(path, logger):
LDOUBLEV's avatar
LDOUBLEV committed
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))


47
def load_model(config, model, optimizer=None):
LDOUBLEV's avatar
LDOUBLEV committed
48
49
50
    """
    load model from checkpoint or pretrained_model
    """
littletomatodonkey's avatar
littletomatodonkey committed
51
    logger = get_logger()
YukSing's avatar
YukSing committed
52
53
54
    global_config = config['Global']
    checkpoints = global_config.get('checkpoints')
    pretrained_model = global_config.get('pretrained_model')
WenmuZhou's avatar
WenmuZhou committed
55
    best_model_dict = {}
LDOUBLEV's avatar
LDOUBLEV committed
56
    if checkpoints:
57
58
        if checkpoints.endswith('pdparams'):
            checkpoints = checkpoints.replace('.pdparams', '')
59
        assert os.path.exists(checkpoints + ".pdparams"), \
60
61
            "The {}.pdparams does not exists!".format(checkpoints)

62
63
64
65
66
67
        # 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:
68
69
                logger.warning("{} not in loaded params {} !".format(
                    key, params.keys()))
70
71
72
73
74
            pre_value = params[key]
            if list(value.shape) == list(pre_value.shape):
                new_state_dict[key] = pre_value
            else:
                logger.warning(
75
76
                    "The shape of model params {} {} not matched with loaded params shape {} !".
                    format(key, value.shape, pre_value.shape))
77
78
        model.set_state_dict(new_state_dict)

79
        optim_dict = paddle.load(checkpoints + '.pdopt')
WenmuZhou's avatar
WenmuZhou committed
80
        if optimizer is not None:
81
            optimizer.set_state_dict(optim_dict)
WenmuZhou's avatar
WenmuZhou committed
82
83
84
85
86
87
88
89
90
91

        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:
92
        load_pretrained_params(model, pretrained_model)
93
    else:
WenmuZhou's avatar
WenmuZhou committed
94
95
        logger.info('train from scratch')
    return best_model_dict
LDOUBLEV's avatar
LDOUBLEV committed
96
97


LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
98
def load_pretrained_params(model, path):
99
100
101
102
    logger = get_logger()
    if path.endswith('pdparams'):
        path = path.replace('.pdparams', '')
    assert os.path.exists(path + ".pdparams"), \
103
        "The {}.pdparams does not exists!".format(path)
104
105

    params = paddle.load(path + '.pdparams')
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
106
107
108
109
110
    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]
LDOUBLEV's avatar
LDOUBLEV committed
111
        else:
112
            logger.warning(
113
114
                "The shape of model params {} {} not matched with loaded params {} {} !".
                format(k1, state_dict[k1].shape, k2, params[k2].shape))
LDOUBLEV's avatar
fix bug  
LDOUBLEV committed
115
    model.set_state_dict(new_state_dict)
116
    logger.info("load pretrain successful from {}".format(path))
LDOUBLEV's avatar
LDOUBLEV committed
117
    return model
Double_V's avatar
Double_V committed
118

119

120
def save_model(model,
WenmuZhou's avatar
WenmuZhou committed
121
122
123
124
125
126
               optimizer,
               model_path,
               logger,
               is_best=False,
               prefix='ppocr',
               **kwargs):
LDOUBLEV's avatar
LDOUBLEV committed
127
128
129
    """
    save model to the target path
    """
WenmuZhou's avatar
WenmuZhou committed
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')
WenmuZhou's avatar
WenmuZhou committed
133
    paddle.save(optimizer.state_dict(), model_prefix + '.pdopt')
WenmuZhou's avatar
WenmuZhou committed
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))