CODEWITHSUNDEEP
pytorchassertiondimensionsonnx
dspyz
dspyz

Reputation: 5524

How do I get `dynamo_export` to not export shape-asserts?

I'm working with a large PyTorch codebase that has frequent asserts to check that tensors have the expected dtype and shape (all our shapes are static. We don't use dynamic shapes for anything):

Eg.

def foo(tensor: torch.Tensor, X: int, Y: int, Z: int):
  assert tensor.dtype == torch.int32
  assert tensor.shape == (2, X, Y, Z)
  ...

We export all our models to onnx using torch.onnx.export (and then convert to a saved TRT engine) for inference, but it seems to have bugs and doesn't always output valid onnx graphs. It was recently suggested to me that I should instead use torch.onnx.dynamo_export instead, but when I do that, it fails pointing to the shape-asserts with:

torch._dynamo.exc.Unsupported: comparison GetAttrVariable(TupleVariable(), shape) <built-in function eq> SizeVariable()

Is there some way to have dynamo_export run the shape-asserts (failing if any dimensions are dynamic or having the dynamic dimensions be None or something) but not put them in the onnx graph?

Upvotes: 0

Views: 168

Answers (0)

Related Questions