Python / PyTorch Fundamentals Interview Questions
What is torch.compile() vs TorchScript and how do you export a PyTorch model for production deployment?
PyTorch offers two main paths for production deployment beyond running the Python interpreter: TorchScript (serialises the model as a language-independent IR) and torch.compile() (JIT compiles for speed within Python). For cross-language/cross-framework deployment, ONNX export is also widely used.
| Method | Best for | Requires Python runtime? | Portable across languages? |
|---|---|---|---|
| torch.compile() | Fastest Python inference; no code changes | Yes | No |
| TorchScript (trace) | Production servers; models with fixed control flow | No | Yes (C++ API) |
| TorchScript (script) | Models with data-dependent control flow (if/loops) | No | Yes |
| ONNX export | Cross-framework deployment (TensorRT, ONNX Runtime, CoreML) | No | Yes (many runtimes) |
import torch import torch.nn as nn model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5)) model.eval() # ALWAYS call eval() before exporting # Γ’ΒΒΓ’ΒΒ 1. torch.compile() Γ’ΒΒ fast Python-based inference (PyTorch 2.0+) compiled = torch.compile(model) with torch.no_grad(): out = compiled(torch.randn(4, 10)) # Γ’ΒΒΓ’ΒΒ 2. TorchScript trace Γ’ΒΒ captures a concrete execution trace # Works best when control flow does NOT depend on input data example_input = torch.randn(1, 10) traced = torch.jit.trace(model, example_input) torch.jit.save(traced, "model_traced.pt") # Load and run without the original Python class loaded_traced = torch.jit.load("model_traced.pt") out = loaded_traced(torch.randn(4, 10)) # Γ’ΒΒΓ’ΒΒ 3. TorchScript script Γ’ΒΒ handles dynamic control flow class DynamicModel(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: if x.mean() > 0: # data-dependent branch Γ’ΒΒ trace would miss this! return torch.relu(x) return torch.tanh(x) scripted = torch.jit.script(DynamicModel()) torch.jit.save(scripted, "model_scripted.pt") # Γ’ΒΒΓ’ΒΒ 4. ONNX export Γ’ΒΒ deploy with ONNX Runtime, TensorRT, CoreML torch.onnx.export( model, example_input, "model.onnx", input_names=["features"], output_names=["logits"], dynamic_axes={"features": {0: "batch_size"}}, # variable batch size opset_version=17, ) # Inference with ONNX Runtime (no PyTorch dependency on deployment host!) # import onnxruntime as ort # sess = ort.InferenceSession("model.onnx") # out = sess.run(["logits"], {"features": x.numpy()})
More Related questions...