Optimize a model with dynamic input shapes using Blade - Platform For AI

Standard inference optimization assumes fixed input shapes. When input shapes vary at runtime — different batch sizes in an online serving system, variable-resolution images in a preprocessing pipeline, or variable-length sequences in a BERT-style NLP model — static optimization loses most of its effect. Blade supports dynamic shape optimization so the compiled model stays fast across a defined range of input shapes, without requiring a separate compiled artifact per shape.

This tutorial walks through optimizing a ResNet50 model with a variable batch dimension using Blade, from setup through deployment.

Prerequisites

Before you begin, make sure you have:

Python 3.6 or later in a Linux environment
PyTorch 1.7.1
NVIDIA T4 GPU with CUDA 11.0
Blade 3.17.0 or later installed

Step 1: Prepare the model and test data

Download the pre-trained ResNet50 checkpoint and a sample from the ImageNet-1k validation set. The checkpoint comes from torchvision and is hosted in Object Storage Service (OSS) for faster download. The sample data is pre-processed and ready to use.

wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/resnet50-19c8e357.pth -O resnet50-19c8e357.pth
wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/imagenet_val_example.pt -O imagenet_val_example.pt

Build the model, load the checkpoint, and export to TorchScript:

import torch
import torchvision

# Build ResNet50 and move to GPU.
model = torchvision.models.resnet50().eval().cuda()
# Load the pre-trained checkpoint.
ckpt = torch.load('resnet50-19c8e357.pth')
model.load_state_dict(ckpt)
# Load the test data.
example_input = torch.load('imagenet_val_example.pt').cuda()
# Export to TorchScript.
traced_model = torch.jit.trace(model, example_input).cuda().eval()

Step 2: Configure the dynamic shape range

Blade optimizes for a declared range of input shapes. Define the range, then build the Blade config from it.

Define the shape range

A dynamic range requires three fields:

Field	Description
`min`	Lower bound of the input shape
`max`	Upper bound of the input shape
`opts`	Shapes that receive extra tuning; inference speed is typically highest at these shapes

All three fields must satisfy min_num <= opt_num <= max_num at each corresponding dimension. The number of shape groups in each field must equal the number of network inputs.

This example varies only the batch dimension (dimension 0), keeping the image dimensions fixed at [3, 224, 224]:

shapes = {
    "min": [[1, 3, 224, 224]],
    "max": [[10, 3, 224, 224]],
    "opts": [
        [[5, 3, 224, 224]],
        [[8, 3, 224, 224]],
    ]
}

If the gap between min and max is large, the speedup may be limited. Split a large range into multiple smaller ranges for better results. See Appendix: Set multiple dynamic ranges.

Build the Blade config

import blade
import blade.torch as blade_torch

# Attach the shape range to the BladeTorch config.
blade_torch_cfg = blade_torch.Config()
blade_torch_cfg.dynamic_tuning_shapes = shapes

# Disable the FP16 accuracy check to maximize acceleration.
gpu_config = {
    "disable_fp16_accuracy_check": True,
}
blade_config = blade.Config(gpu_config=gpu_config)

Step 3: Optimize the model

Call blade.optimize inside the blade_torch_cfg context. The function returns the optimized model, an unused intermediate value, and an optimization report.

with blade_torch_cfg:
    optimized_model, _, report = blade.optimize(
        traced_model,                    # Input TorchScript model.
        'o1',                            # o1 = lossless optimization.
        config=blade_config,
        device_type='gpu',               # Target GPU inference.
        test_data=[(example_input,)]     # Must be within the declared shape range.
    )

The optimized model has the same type as the input. Because the input is a TorchScript, the output is also a TorchScript.

Print the optimization report to see which optimizations were applied and the measured speedup:

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

Expected output (your results will vary):

{
  "software_context": [
    {"software": "pytorch", "version": "1.7.1+cu110"},
    {"software": "cuda", "version": "11.0.0"}
  ],
  "hardware_context": {
    "device_type": "gpu",
    "microarchitecture": "T4"
  },
  "user_config": "",
  "diagnosis": {
    "model": "unnamed.pt",
    "test_data_source": "user provided",
    "shape_variation": "undefined",
    "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
    "test_data_info": "0 shape: (1, 3, 224, 224) data type: float32"
  },
  "optimizations": [
    {
      "name": "PtTrtPassFp16",
      "status": "effective",
      "speedup": "4.06",
      "pre_run": "6.55 ms",
      "post_run": "1.61 ms"
    }
  ],
  "overall": {
    "baseline": "6.54 ms",
    "optimized": "1.61 ms",
    "speedup": "4.06"
  },
  "model_info": {"input_format": "torch_script"},
  "compatibility_list": [
    {"device_type": "gpu", "microarchitecture": "T4"}
  ],
  "model_sdk": {}
}

In this example, the PtTrtPassFp16 optimization reduces inference time from 6.55 ms to 1.61 ms — a 4.06x speedup on the test shape. These numbers reflect this specific setup; actual results depend on your model and hardware. For a description of each field, see Optimization report.

Save the optimized model to disk:

file_name = "resnet50_opt.pt"
torch.jit.save(optimized_model, file_name)
# Reload from disk to verify the saved file is valid.
optimized_model = torch.jit.load(file_name)

