ParseQ模型PyTorch转CoreML时遇None类型转int错误求助

问题描述 #

在Colab环境中将ParseQ模型从PyTorch转换为CoreML格式时，出现错误：Error during Core ML conversion: None has type NoneType, but expected one of: int。模型本地运行正常，调试forward方法输出均为有效Tensor，但转换流程失败。使用的转换代码如下：

import torch
from PIL import Image
import coremltools as ct
from google.colab import drive
import pytorch_lightning as pl
from pytorch_lightning import Trainer

# Mount Google Drive
drive.mount("/content/drive", force_remount=True)

# Ensure the correct path to the local repository in Google Drive
local_repo_path = '/content/drive/My Drive/parseq'

# Initialize the Trainer
trainer = Trainer(accelerator='gpu' if torch.cuda.is_available() else 'cpu', devices=1)

# Load the model from the local repository
parseq = torch.hub.load(local_repo_path, 'parseq', source='local', pretrained=True)
parseq.trainer = trainer  # Attach the trainer to the model
parseq.eval()

# Get the image size from the model's parameters
img_size = [32, 128]  # Defined image size for the parseq model

# Example input tensor with the correct size
example_input = torch.rand(1, 3, img_size[0], img_size[1])

# Trace the model with example input
traced_model = torch.jit.trace(parseq, example_input)

# Debug: Check the output of the traced model
try:
    out = traced_model(example_input)
    print("Tracing successful. Output shape:", out.shape)
except Exception as e:
    print("Error during tracing:", e)

# Convert the traced model to Core ML
try:
    model = ct.convert(
        traced_model,
        convert_to="mlprogram",
        inputs=[ct.TensorType(shape=example_input.shape)],
        debug=True
    )
    print("Conversion to Core ML successful.")
except Exception as e:
    print("Error during Core ML conversion:", e)

# Define the path in Google Drive where you want to save the model
drive_path = '/content/drive/My Drive/new_parseq_model.mlpackage'

# Save the converted model to Google Drive
try:
    model.save(drive_path)
    print(f"Model saved to {drive_path}")
except Exception as e:
    print("Error saving the Core ML model:", e)

解决方案 #

• 移除不必要的Trainer绑定
  ParseQ作为PyTorch Lightning模型，转换CoreML时无需绑定Trainer实例。删除parseq.trainer = trainer这一行，避免模型内部引入无关属性干扰转换流程。

• 改用torch.jit.script替代trace
  动态结构的模型用trace容易丢失逻辑信息，改用script更可靠：

  # 替换原trace代码
  scripted_model = torch.jit.script(parseq)
  # 验证脚本模型输出
  out = scripted_model(example_input)
  print("Scripting successful. Output shape:", out.shape)
  

  后续CoreML转换使用scripted_model替代traced_model。

• 显式指定输入输出格式
  转换时明确输入类型（用ImageType贴合OCR场景）和输出类型，帮助CoreML工具正确解析：

  model = ct.convert(
      scripted_model,
      convert_to="mlprogram",
      inputs=[ct.ImageType(name="input_image", shape=example_input.shape, scale=1/255.0)],
      outputs=[ct.TensorType(name="logits")],
      debug=True
  )
  

  其中scale=1/255.0对应图像预处理的归一化操作，需和模型训练时的预处理逻辑保持一致。

• 覆盖模型forward方法的动态分支
  脚本化前用多个不同尺寸的输入测试模型，确保forward方法中所有动态分支都被执行到，避免脚本化时遗漏逻辑导致转换异常。

内容的提问来源于stack exchange，提问作者THESKRILL