"...onnxruntime/textual_inversion/textual_inversion.py" did not exist on "04ad9486738156d6cf5aed4cf2397fc356dbaeaa"
save_load.py 5.73 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

WenmuZhou's avatar
WenmuZhou committed
26
__all__ = ['init_model', 'save_model', 'load_dygraph_pretrain']
LDOUBLEV's avatar
LDOUBLEV committed
27
28


WenmuZhou's avatar
WenmuZhou committed
29
def _mkdir_if_not_exist(path, logger):
LDOUBLEV's avatar
LDOUBLEV committed
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
    """
    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))


WenmuZhou's avatar
WenmuZhou committed
45
46
47
48
49
def load_dygraph_pretrain(
        model,
        logger,
        path=None,
        load_static_weights=False, ):
LDOUBLEV's avatar
LDOUBLEV committed
50
51
52
    if not (os.path.isdir(path) or os.path.exists(path + '.pdparams')):
        raise ValueError("Model pretrain path {} does not "
                         "exists.".format(path))
WenmuZhou's avatar
WenmuZhou committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
    if load_static_weights:
        pre_state_dict = paddle.io.load_program_state(path)
        param_state_dict = {}
        model_dict = model.state_dict()
        for key in model_dict.keys():
            weight_name = model_dict[key].name
            weight_name = weight_name.replace('binarize', '').replace(
                'thresh', '')  # for DB
            if weight_name in pre_state_dict.keys():
                logger.info('Load weight: {}, shape: {}'.format(
                    weight_name, pre_state_dict[weight_name].shape))
                if 'encoder_rnn' in key:
                    # delete axis which is 1
                    pre_state_dict[weight_name] = pre_state_dict[
                        weight_name].squeeze()
                    # change axis
                    if len(pre_state_dict[weight_name].shape) > 1:
                        pre_state_dict[weight_name] = pre_state_dict[
                            weight_name].transpose((1, 0))
                param_state_dict[key] = pre_state_dict[weight_name]
            else:
                param_state_dict[key] = model_dict[key]
        model.set_dict(param_state_dict)
        return

    param_state_dict, optim_state_dict = paddle.load(path)
    model.set_dict(param_state_dict)
    return
LDOUBLEV's avatar
LDOUBLEV committed
81

WenmuZhou's avatar
WenmuZhou committed
82
83

def init_model(config, model, logger, optimizer=None, lr_scheduler=None):
LDOUBLEV's avatar
LDOUBLEV committed
84
85
86
    """
    load model from checkpoint or pretrained_model
    """
WenmuZhou's avatar
WenmuZhou committed
87
88
89
90
    gloabl_config = config['Global']
    checkpoints = gloabl_config.get('checkpoints')
    pretrained_model = gloabl_config.get('pretrained_model')
    best_model_dict = {}
LDOUBLEV's avatar
LDOUBLEV committed
91
    if checkpoints:
WenmuZhou's avatar
WenmuZhou committed
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
120
121
122
123
124
125
126
127
        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)
        para_dict, opti_dict = paddle.load(checkpoints)
        model.set_dict(para_dict)
        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
            best_model_dict['start_epoch'] = best_model_dict['best_epoch'] + 1

        logger.info("resume from {}".format(checkpoints))
    elif pretrained_model:
        load_static_weights = gloabl_config.get('load_static_weights', False)
        if pretrained_model:
            if not isinstance(pretrained_model, list):
                pretrained_model = [pretrained_model]
            if not isinstance(load_static_weights, list):
                load_static_weights = [load_static_weights] * len(
                    pretrained_model)
            for idx, pretrained in enumerate(pretrained_model):
                load_static = load_static_weights[idx]
                load_dygraph_pretrain(
                    model,
                    logger,
                    path=pretrained,
                    load_static_weights=load_static)
                logger.info("load pretrained model from {}".format(
                    pretrained_model))
128
    else:
WenmuZhou's avatar
WenmuZhou committed
129
130
        logger.info('train from scratch')
    return best_model_dict
LDOUBLEV's avatar
LDOUBLEV committed
131
132


WenmuZhou's avatar
WenmuZhou committed
133
134
135
136
137
138
139
def save_model(net,
               optimizer,
               model_path,
               logger,
               is_best=False,
               prefix='ppocr',
               **kwargs):
LDOUBLEV's avatar
LDOUBLEV committed
140
141
142
    """
    save model to the target path
    """
WenmuZhou's avatar
WenmuZhou committed
143
144
145
146
147
148
149
150
151
152
153
154
    _mkdir_if_not_exist(model_path, logger)
    model_prefix = os.path.join(model_path, prefix)
    paddle.save(net.state_dict(), model_prefix)
    paddle.save(optimizer.state_dict(), model_prefix)

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