gpt4 book ai didi

python - 如何将自定义 Pytorch 模型转换为 torchscript(pth 到 pt 模型)?

转载 作者:行者123 更新时间:2023-12-04 08:04:17 35 4
gpt4 key购买 nike

我使用 colab 环境使用 PyTorch 训练了一个自定义模型。我成功地将训练好的模型保存到 Google Drive,名称为 model_final.pth .我要转换 model_final.pthmodel_final.pt以便它可以在移动设备上使用。
我用来训练模型的代码如下:

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("mouse_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 1000
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.OUTPUT_DIR="drive/Detectron2/"

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
我用来转换模型的代码如下:
from detectron2.modeling import build_model
import torch
import torchvision

print("cfg.MODEL.WEIGHTS: ",cfg.MODEL.WEIGHTS) ## RETURNS : cfg.MODEL.WEIGHTS: drive/Detectron2/model_final.pth
model = build_model(cfg)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("drive/Detectron2/model-final.pt")
但我收到此错误 索引错误:维度 3 的张量索引过多 :
cfg.MODEL.WEIGHTS:  drive/Detectron2/model_final.pth
/usr/local/lib/python3.6/dist-packages/torch/tensor.py:593: RuntimeWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).
'incorrect results).', category=RuntimeWarning)
---------------------------------------------------------------------------

IndexError Traceback (most recent call last)
<ipython-input-17-8e544c0f39c8> in <module>()
7 model.eval()
8 example = torch.rand(1, 3, 224, 224)
----> 9 traced_script_module = torch.jit.trace(model, example)
10 traced_script_module.save("drive/Detectron2/model_final.pt")

7 frames
/usr/local/lib/python3.6/dist-packages/detectron2/modeling/meta_arch/rcnn.py in <listcomp>(.0)
219 Normalize, pad and batch the input images.
220 """
--> 221 images = [x["image"].to(self.device) for x in batched_inputs]
222 images = [(x - self.pixel_mean) / self.pixel_std for x in images]
223 images = ImageList.from_tensors(images, self.backbone.size_divisibility)

IndexError: too many indices for tensor of dimension 3

最佳答案

Detectron2 models expect a dictionary or a list of dictionaries as input默认情况下。
所以不能直接使用torch.jit.trace功能。但是他们提供了一个包装器,名为 TracingAdapter ,这允许模型将张量或张量元组作为输入。您可以在他们的 torchscript tests 中了解如何使用它。 .
跟踪您的 Mask RCNN 模型的代码可能是(我没有尝试过):

import torch
import torchvision
from detectron2.export.flatten import TracingAdapter

def inference_func(model, image):
inputs = [{"image": image}]
return model.inference(inputs, do_postprocess=False)[0]

print("cfg.MODEL.WEIGHTS: ",cfg.MODEL.WEIGHTS) ## RETURNS : cfg.MODEL.WEIGHTS: drive/Detectron2/model_final.pth
model = build_model(cfg)
example = torch.rand(1, 3, 224, 224)
wrapper = TracingAdapter(model, example, inference_func)
wrapper.eval()
traced_script_module = torch.jit.trace(wrapper, (example,))
traced_script_module.save("drive/Detectron2/model-final.pt")
可以在 here 中找到有关具有跟踪功能的检测器 2 部署的更多信息.

关于python - 如何将自定义 Pytorch 模型转换为 torchscript(pth 到 pt 模型)?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/66303312/

35 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com