使用 AI Hub 优化模型

import qai_hub as hub
import torch
from torchvision.models import mobilenet_v2
import numpy as np

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

# Trace model (for on-device deployment)
input_shape = (1, 3, 224, 224)
example_input = torch.rand(input_shape)
traced_torch_model = torch.jit.trace(torch_model, example_input)

# Compile and optimize the model for a specific device
compile_job = hub.submit_compile_job(
    model=traced_torch_model,
    device=hub.Device("Dragonwing RB3 Gen 2 Vision Kit"),
    input_specs=dict(image=input_shape),
    #compile_options="--target_runtime tflite",
)

# Profiling Job
profile_job = hub.submit_profile_job(
    model=compile_job.get_target_model(),
    device=hub.Device("Dragonwing RB3 Gen 2 Vision Kit"),
)

sample = np.random.random((1, 3, 224, 224)).astype(np.float32)

# Inference Job
inference_job = hub.submit_inference_job(
    model=compile_job.get_target_model(),
    device=hub.Device("Dragonwing RB3 Gen 2 Vision Kit"),
    inputs=dict(image=[sample]),
)

# Download model
compile_job.download_target_model(filename="/tmp/mobilenetv2.tflite")