提交 77ac6458 编写于 作者: W wuzewu

Show model name while loading pretrained model

上级 248d59c2
......@@ -50,8 +50,8 @@ def load_entire_model(model, pretrained):
if os.path.exists(pretrained):
load_pretrained_model(model, pretrained)
else:
raise Exception('Pretrained model is not found: {}'.format(
pretrained))
raise Exception(
'Pretrained model is not found: {}'.format(pretrained))
else:
logger.warning('Not all pretrained params of {} to load, '\
'training from scratch or a pretrained backbone'.format(model.__class__.__name__))
......@@ -95,15 +95,18 @@ def load_pretrained_model(model, pretrained_model):
model_state_dict[k] = para_state_dict[k]
num_params_loaded += 1
model.set_dict(model_state_dict)
logger.info("There are {}/{} variables are loaded.".format(
num_params_loaded, len(model_state_dict)))
logger.info("There are {}/{} variables are loaded into {}.".format(
num_params_loaded, len(model_state_dict),
model.__class__.__name__))
else:
raise ValueError(
'The pretrained model directory is not Found: {}'.format(
pretrained_model))
else:
logger.warning('No pretrained model to load, train from scratch')
logger.info(
'No pretrained model to load, {} will be train from scratch.'.
format(model.__class__.__name__))
def resume(model, optimizer, resume_model):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册