CoreML:为 ONNX RandomNormal 创建自定义层

CoreML: creating a custom layer for ONNX RandomNormal

我已经在 PyTorch 中训练了一个 VAE,我需要将其转换为 CoreML。从这个线程 我能够导出 ONNX 模型,但是,这只是将问题进一步推向了 ONNX-CoreML 阶段。

包含 torch.randn() 调用的原始函数是重新参数化函数:

def reparametrize(self, mu, logvar):
    std = logvar.mul(0.5).exp_()
    if self.have_cuda:
        eps = torch.randn(self.bs, self.nz, device='cuda')
    else:
        eps = torch.randn(self.bs, self.nz)
    return eps.mul(std).add_(mu)

当然,解决方案是创建一个自定义层,但我在创建没有输入的层时遇到了问题(即,它只是一个 randn() 调用)。

我可以使用此定义完成 CoreML 转换:

def convert_randn(node):
    params = NeuralNetwork_pb2.CustomLayerParams()
    params.className = "RandomNormal"
    params.description = "Random normal distribution generator"
    params.parameters["dtype"].intValue = node.attrs.get('dtype', 1)
    params.parameters["bs"].intValue = node.attrs.get("shape")[0]
    params.parameters["nz"].intValue = node.attrs.get("shape")[1]
    return params

我用以下方法进行转换:

coreml_model = convert(onnx_model, add_custom_layers=True, 
    image_input_names = ['input'], 
    custom_conversion_functions={"RandomNormal": convert_randn})

我还应该注意到,在 mlmodel 导出完成时,会打印以下内容:

Custom layers have been added to the CoreML model corresponding to the 
following ops in the onnx model: 
1/1: op type: RandomNormal, op input names and shapes: [], op output     
names and shapes: [('62', 'Shape not available')]

.mlmodel 带入 Xcode 会抱怨 Layer '62' of type 500 has 0 inputs but expects at least 1. 所以我想知道如何为层指定一种 "dummy" 输入,因为它没有实际上没有输入——它只是 torch.randn() 的包装器(或者更具体地说,onnx RandonNormal 操作)。我应该澄清一下,我确实需要整个 VAE,而不仅仅是解码器,因为我实际上正在使用整个过程来 "error correct" 我的输入(即编码器估计我的 z 向量,基于输入,然后解码器生成输入的最接近的可泛化预测。

非常感谢任何帮助。

更新: 好的,我终于在 Xcode 中加载了一个版本(感谢@MattijsHollemans 和他的书!)。 originalConversion.mlmodel 是将我的模型从 ONNX 转换为 CoreML 的初始输出。为此,我必须手动插入 RandomNormal 层的输入。我无缘无故地做到了 (64, 28, 28) — 我知道我的批量大小是 64,我的输入是 28 x 28(但大概它也可能是 (1, 1, 1),因为它是 "dummy"):

spec = coremltools.utils.load_spec('originalConversion.mlmodel')
nn = spec.neuralNetwork
layers = {l.name:i for i,l in enumerate(nn.layers)}
layer_idx = layers["62"] # '62' is the name of the layer -- see above
layer = nn.layers[layer_idx]
layer.input.extend(["dummy_input"])

inp = spec.description.input.add()
inp.name = "dummy_input"
inp.type.multiArrayType.SetInParent()
spec.description.input[1].type.multiArrayType.shape.append(64)
spec.description.input[1].type.multiArrayType.shape.append(28)
spec.description.input[1].type.multiArrayType.shape.append(28)
spec.description.input[1].type.multiArrayType.dataType = ft.ArrayFeatureType.DOUBLE

coremltools.utils.save_spec(spec, "modelWithInsertedInput.mlmodel") 

这会在 Xcode 中加载,但我尚未在我的应用程序中测试该模型的功能。由于附加层很简单,而且输入实际上是伪造的、无功能的输入(只是为了 Xcode 高兴),我不认为这会是个问题,但我会 post 再次 运行 正确。

更新 2: 不幸的是,模型在 运行 时没有加载。它失败了[espresso] [Espresso::handle_ex_plan] exception=Failed in 2nd reshape after missing custom layer info.我发现非常奇怪和令人困惑的是,检查model.espresso.shape,我看到几乎每个节点都有这样的形状:

"62" : {
  "k" : 0,
  "w" : 0,
  "n" : 0,
  "seq" : 0,
  "h" : 0
}

我有两个 question/concerns:1) 最明显的是,为什么所有的值都是零(除了输入节点之外的所有值都是这种情况),以及 2) 为什么它看起来是一个顺序模型,当它只是一个相当传统的 VAE 时?在同一应用程序中打开功能齐全的 GAN model.espresso.shape,我看到节点的格式为:

