Record the little trick when defining the model loading --state_dict or the strict function of torch.nn.Module.load_state_dict()

First briefly write the code to load the model:

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 variable to change the way of loading model parameters. Why does this problem occur? Because some networks will define many model function objects under class init, but it does not mean that the definition will be called. Assume that it is not called in forward. Then if strict is not false, the model will report an error when it encounters an unused key value. So after we fully understand the network, we can manually cancel this strict.

Guess you like

Origin blog.csdn.net/qq_40962125/article/details/134639968