Skip to content

Commit

Permalink
Use height and width from preprocessor
Browse files Browse the repository at this point in the history
  • Loading branch information
xenova committed Dec 9, 2023
1 parent be9c707 commit c4d6bc2
Showing 1 changed file with 36 additions and 1 deletion.
37 changes: 36 additions & 1 deletion optimum/exporters/onnx/model_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -708,8 +708,43 @@ class ConvNextV2OnnxConfig(ViTOnnxConfig):
pass


class Dinov2DummyInputGenerator(DummyVisionInputGenerator):
def __init__(
self,
task: str,
normalized_config: NormalizedVisionConfig,
batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"],
num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"],
width: int = DEFAULT_DUMMY_SHAPES["width"],
height: int = DEFAULT_DUMMY_SHAPES["height"],
**kwargs,
):
super().__init__(
task=task,
normalized_config=normalized_config,
batch_size=batch_size,
num_channels=num_channels,
width=width,
height=height,
**kwargs,
)

from transformers.onnx.utils import get_preprocessor

preprocessor = get_preprocessor(normalized_config._name_or_path)
if preprocessor is not None and hasattr(preprocessor, "crop_size"):
self.height = preprocessor.crop_size.get("height", self.height)
self.width = preprocessor.crop_size.get("width", self.width)

def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
input_ = super().generate(
input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype
)
return input_


class Dinov2OnnxConfig(ViTOnnxConfig):
pass
DUMMY_INPUT_GENERATOR_CLASSES = (Dinov2DummyInputGenerator, )


class MobileViTOnnxConfig(ViTOnnxConfig):
Expand Down

0 comments on commit c4d6bc2

Please sign in to comment.