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

Show model name while loading pretrained model

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