I am concerned about the overhead. If I train a model using PyTorch Lightning, does the framework add any latency when I deploy that model for real-time inference? We have strict 50ms requirements for our API. Should I export back to raw PyTorch or ONNX before going live?
3 answers
There is virtually zero overhead during inference because a 'LightningModule' is just a 'torch.nn.Module'. When you call 'model(x)' for a prediction, you are executing the exact same underlying PyTorch code. However, for production-grade APIs with strict 50ms requirements, the best practice is to export your model to TorchScript or ONNX. Lightning makes this incredibly easy with built-in methods like 'to_torchscript()' and 'to_onnx()'. This strips away any Python-level logic and allows you to run the model in a C++ environment or a specialized inference engine like TensorRT for maximum performance.
Are you planning to run inference on the same GPU hardware you used for training, or are you moving to a CPU-based cloud instance?
The 'Trainer.predict' method is great for batch inference, but for a live API, definitely export to a serialized format for the best results.
Alice is spot on. We use Lightning for the research and training phase, then export to ONNX for our production Go-lang microservices. It's the best of both worlds.
James, that's a critical question. If they are moving to CPU, they should definitely look into Quantization. Lightning has excellent support for post-training quantization. By converting weights from float32 to int8, you can often get a 2-4x speedup on CPUs. The framework handles the "calibration" step for you, which ensures that you don't lose too much accuracy during the conversion process. This is vital for meeting those tight 50ms API deadlines.