标签:resnet18 10 19 onnx torch state dict 2021 model
江大白共学计划课程–Pytorch课程作业(1)
import torch
import torchvision.models as models
from thop import profile
# 加载模型结构
model = models.resnet18()
# 读取权重和载入权重
pretrained_state_dict = torch.load('Lesson2/weights/resnet18-5c106cde.pth')
model_state_dict = model.load_state_dict(pretrained_state_dict, strict=False)
# 让模型变为推理状态
model.eval()
# 将模型放置到cuda上
model.to(torch.device('cuda'))
# 构建一个项目推理时需要的输入大小的单精度Tensor,并且放置模型所在的设备(cpu或cuda)
inputs = torch.ones([1, 3, 224, 224]).type(torch.float32).to(torch.device('cuda'))
flops, params = profile(model=model, inputs=(inputs))
print('Model:{:.2f} GFLOPs and {:.2f}M parameters'.format(flops/1e9, params/1e6))
# 生成onnx:将训练好的模型(包括结构和权重)保存成onnx。
torch.onnx.export(model, inputs, 'Lesson2/weights/resnet18.onnx', verbose = False)
输入:[1, 3, 224, 224]
输出:resnet18模型的参数量和计算量
Model:0.31 GFLOPs and 3.50M parameters
输入:[1, 3, 448, 448]
输出:resnet18模型的参数量和计算量
Model:1.25 GFLOPs and 3.50M parameters
Netron可视化ResNet18的网络结构
标签:resnet18,10,19,onnx,torch,state,dict,2021,model 来源: https://blog.csdn.net/weixin_42538848/article/details/120855207
本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享; 2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关; 3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关; 4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除; 5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。