Platform For AI:Optimize a PyTorch model

Step 1: Import dependencies

import os
import time
import torch
import torchvision.models as models
import blade

Step 2: Load and convert the model

PAI-Blade only accepts torch.jit.ScriptModule objects — it does not work with standard nn.Module models. Convert your model using torch.jit.script before passing it to PAI-Blade.

model = models.resnet50().float().cuda()  # Load the ResNet50 model.
model = torch.jit.script(model).eval()   # Convert to ScriptModule.
dummy = torch.rand(1, 3, 224, 224).cuda() # Create test input.

Step 3: Run optimization

Call blade.optimize to optimize the model. The method accepts the following key parameters:

For the full parameter reference, see Python method.

optimized_model, opt_spec, report = blade.optimize(
    model,                 # The model to optimize.
    'o1',                  # Optimization level.
    device_type='gpu',     # Target device type.
    test_data=[(dummy,)],  # Test input data.
)

blade.optimize returns three objects:

• `optimized_model`: The optimized model as a torch.jit.ScriptModule object, ready for inference.

• `opt_spec`: The external dependencies (configuration, environment variables, resource files) needed to reproduce the optimization. Use a with statement in Python to activate these dependencies.

• `report`: The optimization report, which you can print directly. See Optimization report for details on each field.

Optimization takes time. PAI-Blade runs multiple analysis phases before producing the optimized model. Progress is shown as optimization runs:

[Progress] 5%, phase: user_test_data_validation.
[Progress] 10%, phase: test_data_deduction.
[Progress] 15%, phase: CombinedSwitch_4.
[Progress] 95%, phase: model_collecting.

This compilation is a one-time cost. Once optimization completes, inference with optimized_model is consistently faster than the original model.

Step 4: Review the optimization report

print("Report: {}".format(report))

The report shows which optimizations were applied and their effect on latency. A sample report:

{
  "optimizations": [
    {
      "name": "PtTrtPassFp32",
      "status": "effective",
      "speedup": "1.50",     // Speedup ratio for this optimization.
      "pre_run": "5.29 ms",  // Latency before optimization.
      "post_run": "3.54 ms"  // Latency after optimization.
    }
  ],
  "overall": {
    "baseline": "5.30 ms",    // Original model latency.
    "optimized": "3.59 ms",   // Optimized model latency.
    "speedup": "1.48"         // End-to-end speedup ratio.
  }
}

In this example, the end-to-end speedup is 1.48x — the optimized model runs in 3.59 ms versus 5.30 ms for the original. Check "status": "effective" to identify which optimization passes had a measurable impact.

Step 5: Benchmark inference performance

Use the following function to measure steady-state latency for both the original and optimized models. The first 100 runs warm up the GPU; the subsequent 200 runs are timed.

@torch.no_grad()
def benchmark(model, inp):
    for i in range(100):    # Warm up.
        model(inp)
    start = time.time()
    for i in range(200):    # Timed runs.
        model(inp)
    elapsed_ms = (time.time() - start) * 1000
    print("Latency: {:.2f}".format(elapsed_ms / 200))

benchmark(model, dummy)           # Original model.
benchmark(optimized_model, dummy) # Optimized model.