Paddle.fluid之模型freeze

model.freeze

#-*- coding:utf-8 -*-
import os
import paddle 
import paddle.fluid as fluid
from model_zoos import antisp as model


def gen_exe(resume = '', use_cuda = False):

    place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
    exe = fluid.Executor(place)
    exe.run(fluid.default_startup_program())

    if resume:
        def if_exist(var):
            has_var = os.path.exists(os.path.join(resume, var.name))
            if has_var:
                print('var: %s found' % (var.name))
                return has_var
        fluid.io.load_vars(exe, resume, predicate=if_exist)
        print('init weight finished!')
    return exe


def freeze_model(exe, save_dir, tvars):

    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    
    if save_dir is not None:
        fluid.io.save_inference_model(
            dirname = save_dir,
            feeded_var_names = ['image'], 
            target_vars = [tvars], 
            executor = exe,
            model_filename = 'model', 
            params_filename = 'params')
    print('paddle inference model saved')


if __name__ == "__main__":
    
    image = fluid.layers.data(name='image', shape=[1, 3, img_h, img_w], dtype='float32', append_batch_size=False)
    output, _, _ = model(image, True, 1.0)
    resume = 'checkpoints/param_weights/iter_xxx_x_eppch'
    save_dir = 'checkpoints/freeze_model'
    exe = gen_exe(resume)
    freeze_model(exe, save_dir, output)