Open the-neural-networker opened 11 months ago
We only have "expiremental" support for PyTorch models which have not been created by torch.jit.trace
.
The same issue persists if the PyTorch model Experiment
is traced, but the function experiment
is scripted. Does this still fall under experimental support?
@the-neural-networker - the following works for me:
import numpy as np
import torch
import torch.nn as nn
import coremltools as ct
def experiment(val: torch.Tensor):
val = int(val.item())
result = torch.zeros(val)
for i in range(val):
result[i] = i
return result
class Experiment(nn.Module):
def forward(self, x):
return experiment(x)
exp = Experiment().eval()
# Use a tensor as input, not an integer
input_tensor = torch.tensor(100)
traced_model = torch.jit.trace(exp, input_tensor)
y_t = traced_model(input_tensor)
# Specify the input type as ct.TensorType(name="x", shape=(1,))
coreml_exp = ct.convert(
traced_model,
source="pytorch",
inputs=[ct.TensorType(name="x", shape=(1,))],
convert_to="mlprogram"
)
# Create an input dictionary with the necessary input data
input_data = {
'x': np.array([100.0]),
}
# Make a prediction using the model
coreml_output = coreml_exp.predict(input_data)
coreml_output = list(coreml_output.values())[0]
assert all(coreml_output == y_t.numpy())
But won't this fail for other inputs (not 100)? Because technically the experiment function is traced which has a for
loop and the function's for
loop depends on val
, it is not static.
This example will give the wrong prediction when the input is not 100
. However that is a result of the tracing not the conversion to Core ML. The PyTorch traced model gives the wrong the results. The model is converted to Core ML correctly.
If you want to do something like this in Core ML, I think you'll need to write your own mlprogram. Take a look at our set of MIL ops:
https://apple.github.io/coremltools/source/coremltools.converters.mil.mil.ops.defs.html
Specifically, I think you'll want to look at the ones under:
coremltools.converters.mil.mil.ops.defs.iOS15.control_flow.*
πDescription
When torchscript code contains for loops and is converted to coreML, the for loops are skipped and the result before the for loop is returned.
To Reproduce
System environment (please complete the following information):
Additional Context
The for loop can be easily vectorized, but for simplicity, I wanted to show that the for loops get skipped in coreML conversion.