Step 4: Verify performance and correctness

Before deploying, confirm that the optimized model is faster and produces the same predictions across multiple input shapes.

Benchmark inference speed

Define a benchmark function that warms up the model 10 times, then measures the average latency over 100 runs:

import time

@torch.no_grad()
def benchmark(model, test_data):
    model = model.eval()
    # Warm up.
    for i in range(10):
        model(test_data)
    # Timed runs.
    num_runs = 100
    start = time.time()
    for i in range(num_runs):
        model(test_data)
    torch.cuda.synchronize()
    elapsed = time.time() - start
    rt_ms = elapsed / num_runs * 1000.0
    print("{:.2f} ms.".format(rt_ms))
    return rt_ms

Run the benchmark across batch sizes 1, 3, 5, 7, and 9 — all within the declared [1, 10] range:

dummy_inputs = []
batch_num = [1, 3, 5, 7, 9]
for n in batch_num:
    dummy_inputs.append(torch.randn(n, 3, 224, 224).cuda())

for inp in dummy_inputs:
    print(f'--------------test with shape {list(inp.shape)}--------------')
    print("  Origin model inference cost:     ", end='')
    origin_rt = benchmark(traced_model, inp)
    print("  Optimized model inference cost:  ", end='')
    opt_rt = benchmark(optimized_model, inp)
    speedup = origin_rt / opt_rt
    print('  Speed up: {:.2f}'.format(speedup))
    print('')

Expected output:

--------------test with shape [1, 3, 224, 224]--------------
  Origin model inference cost:     6.54 ms.
  Optimized model inference cost:  1.66 ms.
  Speed up: 3.94

--------------test with shape [3, 3, 224, 224]--------------
  Origin model inference cost:     10.79 ms.
  Optimized model inference cost:  2.40 ms.
  Speed up: 4.49

--------------test with shape [5, 3, 224, 224]--------------
  Origin model inference cost:     16.27 ms.
  Optimized model inference cost:  3.25 ms.
  Speed up: 5.01

--------------test with shape [7, 3, 224, 224]--------------
  Origin model inference cost:     22.62 ms.
  Optimized model inference cost:  4.39 ms.
  Speed up: 5.16

--------------test with shape [9, 3, 224, 224]--------------
  Origin model inference cost:     28.83 ms.
  Optimized model inference cost:  5.25 ms.
  Speed up: 5.49

The optimized model is 3.94–5.49x faster across all tested batch sizes. Speedup is typically higher near the opts shapes (5 and 8 in this example). These results are specific to this setup; actual results will vary.

Verify prediction correctness

Run both models on the same input and compare their top-1 predictions:

origin_output = traced_model(example_input)
_, pred = origin_output.topk(1, 1, True, True)
print("origin model output: {}".format(pred))

opt_output = optimized_model(example_input)
_, pred = opt_output.topk(1, 1, True, True)
print("optimized model output: {}".format(pred))

Expected output:

origin model output: tensor([[834]], device='cuda:0')
optimized model output: tensor([[834]], device='cuda:0')

Both models predict class 834 for the test input, confirming that the optimization preserves prediction accuracy.

Step 5: Deploy the optimized model

Blade provides Python and C++ SDKs for deployment. This section covers the Python SDK. For the C++ SDK, see Use an SDK to deploy a TensorFlow model for inference.

(Optional) During the trial period, set the following environment variable to prevent unexpected exits from authentication failures:
```
export BLADE_AUTH_USE_COUNTING=1
```

Authenticate with PAI-Blade by setting your region and token:

Placeholder	Description
`<region>`	The region where you use PAI-Blade. Join the PAI-Blade DingTalk group to get the list of available regions. For the QR code, see Install PAI-Blade.
`<token>`	The authentication token for PAI-Blade. Join the PAI-Blade DingTalk group to get your token. For the QR code, see Install PAI-Blade.

export BLADE_REGION=<region>
export BLADE_TOKEN=<token>

Replace the placeholders with your values:

Load and run the optimized model. Add import blade.runtime.torch — no other changes to your inference code are needed:

import torch
import blade.runtime.torch

# Replace <your_optimized_model_path> with the path to your saved model file.
opt_model_dir = "<your_optimized_model_path>"
# Replace <your_infer_data> with your inference input tensor.
infer_data = "<your_infer_data>"

model = torch.jit.load(opt_model_dir)
output = model(infer_data)

Appendix: Set multiple dynamic ranges

If the gap between min and max is large, the overall speedup may be limited because Blade has to accommodate too wide a range. Split the range into multiple smaller ranges to improve acceleration.

For example, split [1, 10] into [1, 5] and [5, 10]:

shapes1 = {
    "min": [[1, 3, 224, 224]],
    "max": [[5, 3, 224, 224]],
    "opts": [
        [[5, 3, 224, 224]],
    ]
}
shapes2 = {
    "min": [[5, 3, 224, 224]],
    "max": [[10, 3, 224, 224]],
    "opts": [
        [[8, 3, 224, 224]],
    ]
}
shapes = [shapes1, shapes2]

Pass this shapes list to blade_torch_cfg.dynamic_tuning_shapes as described in Step 2.