Pytorch目标检测模型优化

2024-04-28 21:52:49 发布

您现在位置:Python中文网/ 问答频道 /正文

我想减小对象检测模型的大小。同样,我尝试使用pytorch mobile Optimizer优化更快的R-CNN模型进行目标检测,但生成的.pt{}文件大小与原始模型大小相同

我使用了下面提到的代码

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

model.eval()
script_model = torch.jit.script(model)
from torch.utils.mobile_optimizer import optimize_for_mobile
script_model_vulkan = optimize_for_mobile(script_model, backend='Vulkan')
torch.jit.save(script_model_vulkan, "frcnn.pth")

Tags: 对象from模型importformodelscriptutils
1条回答
网友
1楼 · 发布于 2024-04-28 21:52:49

您必须首先量化您的模型
遵循以下步骤here
&;然后使用这些方法

from torch.utils.mobile_optimizer import optimize_for_mobile
script_model_vulkan = optimize_for_mobile(script_model, backend='Vulkan')
torch.jit.save(script_model_vulkan, "frcnn.pth")

编辑:

resnet50模型的量化处理

import torchvision
model = torchvision.models.resnet50(pretrained=True)
import os
import torch

def print_model_size(mdl):
    torch.save(mdl.state_dict(), "tmp.pt")
    print("%.2f MB" %(os.path.getsize("tmp.pt")/1e6))
    os.remove('tmp.pt')
print_model_size(model) # will print original model size
backend = "qnnpack"
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend
model_static_quantized = torch.quantization.prepare(model, inplace=False)
model_static_quantized = torch.quantization.convert(model_static_quantized, inplace=False)



print_model_size(model_static_quantized) ## will print quantized model size

相关问题 更多 >