Fix ONNX exports for Optimum compatible models

[merveenoyan]

@amyeroberts as discussed and also pinging @xenova for review :') (who also fixed DPT)

I prioritized Optimum compatible ones because I'm launching a project where there's Optimum examples for vision models. I will have a separate PR for the models that aren't compatible with Optimum. Rest of the Optimum compatible models export well without a problem.

[xenova]

It might be necessary to extract this cast function to a utility file, and then re-use across different models. I have a feeling that other instances, like:

will break execution, since during inference, it may require normal python types

[xenova]

Same comment as above - doesn't interpolate require (int, int) when not tracing?

[merveenoyan]

I'll check tracing, thanks for the heads up

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

Very nice! Thanks for fixing this for all these models ❤️

Just a few small comments

[amyeroberts]

Docstrings would be helpful here e.g. for inspecting in IDEs: what does it mean for an int to be safe?

[xenova]

Indeed a better name is probably a good idea 😅 I called it safe_int in a way to "safely cast some value (which could be a python number or tensor) to an integer in a way that respects tracing"

[merveenoyan]

I'll swap with torch_int and torch_float

[amyeroberts]

Do we need the conditional here? This is already handled in the safe_float and safe_int functions

[merveenoyan]

I think it's required for torch.ceil no?

[amyeroberts]

tbh, I don't know, is there a reason we couldn't usetorch.ceil directly?

[merveenoyan]

if I'm passing an int or float, torch.ceil will be called first and it will fail because torch.ceil can only be called with tensors AFAIK

[amyeroberts]

Makes sense!

[amyeroberts]

Only other Q here then is why do we use a float when tracing and int otherwise?

[merveenoyan]

sorry I think I was mistaken with that one, you're right, I fixed it :)

[merveenoyan]

@amyeroberts the failing tests seem irrelevant to this PR, I can't re-run them, can you re-run?

[amyeroberts]

@merveenoyan si si - done!

[amyeroberts]

Thanks for fixing all of these!

Just for my own understanding - it there any reason to not use the torch compatible float/int when were not tracing?

[merveenoyan]

@amyeroberts to my understanding, torch ONNX export internally calls trace thus to capture the graph properly all needs to be torch tensor, one needs to check tracing, and convert to torch tensor/convert ops internally if it's not torch

[merveenoyan]

@amyeroberts can you merge if you think it's ok?

[amyeroberts]

@amyeroberts to my understanding, torch ONNX export internally calls trace thus to capture the graph properly all needs to be torch tensor, one needs to check tracing, and convert to torch tensor/convert ops internally if it's not torch

Right, I see why we need to do it for the onnx export, but for day-to-day use could we just use torch primitives instead of a python int or float i.e. do we need to maintain this if/else structure or can we switch everything to torch land regardless of whether we're tracing or not?

[merveenoyan]

@amyeroberts I guess if it's just torch modelling code then yes. Would you like me to swap everything?

[merveenoyan]

also asking the same question to @xenova

[amyeroberts]

Would you like me to swap everything?

@merveenoyan Yes please! This will be cleaner and easier to follow in the code :)

[xenova]

also asking the same question to @xenova

I agree with @amyeroberts - if there is a way to "do everything in torch land", that's the best solution! However, there are cases where I'm not entirely sure how to do this. For example, with torch.nn.interpolate:

• during inference, the sizes/scales CANNOT be tensors. They must be python ints/floats.
• during tracing, the sizes/scales SHOULD be tensors. Dynamic shapes usually break if they are python ints/floats due to the loss of tracing.

See here for example code (DinoV2 backbone):

        if torch.jit.is_tracing():
          sqrt_N = N ** 0.5
          patch_pos_embed = nn.functional.interpolate(
              patch_pos_embed.reshape(1, (sqrt_N).to(torch.int64), (sqrt_N).to(torch.int64), dim).permute(0, 3, 1, 2),
              size=(w0, h0),
              mode="bicubic",
              antialias=self.interpolate_antialias,
          )
        else:
          sqrt_N = math.sqrt(N)
          sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
          patch_pos_embed = nn.functional.interpolate(
              patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
              scale_factor=(sx, sy),
              mode="bicubic",
              antialias=self.interpolate_antialias,
          )

Very ugly... I know :/

[merveenoyan]

@xenova sounds good, very glad to work with you tbh I didn't know that it would be required in inference.
@amyeroberts should we merge?

[amyeroberts]

@merveenoyan My understanding from above was that the PR would be updated to remove all the if/else structures wherever possible (but as @xenova points out isn't everywhere unfortunately)

[merveenoyan]

@amyeroberts from what I understood we should still keep them in if/else not to break the inference (I'm also scared of edge cases if there is etc) so I'd rather keep them. what I can do is to test all of them to see if they break or not when all are tensors and remove where it doesn't have to be a python type

[amyeroberts]

@merveenoyan OK. Let's just merge then and we can follow up in future PRs 👍