In macOS 15 beta the gridsample function from PyTorch is not executing as expected on the Apple Neural Engine in MacBook Pro M2.
Please find below a Python code snippet that demonstrates the problem:
import coremltools as ct
import torch.nn as nn
import torch.nn.functional as F
class PytorchGridSample(torch.nn.Module):
def __init__(self, grids):
super(PytorchGridSample, self).__init__()
self.upsample1 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
self.upsample2 = nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1)
self.upsample3 = nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)
self.upsample4 = nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
self.upsample5 = nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1)
self.grids = grids
def forward(self, x):
x = self.upsample1(x)
x = F.grid_sample(x, self.grids[0], padding_mode='reflection', align_corners=False)
x = self.upsample2(x)
x = F.grid_sample(x, self.grids[1], padding_mode='reflection', align_corners=False)
x = self.upsample3(x)
x = F.grid_sample(x, self.grids[2], padding_mode='reflection', align_corners=False)
x = self.upsample4(x)
x = F.grid_sample(x, self.grids[3], padding_mode='reflection', align_corners=False)
x = self.upsample5(x)
x = F.grid_sample(x, self.grids[4], padding_mode='reflection', align_corners=False)
return x
def convert_to_coreml(model, input_):
traced_model = torch.jit.trace(model, example_inputs=input_, strict=False)
coreml_model = ct.converters.convert(
traced_model,
inputs=[ct.TensorType(shape=input_.shape)],
compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.macOS14,
compute_units=ct.ComputeUnit.ALL
)
return coreml_model
def main(pt_model, input_):
coreml_model = convert_to_coreml(pt_model, input_)
coreml_model.save("grid_sample.mlpackage")
if __name__ == "__main__":
input_tensor = torch.randn(1, 512, 4, 4)
grids = [torch.randn(1, 2*i, 2*i, 2) for i in [4, 8, 16, 32, 64, 128]]
pt_model = PytorchGridSample(grids)
main(pt_model, input_tensor)