本次主要对pytorch中的模型加载方法、各个层对应的名称及tensor值查看的方法做一个总结。
模型加载
import torch
# 模型文件位置
model_file = 'deeplab-mobilenet.pth.tar' # 或者.pth格式的模型文件
# 创建模型对象
model = MyDeepLab()
# 加载模型参数,若为cpu加载,则后面添加参数: map_location='cpu'
ckpt = torch.load(model_file) # cpu加载方式 ckpt = torch.load(model_file, map_location='cpu')
# 模型参数恢复
model.load_state_dict(ckpt) # 若要忽略名称不匹配的层:model.load_state_dict(ckpt, strict=False)
模型文件参数查看
# 使用上面代码块中的ckpt,其实质上是一个dict, key是layer名称, value是对应的tensor值
# 获取所有layer名称
layer_name = list(ckpt.keys()) # 获取预训练模型各层的名称,并转为list
# 查看指定layer的tensor值
print(ckpt[ layer_name[2]) # 查看第2个层的参数值
模型对象参数查看
# 这里与模型文件参数的区别在于:这是针对代码中创建的模型对象,查看其各个layer的名称与tensor值
# 获取模型中所有layer的名称
layer_name = list(model.state_dict().keys())
# 查看指定layer的tensor值
print( model.state_dict()[ layer_name[2] ])
模型文件中layer名称与模型对象中layer名称的区别
主要区别在于模型对象中的名称最开始会多一个module.
例如:
模型文件中layer名称: conv5.weight
模型对象中layer名称: module.conv5.weight