Torch.onnx.export of PyTorch model is slow - expected completion time?

sevagh: <p>Hello,</p>
<p>I’m trying to speed up my model inference. It’s a PyTorch module, pretty standard - no special ops, just PyTorch convolution layers.</p>
<p>The export code is copied from this tutorial <a href="https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html" class="inline-onebox" rel="noopener nofollow ugc">(optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime — PyTorch Tutorials 1.9.0+cu102 documentation</a> :</p>
<pre><code class="lang-auto">if __name__ == '__main__':
    model_str_or_path = sys.argv[1]
    model_path = Path(model_str_or_path).expanduser()

    # when path exists, we assume its a custom model saved locally
    if model_path.exists():
        with open(Path(model_path, "separator.json"), "r") as stream:
            enc_conf = json.load(stream)

        xumx_model, model_nsgt, jagged_slicq_sample = load_target_models(
            model_str_or_path=model_path, pretrained=True, sample_rate=enc_conf["sample_rate"], device="cpu"
        )

    xumx_model.eval()

    # Input to the model
    torch_out = xumx_model(jagged_slicq_sample)

    # Export the model
    torch.onnx.export(xumx_model,                # model being run
                      jagged_slicq_sample,       # model input (or a tuple for multiple inputs)
                      "xumx_slicq.onnx",         # where to save the model (can be a file or file-like object)
                      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'}})
</code></pre>
<p>My model has ~6.7 million parameters (torchinfo output):</p>
<pre><code class="lang-auto">===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
OpenUnmix                                     --                        --
Total params: 6,669,912
Trainable params: 6,669,912
Non-trainable params: 0
Total mult-adds (G): 194.27
===============================================================================================
Input size (MB): 28.63
Forward/backward pass size (MB): 9359.33
Params size (MB): 26.68
Estimated Total Size (MB): 9414.64
</code></pre>
<p>The trained size for the pth weights file is 28M.</p>
<p>My export command has been running for 3 hours, and has not completed yet. I’m wondering if that’s an expected export time for my size of model.</p>
<p>Thanks.</p>

ptrblck: <p>No, I guess your script might hang, as 3h are not expected.<br>
In case the process is still running, you could try to attach <code>gdb</code> to it and check the backtrace to see where it’s hanging.</p>

sevagh: <p>It actually ended up working. I used strace to see if it was hung or progressing, and could see some activity:</p>
<pre><code class="lang-auto">[pid 238336] munmap(0x7fb9bb300000, 262144) = 0
[pid 238336] munmap(0x7fb9bb7c0000, 262144) = 0
[pid 238336] munmap(0x7fb9bb800000, 262144) = 0
[pid 238336] munmap(0x7fb9bb700000, 262144) = 0
[pid 238336] munmap(0x7fb9bb380000, 262144) = 0
[pid 238336] munmap(0x7fb9bb180000, 262144) = 0
[pid 238336] mmap(NULL, 262144, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0) = 0x7fb9bb8c0000
[pid 238336] mmap(NULL, 262144, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0) = 0x7fb9bb880000
[pid 238336] mmap(NULL, 262144, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0) = 0x7fb9bb800000
[pid 238336] mmap(NULL, 262144, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0) = 0x7fb9bb7c0000
</code></pre>
<p>The whole thing ended up taking 8 (!) hours:</p>
<pre><code class="lang-auto">(umx-gpu-numba) sevagh:xumx-sliCQ $ ls -latrh xumx_slicq.onnx
-rw-r--r--. 1 sevagh sevagh 39M Jul 29 04:28 xumx_slicq.onnx
</code></pre>
<p>The verbose output shows a lot of info:</p>
<pre><code class="lang-auto">%247701 : Long(device=cpu) = onnx::Constant[value={0}]()
  %247702 : Long(device=cpu) = onnx::Constant[value={1}]()
  %247703 : Long(1, device=cpu) = onnx::Range(%247701, %247700, %247702)
  %247704 : Long(4, strides=[1], device=cpu) = onnx::Shape(%247669)
  %247705 : Long(device=cpu) = onnx::Constant[value={3}]()
  %247706 : Long(device=cpu) = onnx::Gather[axis=0](%247704, %247705)
  %247707 : Long(device=cpu) = onnx::Cast[to=7](%247706)
  %247708 : Long(device=cpu) = onnx::Constant[value={0}]()
  %247709 : Long(device=cpu) = onnx::Constant[value={1}]()
  %247710 : Long(644, device=cpu) = onnx::Range(%247708, %247707, %247709)
  %247712 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[axes=[0]](%247670)
  %247713 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[axes=[0]](%247671)
  %247715 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={1}]()
  %247716 : Long(184, device=cpu) = onnx::Slice(%247710, %247712, %247713, %263646, %247715)
  %247717 : Long(4, strides=[1], device=cpu) = onnx::Constant[value=-1  1  1  1 [ CPULongType{4} ]]()
  %247718 : Long(1, 1, 1, 1, strides=[1, 1, 1, 1], device=cpu) = onnx::Reshape(%247689, %247717)
  %247719 : Long(3, strides=[1], device=cpu) = onnx::Constant[value=-1  1  1 [ CPULongType{3} ]]()
  %247720 : Long(2, 1, 1, strides=[1, 1, 1], device=cpu) = onnx::Reshape(%247696, %247719)
  %247721 : Long(2, strides=[1], device=cpu) = onnx::Constant[value=-1  1 [ CPULongType{2} ]]()
  %247722 : Long(1, 1, strides=[1, 1], device=cpu) = onnx::Reshape(%247703, %247721)
  %247723 : Long(1, 2, 1, 1, strides=[2, 1, 1, 1], device=cpu) = onnx::Add(%247718, %247720)
  %247724 : Long(1, 2, 1, 1, strides=[2, 1, 1, 1], device=cpu) = onnx::Add(%247723, %247722)
</code></pre>
<p>It would be nice if onnx.export had a progress bar/tqdm feature built in, though, so that it’s not silent for 8 hours.</p>

ptrblck: <p>Creating a 39MB file in 8h seems terrible and I still think that this cannot be expected.<br>
Good to hear the script finally finished, but I guess the export might have been slower than the model training? <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=10" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"><br>
In case you can reproduce it, could you create an issue on <a href="https://github.com/pytorch/pytorch/issues">GitHub</a>, please?</p>

sevagh: <p>I’m a bit late here but I filed it on GitHub finally: <a href="https://github.com/pytorch/pytorch/issues/63734" class="inline-onebox" rel="noopener nofollow ugc">Torch.onnx.export is very slow, takes 8 hours · Issue #63734 · pytorch/pytorch · GitHub</a><br>
It should be reproducible.</p>