Hello @vict0rsch,

A couple of comments on this script:

• in torch_xla, most of the python code only builds an IR graph and sends it to device. This is very fast. However, the fact that a line is executed does not mean the results have been materialized, due to async execution.
• it's only when you mark the step w/ xm.mark_step(), or send the tensor to device as you're doing here, that the execution happens.
• Therefore, timing the code is not very straightforward.
• On top of this, when TPUs run the program for the first time, it has to compile the graph to a lower representation before executing on the device. This is known to be slow.
• So, this code really does the following underneath:
  1. Timer('inference') only measures construction of the ir graph.
  2. Timer('back from device') measures compilation, execution, and data transfer.
• If you did this forward pass 1000 times, the first few ones will be by far slowest due to compilation. So I suggest the following modifications to above:

# DO NOT MEASURE THIS -- GETS RID OF COMPILATION OVERHEAD
for i in range(20):
    y=model(data)
    xm.mark_step()  # at this point results get materialized.

# MEASURE THIS
with Timer('inference'):
    y = model(data)
    xm.mark_step()
    print(y.shape)
with Timer("back from device"):
    y = y.cpu().numpy()

@taylanbil this changes everything... thanks for pointing it out. I should always call xm.mark_step() then, right? I think I was not. I'm sorry for having msised this!

I'll update the above chart accordingly and get back to you asap. Thank you very much

@taylanbil then I guess the way to measure the inference time of 2 models would be

with Timer("Total inference"):
    with Timer("masker"):
        y = masker(x)
        xm.mark_step()
    with Timer("painter"):
        z = painter(y)
        xm.mark_step()

right?

in an ordinary training script, such as our imagenet example, you'll notice that we don't actually use xm.mark_step. This is because it's called every time the xla dataloader yields. Here in this script, you're not using a dataloader so you'll have to call it yourself.

As to your last comment; if you'd like to measure the two models separately, you'd do that, yes. I'd do the whole thing 10-20 times first to get rid of compilation overhead, and then measure as pointed out in my first comment.

Also, note that every time you mark step, you'll force execution, and most of the time you'll get better performance if you don't do that; i.e.

y=masker(x)
z=painter(y)
xm.mark_step()

should be faster than

y=masker(x)
xm.mark_step()
z=painter(y)
xm.mark_step()

The exception here is if the total IR graph the first snippet creates is too large and executing the graph actually causes an OOM or swapping in the TPU host.