ONNX export in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

ONNX export output shape verification

Consider the following PyTorch model and ONNX export code. What will be the shape of the output tensor when running the ONNX model with the given input shape?

PyTorch

import torch
import torch.nn as nn
import onnx

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
    def forward(self, x):
        return self.linear(x)

model = SimpleModel()
model.eval()
x = torch.randn(3, 10)

# Export to ONNX
onnx_path = "simple_model.onnx"
torch.onnx.export(model, x, onnx_path, input_names=["input"], output_names=["output"], opset_version=11)

# Assume we load and run the ONNX model with input shape (3, 10)
# What is the shape of the output tensor?

A (3, 5)

B (10, 5)

C (3, 10)

D (5, 3)

💡 Hint

Think about how the linear layer transforms the input shape.

The linear layer transforms input of shape (batch_size, 10) to (batch_size, 5). Since input batch size is 3, output shape is (3, 5).

❓ Model Choice

intermediate

Choosing correct opset version for ONNX export

You want to export a PyTorch model using torch.onnx.export. Which opset_version should you choose to ensure compatibility with most ONNX runtimes and support for common operators?

A opset_version=7

B opset_version=11

C opset_version=15

D opset_version=9

💡 Hint

Check the default recommended opset version in PyTorch documentation.

PyTorch recommends opset_version=11 for broad compatibility and support of common operators.

🔧 Debug

advanced

ONNX export error due to dynamic axes

You try to export a PyTorch model to ONNX with dynamic batch size using this code: torch.onnx.export(model, x, "model.onnx", dynamic_axes={"input": {0: "batch_size"}}) But you get an error about missing output dynamic axes. What is the most likely cause?

A You must specify dynamic axes for both input and output tensors.

B The model does not support dynamic batch size.

C You forgot to set opset_version=15.

D You need to convert the model to eval mode before export.

💡 Hint

Dynamic axes must be declared for all tensors that vary in size.

When using dynamic axes, you must specify them for both inputs and outputs if their batch size varies.

❓ Metrics

advanced

Verifying ONNX model correctness with output comparison

After exporting a PyTorch model to ONNX, you want to verify the ONNX model produces the same output as the PyTorch model for the same input. Which metric is best to check this?

A Cross-entropy loss between PyTorch and ONNX outputs

B Accuracy score between PyTorch and ONNX outputs

C Mean squared error (MSE) between PyTorch and ONNX outputs

D Confusion matrix of PyTorch and ONNX outputs

💡 Hint

You want to measure numerical closeness of outputs.

MSE measures the average squared difference between outputs, suitable for regression outputs comparison.

🧠 Conceptual

expert

Understanding ONNX export limitations with control flow

Which of the following statements about exporting PyTorch models with control flow (like loops or conditionals) to ONNX is TRUE?

A Control flow must be removed from the model before exporting to ONNX.

B ONNX export can handle control flow only if it is traced, not scripted.

C ONNX supports all Python control flow constructs natively during export.

D ONNX supports control flow if the model is scripted with torch.jit.script before export.

💡 Hint

Think about how PyTorch JIT scripting helps ONNX export.

ONNX supports control flow when exporting scripted models, as scripting preserves control flow structure.