ONNX export in PyTorch - Interactive Code Practice

Practice - 5 Tasks

1fill in blank

easy

Complete the code to export a PyTorch model to ONNX format.

PyTorch

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
dummy_input = torch.randn(1, 10)
torch.onnx.export(model, dummy_input, [1])

Drag options to blanks, or click blank then click option'

A"model.onnx"

Bmodel.onnx

Cdummy_input

Dtorch.onnx

Attempts:

3 left

💡 Hint

Common Mistakes

Not providing the file path as a string.

Passing the model or dummy input as the file path.

The second argument to torch.onnx.export is the dummy input, and the third argument is the file path as a string where the ONNX model will be saved.

2fill in blank

medium

Complete the code to specify the input names when exporting to ONNX.

PyTorch

torch.onnx.export(model, dummy_input, "model.onnx", input_names=[[1]])

Drag options to blanks, or click blank then click option'

Ainput1

B["input1"]

C"input1"

D'input1'

Attempts:

3 left

💡 Hint

Common Mistakes

Passing the input name without quotes.

Passing a list inside the blank when brackets are already present.

The input_names argument expects a list of strings. Since the list brackets are already in the code, just provide the string name with quotes.

3fill in blank

hard

Fix the error in the code to export the model with dynamic axes for variable batch size.

PyTorch

