使用 AI Hub 优化模型

import torch
import torchvision

import qai_hub as hub

client = hub.Client()

# Using pre-trained MobileNet
torch_model = torchvision.models.mobilenet_v2(pretrained=True)
torch_model.eval()

# Trace model
input_shape: tuple[int, ...] = (1, 3, 224, 224)
example_input = torch.rand(input_shape)
with torch.no_grad():
    pt2_model = torch.export.export(torch_model, (example_input,))

# Compile model on a specific device
compile_job = client.submit_compile_job(
    pt2_model,
    name="MobileNet_V2",
    device=hub.Device("Dragonwing RB3 Gen 2 Vision Kit"),
    input_specs=dict(image=input_shape),
)

# Download the optimized compiled model
compile_job.download_target_model("MobileNet_V2.tflite")