Как преобразовать сегментацию Pytorch в CoreML ⇐ Python

[Anonymous]

Скрипт получает модель DeepLabV3 + MobileNet из pytorch.
Затем он сохраняет ее в .pt, который можно открыть в Netron.
Затем он преобразует ее в .mlpackage, который невозможно открыть в Netron - он выдает ошибку «Файл не имеет содержимого», хотя его размер около 22 МБ.
Имейте в виду, что ввод должен быть типом изображения, а не чем несколько массивов.
Я пробовал это со многими различными версиями torch и torchvision. Последняя версия torch, поддерживаемая Coreml, — 2.4, поэтому требуется более ранняя версия torchvision.
pip install "torch==2.4" "torchvision==0.19"
pip install coremltools
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import coremltools as ct

# Load the pretrained DeepLabV3 model with MobileNetV3 backbone
model = torch.hub.load('pytorch/vision:v0.19.0', 'deeplabv3_mobilenet_v3_large', pretrained=True)
model.eval()

# wrapper class
class Deeplabv3Wrapper(nn.Module):
def __init__(self, model):
super(Deeplabv3Wrapper, self).__init__()
self.model = model

def forward(self, x):
outputs = self.model(x)
return outputs['out'] # Return the 'out' tensor for segmentation mask

# Define preprocessing transforms
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # values from pytorch docs
])

dummy_image = Image.new("RGB", (513, 513), color="white")
dummy_tensor = preprocess(dummy_image).unsqueeze(0) # Add batch dimension

# Wrap and trace the model
wrapped_model = Deeplabv3Wrapper(model)

# Trace the model using TorchScript
traced_model = torch.jit.trace(wrapped_model, dummy_tensor)
print("Model tracing completed")

# Save the traced model
torch.jit.save(traced_model, "traced_deeplabv3.pt") # i CAN open this in netron.app
print("Model successfully saved as pt")

traced_model.eval()

example_input = torch.randn(1, 3, 513, 513)

model_from_trace = ct.convert(
traced_model,
inputs=[ct.ImageType(name="input_image", shape=example_input.shape, scale=1/255.0, bias=[-0.485, -0.456, -0.406])],
outputs=[ct.TensorType(name="output")]
)

print("Model conversion completed")

model_from_trace.save("DeepLabV3_MobileNet.mlpackage") # i CANNOT open this in netron.app - Error: "File has no content"

print("Model successfully saved as mlpackage")



Подробнее здесь: https://stackoverflow.com/questions/791 ... -to-coreml