Невозможно передать kwargs в аргументы torch.onnx.export в Pytorch ONNX. ⇐ Python

[Anonymous]

Я хочу экспортировать модель Pytorch в формат onnx. Вот модель:

Код: Выделить всё

from transformers import AutoProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig

model_id = "google/paligemma-3b-mix-224"

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

model = PaliGemmaForConditionalGeneration.from_pretrained(
model_id, quantization_config=quantization_config
).eval()
processor = AutoProcessor.from_pretrained(model_id)

И мой код экспорта onnx:

Код: Выделить всё

torch.onnx.export(
model,
(input_ids, attention_mask, pixel_values),
"model.onnx",
input_names=['input_ids', 'attention_mask', 'pixel_values'],
output_names=['output']
)

Приведенный выше код выдает следующую ошибку:

Код: Выделить всё

/usr/local/lib/python3.10/dist-packages/accelerate/hooks.py in new_forward(module, *args, **kwargs)
164                 output = module._old_forward(*args, **kwargs)
165         else:
--> 166             output = module._old_forward(*args, **kwargs)
167         return module._hf_hook.post_forward(module, output)
168

/usr/local/lib/python3.10/dist-packages/transformers/models/siglip/modeling_siglip.py in forward(self, pixel_values, interpolate_pos_encoding)
302
303     def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor:
--> 304         _, _, height, width = pixel_values.shape
305         patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
306         embeddings = patch_embeds.flatten(2).transpose(1, 2)

ValueError: not enough values to unpack (expected 4, got 2)

Я попробовал несколько решений, я думаю, это потому, что при передаче входных данных в onnx это эквивалентно (этот код выдает ту же ошибку, что и выше):

Код: Выделить всё

model(input_ids, attention_mask, pixel_values)

Однако правильный способ передачи входных данных в модель — использование kwargs.

Код: Выделить всё

model(input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values)

Как передать ключевые слова в средство экспорта onnx или есть ли лучший способ экспорта формата onnx?


Подробнее здесь: https://stackoverflow.com/questions/786 ... torch-onnx