CoreML: создание пользовательского слоя для ONNX RandomNormal

Question

Я обучил VAE, что в PyTorch мне нужно конвертировать в CoreML. Из этого потока PyTorch VAE не удается преобразовать в onnx. Я смог получить модель 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. Поэтому мне интересно, как указать тип "фиктивного" ввода для слоя, так как на самом деле он не имеет ввода - это просто оболочка вокруг torch.randn() (или, более конкретно, onnx RandonNormal оп). Я должен уточнить, что мне нужен весь VAE, а не только декодер, так как я на самом деле использую весь процесс для "исправления ошибок" моих входов (т. Е. Кодер оценивает мои z вектор, на основе входных данных, затем декодер генерирует наиболее близкое обобщаемое предсказание входных данных).

Любая помощь с благодарностью.

ОБНОВЛЕНИЕ: Хорошо, я наконец-то получил версию для загрузки в Xcode (спасибо @MattijsHollemans и его книге!). originalConversion.mlmodel является исходным результатом преобразования моей модели из ONNX в CoreML. Для этого мне пришлось вручную вставить вход для RandomNormal слой. Я сделал это (64, 28, 28) без особой причины - я знаю, что мой размер пакета равен 64, и мои входные данные равны 28 x 28 (но, вероятно, это также может быть (1, 1, 1), так как это "пустышка" "):

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 был доволен), я не думаю, что это будет проблемой, но я опубликую еще раз, если он не запустится должным образом.

ОБНОВЛЕНИЕ 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
}

У меня есть два вопроса / проблемы: 1) Очевидно, почему все значения равны нулю (это имеет место для всех, кроме входных узлов), и 2) Почему это выглядит как последовательная модель, когда это просто довольно условно VAE? открытие model.espresso.shape для полнофункциональной GAN в том же приложении я вижу, что узлы имеют формат:

"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

Answer 1

Чтобы добавить вход в вашу модель 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,
• Я сделал новый фиктивный ввод типа "двойной". Это означает, что ваш пользовательский слой получает двойное значение в качестве входных данных.
• Вам необходимо указать значение для этого фиктивного ввода при использовании модели.