test tflops

import torchvision.models as models
import torch
import onnx

from ptflops import get_model_complexity_info
# from time import sleep



def get_base_info():
	print(torch.__version__)

	# print(torch.cuda.get_device_name())

	if torch.cuda.is_available():
	    print("is availiable")
	else:
	    print("not availiable")




def get_flops_resnet():

    net = models.resnet50()
    macs, params = get_model_complexity_info(
	                    net, 
	                    (3, 224, 224), 
	                    as_strings=True,
	                    print_per_layer_stat=True, 
	                    verbose=True)

    print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
    print('{:<30}  {:<8}'.format('Number of parameters: ', params))



def get_flops_yolov5m():
    # net = torch.hub.load('ultralytics/yolov5', 'yolov5m', pretrained=True)

    path='/disk2T/dyling/code/yolov5-master/'

    net = torch.hub.load(path, 'yolov5m', pretrained=True, source="local")

    imgs = ['https://ultralytics.com/images/zidane.jpg']  # batch of images

    macs, params = get_model_complexity_info(
	                    net, 
	                    (3, 640, 640), 
	                    as_strings=True,
	                    print_per_layer_stat=True, 
	                    verbose=True)

    print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
    print('{:<30}  {:<8}'.format('Number of parameters: ', params))
    

def load_model():
    model = onnx.load("/disk2T/dyling/test_result/detail_info/yolov5m.onnx")
    # model.eval()


if __name__ == "__main__":
    #get_flops_resnet()
    get_flops_yolov5m()