import time, random
n = random.randint(1, 10) * 100

# 获取当前时间 
start = time.time()

# 执行一些操作
print("Sleeping for {} ms".format(n))
time.sleep(n/1000)

# 比较当前时间和起始时间
print(time.time() - start)

# Output
# Sleeping for 500 ms
# 0.5713930130004883

def calc_param(model):
    """
    Calculate the number of parameters in the model.
    :param model: model
    :return: number of parameters
    """
    param_num =  sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(param_num*4/1024/1024, 'MB') # assume dtype float32

print(flop_count_table(FlopCountAnalysis(model, inputs)))

with torch.no_grad():
    with profile(activities=[ProfilerActivity.CUDA], record_shapes=True, use_cuda=True) as prof:
        with record_function("model_inference"):
            model(input_data['ego'])
    print("GPU time sorted operators:")
    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
    print("CPU time sorted operators:")
    print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ]
) as p:
    code_to_profile()
print(p.key_averages().table(
    sort_by="self_cuda_time_total", row_limit=-1))

from deepspeed.profiling.flops_profiler import FlopsProfiler

model = Model()
prof = FlopsProfiler(model)

profile_step = 5
print_profile= True

for step, batch in enumerate(data_loader):
  # start profiling at training step "profile_step"
  if step == profile_step:
    prof.start_profile()

  # forward() method
  loss = model(batch)

  # end profiling and print output
  if step == profile_step: # if using multi nodes, check global_rank == 0 as well
    prof.stop_profile()
    flops = prof.get_total_flops()
    macs = prof.get_total_macs()
    params = prof.get_total_params()
    if print_profile:
        prof.print_model_profile(profile_step=profile_step)
    prof.end_profile()

  # runs backpropagation
  loss.backward()

  # weight update
  optimizer.step()

def inference_throughput_cuda_event(model, data):
    print("start inference throughput performance test")
    run_num = 50
    print('warm up ...\n')
    with torch.no_grad(): # warm up
        for i in range(run_num):
            output = model(data)
    print('warm up done.')

    # synchronize 等待所有 GPU 任务处理完才返回 CPU 主线程
    torch.cuda.synchronize()

    # 设置用于测量时间的 cuda Event, 这是PyTorch 官方推荐的接口
    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    # 初始化一个时间容器
    run_num = 200
    timings = np.zeros((run_num,))

    print('start testing ...\n')
    with torch.no_grad():
        for i in range(run_num):
            starter.record()
            output = model(data)
            ender.record()
            torch.cuda.synchronize() # 等待GPU任务完成
            curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
            timings[i] = curr_time / 1000

    infer_thro = run_num / timings.sum()
    print("inference throughput (cuda event): ", infer_thro)

    return infer_thro

start_time = time.time()
for i in range(run_num):
    output = model(data)
end_time = time.time()
infer_thro = run_num / (end_time - start_time)