pytorch模型参数C语言读取,pytorch模型加载与参数查看

本次主要对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

参考链接