<reponame>spxia/wenet<filename>wenet/utils/checkpoint.py<gh_stars>0
# Copyright 2019 Mobvoi Inc. All Rights Reserved.
# Author: <EMAIL> (<NAME>)

import logging
import os
import re

import yaml
import torch

def transfer_state_dict(pretrained_dict, model_dict):
    #fil_key = ['decoder.embed.0.weight', 'decoder.output_layer.weight', 'decoder.output_layer.bias', 'ctc.ctc_lo.weight', 'ctc.ctc_lo.bias']
    #fil_key = ['decoder','ctc']

    fil_key = [] #Set to null, do not migrate parameter
    state_dict = {}
    for k, v in pretrained_dict.items():
        # if k in model_dict.keys():
        #     state_dict[k] = v
        # else:
        #     print("Missing key(s) in state_dict :{}".format(k))
        # Specify layers not to load

        if k not in fil_key:
            # state_dict[k] = v
        #if k.split('.')[0] not in fil_key:
            state_dict[k] = v
        else:
            print("filter key(s) in state_dict :{}".format(k))
    return state_dict


def load_checkpoint(model: torch.nn.Module, path: str) -> dict:
    if torch.cuda.is_available():
        logging.info('Checkpoint: loading from checkpoint %s for GPU' % path)
        checkpoint = torch.load(path)
    else:
        logging.info('Checkpoint: loading from checkpoint %s for CPU' % path)
        checkpoint = torch.load(path, map_location='cpu')
    # Modify the nodes of the output layer---------
    model_dict = model.state_dict()
    pretrained_dict = transfer_state_dict(checkpoint, model_dict)
    model_dict.update(pretrained_dict)
    checkpoint = model_dict
    #--------
    model.load_state_dict(checkpoint)
    info_path = re.sub('.pt$', '.yaml', path)
    configs = {}
    if os.path.exists(info_path):
        with open(info_path, 'r') as fin:
            configs = yaml.load(fin, Loader=yaml.FullLoader)
    return configs


def save_checkpoint(model: torch.nn.Module, path: str, infos=None):
    '''
    Args:
        infos (dict or None): any info you want to save.
    '''
    logging.info('Checkpoint: save to checkpoint %s' % path)
    if isinstance(model, torch.nn.DataParallel):
        state_dict = model.module.state_dict()
    elif isinstance(model, torch.nn.parallel.DistributedDataParallel):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()
    torch.save(state_dict, path)
    info_path = re.sub('.pt$', '.yaml', path)
    if infos is None:
        infos = {}
    with open(info_path, 'w') as fout:
        data = yaml.dump(infos)
        fout.write(data)
