下面的代码包含用途有:
1.训练时多GPU,推理时所有层多出一个module时替换;
2.训练模型出现层的定义不一致时替换;
3.打印训练过程中学习的参数,可视化对应参数的值。
import torch
from collections import OrderedDict
from your_model import Net
# your net architecture
net = Net()
model_path = "your_model_path"
# load model parameters
state_dict = torch.load(model_path, map_location="cpu")
# define a new dict
new_state_dict = OrderedDict()
for k,v in state_dict.items():
# if you train model with mutil gpu
if "module" in k:
name = k[7:]
else:
name = k
# if some layer's name need to change
if "fc" in name:
name = name.replace("fc", "classifier")
# print your layer params
print("layer : {}, value : {}\n\n".format(name, v))
# store the value to new_state_dict
new_state_dict[name] = v
# load model param to model architecture
net.load_state_dict(new_state_dict)