"54" : {
  "k" : 256,
  "w" : 16,
  "n" : 1,
  "h" : 16
}

也就是说,它们包含合理的形状信息,并且它们没有 seq 个字段。

非常非常困惑...

更新 3: 我也刚刚注意到编译器报告错误:IMPORTANT: new sequence length computation failed, falling back to old path. Your compilation was sucessful, but please file a radar on Core ML | Neural Networks and attach the model that generated this message.

这是原始的 PyTorch 模型:

class VAE(nn.Module):
def __init__(self, bs, nz):
    super(VAE, self).__init__()

    self.nz = nz
    self.bs = bs

    self.encoder = nn.Sequential(
        # input is (nc) x 28 x 28
        nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        # size = (ndf) x 14 x 14
        nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
        nn.BatchNorm2d(ndf * 2),
        nn.LeakyReLU(0.2, inplace=True),
        # size = (ndf*2) x 7 x 7
        nn.Conv2d(ndf * 2, ndf * 4, 3, 2, 1, bias=False),
        nn.BatchNorm2d(ndf * 4),
        nn.LeakyReLU(0.2, inplace=True),
        # size = (ndf*4) x 4 x 4
        nn.Conv2d(ndf * 4, 1024, 4, 1, 0, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
    )

    self.decoder = nn.Sequential(
        # input is Z, going into a convolution
        nn.ConvTranspose2d(     1024, ngf * 8, 4, 1, 0, bias=False),
        nn.BatchNorm2d(ngf * 8),
        nn.ReLU(True),
        # size = (ngf*8) x 4 x 4
        nn.ConvTranspose2d(ngf * 8, ngf * 4, 3, 2, 1, bias=False),
        nn.BatchNorm2d(ngf * 4),
        nn.ReLU(True),
        # size = (ngf*4) x 8 x 8
        nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
        nn.BatchNorm2d(ngf * 2),
        nn.ReLU(True),
        # size = (ngf*2) x 16 x 16
        nn.ConvTranspose2d(ngf * 2,     nc, 4, 2, 1, bias=False),
        nn.Sigmoid()
    )

    self.fc1 = nn.Linear(1024, 512)
    self.fc21 = nn.Linear(512, nz)
    self.fc22 = nn.Linear(512, nz)

    self.fc3 = nn.Linear(nz, 512)
    self.fc4 = nn.Linear(512, 1024)

    self.lrelu = nn.LeakyReLU()
    self.relu = nn.ReLU()

def encode(self, x):
    conv = self.encoder(x);
    h1 = self.fc1(conv.view(-1, 1024))
    return self.fc21(h1), self.fc22(h1)

def decode(self, z):
    h3 = self.relu(self.fc3(z))
    deconv_input = self.fc4(h3)
    deconv_input = deconv_input.view(-1,1024,1,1)
    return self.decoder(deconv_input)

def reparametrize(self, mu, logvar):
    std = logvar.mul(0.5).exp_()
    eps = torch.randn(self.bs, self.nz, device='cuda') # needs custom layer!
    return eps.mul(std).add_(mu)

def forward(self, x):
    # print("x", x.size())
    mu, logvar = self.encode(x)
    z = self.reparametrize(mu, logvar)
    decoded = self.decode(z)
    return decoded, mu, logvar

要向您的 Core ML 模型添加输入,您可以从 Python 执行以下操作:

import coremltools
spec = coremltools.utils.load_spec("YourModel.mlmodel")

nn = spec.neuralNetworkClassifier  # or just spec.neuralNetwork

layers = {l.name:i for i,l in enumerate(nn.layers)}
layer_idx = layers["your_custom_layer"]
layer = nn.layers[layer_idx]
layer.input.extend(["dummy_input"])

inp = spec.description.input.add()
inp.name = "dummy_input"
inp.type.doubleType.SetInParent()

coremltools.utils.save_spec(spec, "NewModel.mlmodel")

此处,"your_custom_layer" 是您要添加虚拟输入的层的名称。在您的模型中,它看起来像是 62。您可以查看 layers 字典以查看模型中所有层的名称。

备注:

  • 如果您的模型不是分类器,请使用 nn = spec.neuralNetwork 而不是 neuralNetworkClassifier
  • 我将新的虚拟输入设置为 "double" 类型。这意味着您的自定义图层将获得双精度值作为输入。
  • 您需要在使用模型时为此虚拟输入指定一个值。