utility.py 4.2 KB
Newer Older
S
shaohua.zhang 已提交
1
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
L
LDOUBLEV 已提交
2
#
S
shaohua.zhang 已提交
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
#
S
shaohua.zhang 已提交
7
#    http://www.apache.org/licenses/LICENSE-2.0
L
LDOUBLEV 已提交
8
#
S
shaohua.zhang 已提交
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

S
shaohua.zhang 已提交
15 16 17
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
L
LDOUBLEV 已提交
18

S
shaohua.zhang 已提交
19 20 21 22 23
import errno
import os
import shutil
import tempfile

S
shaohua.zhang 已提交
24
import paddle
S
shaohua.zhang 已提交
25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
import paddle.fluid as fluid

from .utility import initial_logger
import re
logger = initial_logger()


def _mkdir_if_not_exist(path):
    """
    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_state(path):
    if os.path.exists(path + '.pdopt'):
        # XXX another hack to ignore the optimizer state
        tmp = tempfile.mkdtemp()
        dst = os.path.join(tmp, os.path.basename(os.path.normpath(path)))
        shutil.copy(path + '.pdparams', dst + '.pdparams')
        state = fluid.io.load_program_state(dst)
        shutil.rmtree(tmp)
    else:
        state = fluid.io.load_program_state(path)
    return state


def load_params(exe, prog, path, ignore_params=[]):
    """
    Load model from the given path.
    Args:
        exe (fluid.Executor): The fluid.Executor object.
        prog (fluid.Program): load weight to which Program object.
        path (string): URL string or loca model path.
        ignore_params (list): ignore variable to load when finetuning.
            It can be specified by finetune_exclude_pretrained_params
            and the usage can refer to docs/advanced_tutorials/TRANSFER_LEARNING.md
    """
    if not (os.path.isdir(path) or os.path.exists(path + '.pdparams')):
        raise ValueError("Model pretrain path {} does not "
                         "exists.".format(path))

    logger.info('Loading parameters from {}...'.format(path))

    ignore_set = set()
    state = _load_state(path)

    # ignore the parameter which mismatch the shape
    # between the model and pretrain weight.
    all_var_shape = {}
    for block in prog.blocks:
        for param in block.all_parameters():
            all_var_shape[param.name] = param.shape
    ignore_set.update([
        name for name, shape in all_var_shape.items()
        if name in state and shape != state[name].shape
    ])

    if ignore_params:
        all_var_names = [var.name for var in prog.list_vars()]
        ignore_list = filter(
            lambda var: any([re.match(name, var) for name in ignore_params]),
            all_var_names)
        ignore_set.update(list(ignore_list))

    if len(ignore_set) > 0:
        for k in ignore_set:
            if k in state:
                logger.warning('variable {} not used'.format(k))
                del state[k]
    fluid.io.set_program_state(prog, state)


def init_model(config, program, exe):
    """
    load model from checkpoint or pretrained_model
    """
    checkpoints = config['Global'].get('checkpoints')
    if checkpoints:
        path = checkpoints
        fluid.load(program, path, exe)
        logger.info("Finish initing model from {}".format(path))
S
shaohua.zhang 已提交
116
        return
S
shaohua.zhang 已提交
117 118 119 120 121 122

    pretrain_weights = config['Global'].get('pretrain_weights')
    if pretrain_weights:
        path = pretrain_weights
        load_params(exe, program, path)
        logger.info("Finish initing model from {}".format(path))
S
shaohua.zhang 已提交
123
    return
S
shaohua.zhang 已提交
124 125 126 127 128 129 130 131


def save_model(program, model_path):
    """
    save model to the target path
    """
    fluid.save(program, model_path)
    logger.info("Already save model in {}".format(model_path))