记录一下模型加载定义时的小trick --state_dict或torch.nn.Module.load_state_dict()的strict函数

先简单写一下加载模型的代码:

model = net()

model_data = torch.load("a.pth",cpu())

model.eval()

model.load_state_dict(model_data,strict = False)

eval():stop updata the BN or dropout layer direct use the mean,var for infenece.

load_state_dict(model,strict)

strict 变量 去改变加载模型参数时的方式,为什么出现这个问题,因为有的网络在class init下会定义很多模型函数的对象,但是并不代表定义就一定被调用,假设在forward中没有被调用,那么如果strict不是false的话,模型遇到没有使用的key值就会报错。所以在我们完全理解网络之后,可以手动去取消这个strict。

猜你喜欢

转载自blog.csdn.net/qq_40962125/article/details/134639968
今日推荐