Поделиться через

Преобразование модели PyTorch в формат ONNX

На предыдущем этапе этого руководства мы использовали PyTorch для создания модели машинного обучения. Однако эта модель является файлом .pth . Чтобы интегрировать его с приложением Windows ML, необходимо преобразовать модель в формат ONNX.

Экспорт модели

Для экспорта модели вы будете использовать функцию torch.onnx.export() . Эта функция выполняет модель и записывает трассировку того, какие операторы используются для вычисления выходных данных.

  1. Скопируйте следующий код в DataClassifier.py файл в Visual Studio над основной функцией.
#Function to Convert to ONNX 
def convert(): 

    # set the model to inference mode 
    model.eval() 

    # Let's create a dummy input tensor  
    dummy_input = torch.randn(1, 3, 32, 32, requires_grad=True)  

    # Export the model   
    torch.onnx.export(model,         # model being run 
         dummy_input,       # model input (or a tuple for multiple inputs) 
         "Network.onnx",       # where to save the model  
         export_params=True,  # store the trained parameter weights inside the model file 
         opset_version=11,    # the ONNX version to export the model to 
         do_constant_folding=True,  # whether to execute constant folding for optimization 
         input_names = ['input'],   # the model's input names 
         output_names = ['output'], # the model's output names 
         dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes 
                                'output' : {0 : 'batch_size'}}) 
    print(" ") 
    print('Model has been converted to ONNX')

Важно вызвать model.eval() или model.train(False) перед экспортом модели, так как это переводит модель в режим вывода заключений. Это необходимо, поскольку операторы, такие как dropout или batchnorm, ведут себя по-разному в режиме вывода и обучения.

  1. Чтобы запустить преобразование в ONNX, добавьте вызов функции преобразования в основную функцию. Вам не нужно снова обучить модель, поэтому мы закомментируем некоторые функции, которые нам больше не нужно запускать. Ваша основная функция заключается в следующем.
if __name__ == "__main__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert()
  1. Снова запустите проект, нажав Start Debugging кнопку на панели инструментов или нажав клавишу F5. Нет необходимости повторно обучить модель, просто загрузить существующую модель из папки проекта.

Перейдите в расположение проекта и найдите модель ONNX рядом с моделью .pth .

Замечание

Заинтересованы в обучении больше? Ознакомьтесь с руководством по экспорту mdoel в PyTorch.

Проанализируйте модель.

  1. Откройте файл модели с помощью Neutron.

  2. Выберите узел данных , чтобы открыть свойства модели.

Свойства модели ONNX

Как видно, для модели требуется 32-разрядный объект с плавающей запятой (многомерный массив) в качестве входных данных и возвращает значение Tensor float в виде выходных данных. Выходной массив будет включать вероятность для каждой метки. То, как вы создали модель, метки представлены 3 числами, каждый из которых связан с определенным типом цветка ириса.

Метка 1 метка 2 Метка 3
0 1 2
Ирис-сетоза Ирис многоцветковый (Iris-versicolor) Ирис виргинский

Эти значения необходимо извлечь, чтобы отобразить правильный прогноз с помощью приложения Windows ML.

Дальнейшие шаги

Наша модель готова к развертыванию. Затем для основного события давайте создадим приложение Windows и запустите его локально на устройстве Windows.