import roctracer_py
import numpy as np
import torch
from torchvision import models, transforms

# 创建wrapper实例
tracer = roctracer_py.RoctracerWrapper()


# 添加标记
tracer.mark("Program start")

# 开始一个范围
range_id = tracer.range_start("Computation")


# 内存拷贝（模拟）
tracer.range_push("Memory Copy")
model = models.resnet50(pretrained=False)
device = torch.device("cuda:0")  # 指定设备为CPU
model.to(device) 
tracer.range_pop()


input = torch.zeros((1, 3, 224, 224), dtype=torch.float).to(device)

tracer.range_push("Model infer")

with torch.no_grad():
    outputs = model(input)

tracer.range_pop()

# 停止范围
tracer.range_stop(range_id)

print("Done!")