torch.onnx.export(model, dummy_input, "model.onnx", dynamic_axes={"input": [1])

Drag options to blanks, or click blank then click option'

A{"batch_size": "0"}

B{"0": "batch_size"}

C{"batch_size": 0}

D{"0": 0}

Attempts:

3 left

💡 Hint

Common Mistakes

Using integers instead of strings as keys.

Swapping keys and values in the dictionary.

The dynamic_axes argument expects a dictionary mapping input names to dictionaries that map axis indices (as strings) to axis names.

4fill in blank

hard

Fill both blanks to export the model with verbose output and specifying output names.

PyTorch

torch.onnx.export(model, dummy_input, "model.onnx", verbose=[1], output_names=[[2]])

Drag options to blanks, or click blank then click option'

ATrue

BFalse

C"output1"

Doutput1

Attempts:

3 left

💡 Hint

Common Mistakes

Passing output names without quotes.

Passing verbose as a string instead of boolean.

verbose expects a boolean value, so True or False. output_names expects a list of strings, so the string name with quotes is correct.

5fill in blank

hard

Fill all three blanks to export the model with opset version 11, enable constant folding, and set training mode to evaluation.

PyTorch

torch.onnx.export(model, dummy_input, "model.onnx", opset_version=[1], do_constant_folding=[2], training=[3])

Drag options to blanks, or click blank then click option'

A10

B11

CTrue

DFalse

Attempts:

3 left

💡 Hint

Common Mistakes

Using wrong opset version.

Not enabling constant folding.

Setting training to True instead of evaluation mode.

opset_version should be 11 for compatibility, do_constant_folding is a boolean to optimize constants, and training should be set to torch.onnx.TrainingMode.EVAL which is represented here as False for no training.

Practice

1. What is the main purpose of exporting a PyTorch model to ONNX format?

easy

A. To save the model in a universal format for sharing and deployment

B. To train the model faster on GPUs

C. To convert the model into a TensorFlow model automatically

D. To visualize the model architecture in PyTorch

Solution

1. Step 1: Understand ONNX export purpose

   ONNX is designed to save models in a format that can be used across different frameworks and platforms.

2. Step 2: Compare options

   Only To save the model in a universal format for sharing and deployment correctly describes the universal sharing and deployment purpose of ONNX export.

3. Final Answer:

   To save the model in a universal format for sharing and deployment -> Option A

4. Quick Check:

   ONNX export = universal format [OK]

Hint: ONNX = share and deploy models universally [OK]

Common Mistakes:

• Confusing ONNX export with training speedup
• Thinking ONNX converts models to TensorFlow automatically
• Assuming ONNX export is for visualization only

2. Which of the following is the correct way to prepare a PyTorch model for ONNX export?

easy

A. Call model.train() before export

B. Export without setting any input

C. Use a dummy input tensor matching the model input shape

D. Use model.eval() after export

Solution

1. Step 1: Identify preparation steps for ONNX export

   Model should be in evaluation mode and a dummy input tensor matching input shape is needed for tracing.

2. Step 2: Evaluate options

   Use a dummy input tensor matching the model input shape correctly states the use of a dummy input tensor. Call model.train() before export is wrong because model.train() is for training mode. Export without setting any input misses input, and Use model.eval() after export is incorrect order.

3. Final Answer:

   Use a dummy input tensor matching the model input shape -> Option C

4. Quick Check:

   Dummy input needed = Use a dummy input tensor matching the model input shape [OK]

Hint: Always use dummy input tensor before export [OK]

Common Mistakes:

• Forgetting to set model.eval() before export
• Not providing dummy input tensor
• Calling model.train() instead of eval()

3. Given the code below, what will be the output of print(onnx_model.graph.input[0].name) after export?

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
model.eval()
dummy_input = torch.randn(1, 3)
torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input_tensor"])
import onnx
onnx_model = onnx.load("model.onnx")
print(onnx_model.graph.input[0].name)

medium

A. x

B. input_tensor

C. input0

D. data

Solution

1. Step 1: Check input_names parameter in export

   The export call sets input_names=["input_tensor"], so the input name in ONNX graph should be "input_tensor".

2. Step 2: Confirm printed input name

   Loading the ONNX model and printing the first input name will output "input_tensor" as set.

3. Final Answer:

   input_tensor -> Option B

4. Quick Check:

   input_names param = input_tensor [OK]

Hint: input_names param sets ONNX input name [OK]

Common Mistakes:

• Assuming default input name like 'input0'
• Confusing PyTorch variable name with ONNX input name
• Not setting input_names and expecting custom name

4. You try to export a PyTorch model to ONNX but get an error: RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same. What is the best fix?

medium

A. Export without specifying input_names

B. Use model.train() instead of model.eval()

C. Remove dummy input tensor

D. Move both the model and the dummy input to CPU before export

Solution

1. Step 1: Understand the error cause

   The error means model weights are on CPU but input tensor is on GPU (cuda), causing type mismatch.

2. Step 2: Fix by aligning device

   Moving both the model and the dummy input to CPU ensures both are on the same device (CPU), fixing the mismatch.

3. Final Answer:

   Move both the model and the dummy input to CPU before export -> Option D

4. Quick Check:

   Device mismatch fix = move model and input to CPU [OK]

Hint: Ensure model and input are on same device before export [OK]

Common Mistakes:

• Ignoring device mismatch and exporting anyway
• Switching to train mode instead of fixing device
• Removing dummy input tensor causing other errors

5. You want to export a PyTorch model to ONNX with dynamic batch size support. Which argument should you add to torch.onnx.export to enable this?

hard

A. dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}

B. enable_dynamic_batch=True

C. set_dynamic=True

D. dynamic_batch=True

Solution

1. Step 1: Identify how to specify dynamic axes in ONNX export

   PyTorch ONNX export uses the dynamic_axes argument to mark dimensions as dynamic, e.g., batch size dimension 0.

2. Step 2: Check options for dynamic batch size

   Only dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} correctly uses dynamic_axes with dictionary specifying batch dimension 0 for input and output.

3. Final Answer:

   dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} -> Option A

4. Quick Check:

   Dynamic batch size = dynamic_axes param [OK]

Hint: Use dynamic_axes dict to set dynamic batch size [OK]

Common Mistakes:

• Using nonexistent parameters like enable_dynamic_batch
• Forgetting to specify dynamic axes for output
• Assuming batch size is dynamic by default