从Pytorch模型pth文件中读取参数成numpy矩阵

目的:把训练好的pth模型参数提取出来,然后用其他方式部署到边缘设备。

Pytorch给了很方便的读取参数接口:nn.Module.parameters()

直接看demo:

from torchvision.models.alexnet import alexnet


model = alexnet(pretrained=True).eval().cuda()
parameters = model.parameters()
for p in parameters:
    numpy_para = p.detach().cpu().numpy()
    print(type(numpy_para))
    print(numpy_para.shape)

上面得到的numpy_para就是numpy参数了~

Note:model.parameters()是以一个生成器的形式迭代返回每一层的参数。所以用for循环读取到各层的参数,循环次数就表示层数。而每一层的参数都是torch.nn.parameter.Parameter类型,是Tensor的子类,所以直接用tensor转numpy(即p.detach().cpu().numpy())的方法就可以直接转成numpy矩阵。方便又好用,爆赞~

已标记关键词 清除标记
相关推荐
©️2020 CSDN 皮肤主题: 精致技术 设计师:CSDN官方博客 返回首页