Python 如何将自定义Pytorch模型转换为torchscript(pth到pt模型)?
我使用colab环境使用pytorch训练了一个定制模型。我成功地将经过培训的模型保存到名为model_final.pth的Google Drive。我想将model_final.pth转换为model_final.pt,以便可以在移动设备上使用。 我用来训练模型的代码如下:Python 如何将自定义Pytorch模型转换为torchscript(pth到pt模型)?,python,pytorch,Python,Pytorch,我使用colab环境使用pytorch训练了一个定制模型。我成功地将经过培训的模型保存到名为model_final.pth的Google Drive。我想将model_final.pth转换为model_final.pt,以便可以在移动设备上使用。 我用来训练模型的代码如下: cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_F
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()
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")
我用来转换模型的代码如下
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()
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")
但我得到了这个错误
/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
/usr/local/lib/python3.6/dist-packages/torch/tensor.py:593:RuntimeWarning:在tensor上迭代可能会导致跟踪不正确。传递不同形状的张量不会改变执行的迭代次数(并且可能会导致错误或给出错误的结果)。
“结果不正确)。”,类别=运行时警告)
---------------------------------------------------------------------------
索引器回溯(最后一次最近调用)
在()
7.模型评估()
8示例=torch.rand(1,3,224,224)
---->9跟踪的脚本模块=torch.jit.trace(模型,示例)
10跟踪的脚本模块保存(“驱动器/Detectron2/model\u final.pt”)
7帧
/usr/local/lib/python3.6/dist-packages/detectron2/modeling/meta_arch/rcnn.py in(.0)
219对输入图像进行规格化、填充和批处理。
220 """
-->221图像=[x[“图像”].到(self.device),用于批处理_输入中的x]
222个图像=[(x-self.pixel_平均值)/self.pixel_标准表示图像中的x]
223 images=ImageList.from_张量(images,self.backbone.size_整除性)
索引器:维度3的张量的索引太多
默认情况下,
因此,您不能直接使用torch.jit.trace
函数。但它们提供了一个名为的包装器,允许模型将张量或张量元组作为输入。您可以了解如何在模型中使用它
跟踪掩码RCNN模型的代码可能是(我没有尝试):
可以找到更多关于detectron2部署和跟踪的信息