import os
import numpy as np
import tensorflow.compat.v1 as tf
tf.compat.v1.disable_eager_execution()
 
def _wget_demo_tgz():
    # 此处以下载一个公开的resnet50模型为例。
    url = 'https://cce-ai-native-package-bj.bj.bcebos.com/aiak-inference/examples/models/resnet50.pb'
    local_tgz = os.path.basename(url)
    local_dir = local_tgz.split('.')[0]
    if not os.path.exists(local_dir):
        luno.util.wget(url, local_tgz)
        luno.util.unpack(local_tgz)
    model_path = os.path.abspath(os.path.join(local_dir, "frozen_inference_graph.pb"))
    graph_def = tf.GraphDef()
    with open(model_path, 'rb') as f:
        graph_def.ParseFromString(f.read())
    # 以随机数作为测试数据，可替换为自己的数据集
    test_data = np.random.rand(1, 800, 1000, 3)
    return graph_def, {'image_tensor:0': test_data}
 
graph_def, test_data = _wget_demo_tgz()
 
input_nodes=['image_tensor']
output_nodes = ['detection_boxes', 'detection_scores', 'detection_classes', 'num_detections', 'detection_masks']

import time
 
def benchmark(model):
    tf.reset_default_graph()
    with tf.Session() as sess:
        sess.graph.as_default()
        tf.import_graph_def(model, name="")
        # Warmup!
        for i in range(0, 1000):
            sess.run(['image_tensor:0'], test_data)
        # Benchmark!
        num_runs = 1000
        start = time.time()
        for i in range(0, num_runs):
            sess.run(['image_tensor:0'], test_data)
        elapsed = time.time() - start
        rt_ms = elapsed / num_runs * 1000.0
        # Show the result!
        print("Latency of model: {:.2f} ms.".format(rt_ms))
 
# original graph
print("=====Original Performance=====")
benchmark(graph_def)

import luno
optimized_model = luno.optimize(
    graph_def,                 # 待优化的模型，此处是tf.GraphDef, 也可以配置为SavedModel的路径。
    'o1',                      # 优化级别，o1或o2。
    device_type='gpu',         # 目标设备，gpu/cpu
    outputs=['detection_boxes', 'detection_scores', 'detection_classes', 'num_detections', 'detection_masks']
)

# optimized graph
print("=====Optimized Performance=====")
benchmark(optimized_model)

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

model = models.resnet50().float().cuda()
model = torch.jit.script(model).eval()     # 使用jit转为静态图
dummy = torch.rand(1, 3, 224, 224).cuda()

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

# benchmark before optimization
print("before optimization:")
benchmark(model, dummy)

import luno

optimized_model = luno.optimize(
    model,
    'gpu',
    input_shapes=[[1, 3, 224, 224]],
    test_data=[dummy],
)

# benchmark after optimization
print("after optimization:")
benchmark(optimized_model, dummy)