mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-06 04:35:00 -05:00
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
from extra.models.resnet import ResNet50
|
|
from tinygrad import Tensor
|
|
from tinygrad.helpers import Profiling, Timing, getenv, dedup
|
|
from tinygrad.ops import MetaOps
|
|
from tinygrad.codegen.kernel import Kernel
|
|
|
|
if __name__ == "__main__":
|
|
mdl = ResNet50()
|
|
img = Tensor.empty(64, 3, 224, 224)
|
|
|
|
PROFILE = getenv("PROFILE", 1)
|
|
FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
|
|
|
|
with Profiling(PROFILE):
|
|
with Timing("***** model forward in "):
|
|
out = mdl(img)
|
|
|
|
if not FORWARD_ONLY:
|
|
with Profiling(PROFILE):
|
|
with Timing("***** model schedule in "):
|
|
sched = out.schedule()
|
|
|
|
asts = dedup([x.ast for x in sched if x.ast.op is MetaOps.KERNEL])
|
|
uops = []
|
|
with Profiling(PROFILE):
|
|
with Timing("***** model uops in "):
|
|
for ast in asts:
|
|
k = Kernel(ast)
|
|
k.hand_coded_optimizations()
|
|
k.linearize()
|
|
uops.append((k.name, k.uops))
|
|
|
|
with Profiling(PROFILE, fn="/tmp/schedule.prof"):
|
|
with Timing("***** model linearize in "):
|
|
for _,u in uops: u.linearize()
|
|
|
|
#renderer = Device[Device.DEFAULT].renderer
|
|
#with Profiling(PROFILE, fn="/tmp/schedule.prof"):
|
|
# with Timing("***** model render in "):
|
|
# for n,u in uops: renderer.render(n, u)
|
|
|
|
# snakeviz /tmp/schedule.prof
|
|
#with Profiling(PROFILE, fn="/tmp/schedule.prof"):
|
|
# with Timing("***** model lower in "):
|
|
# eis = list(lower_schedule(sched))
|
|
|
|
# random makes this slow
|
|
#with Profiling(PROFILE):
|
|
# with Timing("***** model run in "):
|
|
# for ei in eis: ei.run()
|
|
|
|
# this is all wait
|
|
#with Profiling(PROFILE):
|
|
# with Timing("***** model finish in "):
|
|
# out.data()
|
|
|