#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@project: 
@File    : testfps
@Author  : qiqq
@create_time    : 2023/11/6 16:40
"""

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import  torch
import  numpy as  np

from taihuyuan_pv.compared_experiment.mySETR.model.setr import MSETR_naive

model = MSETR_naive(
   pretrained="/home/qiqq/q3dl/code/pretrain_weight/pretrained/vit_checkpoint/setrvit/vit_base_p16_224-4e355ebd.pth")

model_path = "/home/qiqq/q3dl/code/applicationproject/taihuyuan_pv/compared_experiment/mySETR/train_val_test/checkpoints/2023-08-08-19.46/best.pth"
model_dict = torch.load(model_path)
model.load_state_dict(model_dict['net'])
print("权重加载")
model.eval()

device = torch.device("cuda")
model.to(device)
dummy_input = torch.randn(1, 3, 512, 512,dtype=torch.float).to(device)
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
repetitions = 300
timings=np.zeros((repetitions,1))
#GPU-WARM-UP
for _ in range(10):
   _ = model(dummy_input)
# MEASURE PERFORMANCE
with torch.no_grad():
  for rep in range(repetitions):
     starter.record()
     _ = model(dummy_input)
     ender.record()
     # WAIT FOR GPU SYNC
     torch.cuda.synchronize()
     curr_time = starter.elapsed_time(ender)
     timings[rep] = curr_time
mean_syn = np.sum(timings) / repetitions
std_syn = np.std(timings)
mean_fps = 1000. / mean_syn
print(' * Mean@1 {mean_syn:.3f}ms Std@5 {std_syn:.3f}ms FPS@1 {mean_fps:.2f}'.format(mean_syn=mean_syn, std_syn=std_syn, mean_fps=mean_fps))
print(mean_syn)


# * Mean@1 9.347ms Std@5 0.468ms FPS@1 106.99