torch.onnx

    这是一个简单的脚本,将torchvision中定义的预训练的AlexNet导出到ONNX中。它运行一轮推理,然后将结果跟踪模型保存到alexnet.proto

    保存文件alexnet.proto是一个二进制protobuf文件,其中包含您导出的模型(在本例中为AlexNet)的网络结构和参数。关键字参数verbose=True导致导出器打印出一个人类可读的网络表示:

    1. # All parameters are encoded explicitly as inputs. By convention,
    2. # learned parameters (ala nn.Module.state_dict) are first, and the
    3. # actual inputs are last.
    4. graph(%1 : Float(64, 3, 11, 11)
    5. %2 : Float(64)
    6. # The definition sites of all variables are annotated with type
    7. # information, specifying the type and size of tensors.
    8. # For example, %3 is a 192 x 64 x 5 x 5 tensor of floats.
    9. %3 : Float(192, 64, 5, 5)
    10. %4 : Float(192)
    11. # ---- omitted for brevity ----
    12. %15 : Float(1000, 4096)
    13. %16 : Float(1000)
    14. %17 : Float(10, 3, 224, 224)) { # the actual input!
    15. # Every statement consists of some output tensors (and their types),
    16. # etc.), its input tensors (%17, %1)
    17. %19 : UNKNOWN_TYPE = Conv[kernels=[11, 11], strides=[4, 4], pads=[2, 2, 2, 2], dilations=[1, 1], group=1](%17, %1), uses = [[%20.i0]];
    18. # UNKNOWN_TYPE: sometimes type information is not known. We hope to eliminate
    19. # all such cases in a later release.
    20. %20 : Float(10, 64, 55, 55) = Add[broadcast=1, axis=1](%19, %2), uses = [%21.i0];
    21. %21 : Float(10, 64, 55, 55) = Relu(%20), uses = [%22.i0];
    22. %22 : Float(10, 64, 27, 27) = MaxPool[kernels=[3, 3], pads=[0, 0, 0, 0], dilations=[1, 1], strides=[2, 2]](%21), uses = [%23.i0];
    23. # ...
    24. # Finally, a network returns some tensors
    25. return (%58);
    26. }

    您也可以使用onnx库来验证protobuf。你可以onnxconda安装:

    然后,你可以运行:

    1. import onnx
    2. model = onnx.load("alexnet.proto")
    3. # Check that the IR is well formed
    4. onnx.checker.check_model(model)
    5. # Print a human readable representation of the graph
    6. onnx.helper.printable_graph(model.graph)

    要用caffe2运行导出的脚本,您将需要三件事情:

    • 2、你需要onnx-caffe2,一个纯Python库,为ONNX提供一个Caffe2后端。onnx-caffe2你可以用pip来安装:

    安装完成后,您可以使用Caffe2的后端:

    1. # ...continuing from above
    2. import onnx_caffe2.backend as backend
    3. import numpy as np
    4. rep = backend.prepare(model, device="CUDA:0") # or "CPU"
    5. # For the Caffe2 backend:
    6. # rep.predict_net is the Caffe2 protobuf for the network
    7. # rep.workspace is the Caffe2 workspace for the network
    8. # (see the class onnx_caffe2.backend.Workspace)
    9. outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
    10. # To run networks with more than one input, pass a tuple
    11. # rather than a single numpy ndarray.
    12. print(outputs[0])

    未来,其他框架也会有后端。

    • ONNX出口是一个基于跟踪的出口,这意味着它通过执行一次模型来运行,并导出在此运行期间实际运行的运营商。这意味着如果您的模型是动态的,例如,根据输入数据改变行为,则导出将不准确。类似地,跟踪可能仅对特定输入大小有效(这是我们为什么需要明确输入跟踪的一个原因)。我们建议检查模型跟踪并确保跟踪的运算符看起来合理。
    • PyTorchCaffe2通常具有一些数字差异的操作符的实现。根据模型结构的不同,这些差异可能是微不足道的,但是它们也会造成行为上的重大分歧(特别是在未经训练的模型上)。在未来的版本中,我们计划让Caffe2直接调用操作员的Torch实现,以帮助您平滑在精度很重要的时候处理这些差异,并记录这些差异。

    支持以下运算符: add (nonzero alpha not supported),sub (nonzero alpha not supported),mul,div,cat,mm,addmm,neg,,sigmoid,mean,t,expand (only when used before a broadcasting ONNX operator; e.g. add),transpose,view,split,squeeze,prelu (single weight shared among input channels not supported),threshold (non-zero threshold/non-zero value not supported),leaky_relu,glu,softmax,avg_pool2d (ceil_mode not supported),log_softmax,unfold (experimental support with ATen-Caffe2 integration),elu,Conv,BatchNorm,MaxPool1d (ceil_mode not supported),MaxPool2d (ceil_mode not supported),MaxPool3d (ceil_mode not supported),Embedding (no optional arguments supported),RNN,ConstantPadNd,Dropout,FeatureDropout (training mode not supported),Index (constant integer and tuple indices supported),Negate

    以上设置的操作足以导出以下型号: AlexNet,DCGAN,DenseNet,Inception (warning: this model is highly sensitive to changes in operator implementation),ResNet,SuperResolution,VGG,word_language_model

    将模型导出为ONNX格式。这个导出器运行你的模型一次,以获得其导出的执行轨迹; 目前,它不支持动态模型(例如,RNN)。

    另见:onnx-export

    参数:

    • 模型(torch.nn.Module) - 要导出的模型。
    • args(参数元组) - 模型的输入,例如,这-model(*args)是模型的有效调用。任何非变量参数将被硬编码到导出的模型中; 任何变量参数都将成为输出模型的输入,按照它们在参数中出现的顺序。如果args是一个变量,这相当于用该变量的一个元组来调用它。(注意:将关键字参数传递给模型目前还不支持,如果需要,给我们留言。)

    • f - 类文件对象(必须实现返回文件描述符的fileno)或包含文件名的字符串。一个二进制Protobuf将被写入这个文件。

    • export_params(布尔,默认为True) - 如果指定,所有参数将被导出。如果要导出未经训练的模型,请将其设置为False。在这种情况下,导出的模型将首先将其所有参数作为参数,按照指定的顺序model.state_dict().values()