[PT FE]: support nested inputs in example_inputs and arg dicts with d… (#18492)

* [PT FE]: support nested inputs in example_inputs and arg dicts with different argtypes

* accept hande lists as inputs

* Update tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py

* update tests and add comments in code

* fix for custom types in annotations and duplicate in mo

* Update tools/mo/openvino/tools/mo/moc_frontend/pytorch_frontend_utils.py
This commit is contained in:
Ekaterina Aidova
2023-07-19 17:01:22 +02:00
committed by GitHub
parent 186b1b6bfc
commit 61504bbfc2
4 changed files with 251 additions and 35 deletions
@@ -3,6 +3,7 @@
import os
from typing import Tuple
import numpy
import numpy as np
import openvino.runtime as ov
@@ -725,6 +726,150 @@ def create_pytorch_module_with_compressed_int8_constant(tmp_dir):
ref_model = Model([conv], [param1], "test")
return traced_model, ref_model, {"example_input": example_input}
def create_pytorch_module_with_nested_inputs(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, z:Tuple[torch.Tensor, torch.Tensor]):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
return torch.cat([z1, zeros1], 1), torch.cat([z2, zeros2], 2)
net = PTModel()
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
ref_model = Model([concat2, concat1], [param1, param2], "test")
return net, ref_model, {"example_input": {"z": (torch.zeros((1, 10)), torch.ones((1, 5, 2)))}}
def create_pytorch_module_with_nested_inputs2(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, x:torch.Tensor, z:Tuple[torch.Tensor, torch.Tensor]):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
return torch.cat([z1, zeros1], 1) + x, torch.cat([z2, zeros2], 2)
net = PTModel()
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
param0 = ov.opset10.parameter(PartialShape([-1, -1]), dtype=np.float32)
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
add = ov.opset10.add(concat1, param0)
ref_model = Model([concat2, add], [param0, param1, param2], "test")
return net, ref_model, {"example_input": {"x": torch.ones((1, 10)), "z": (torch.zeros((1, 10)), torch.ones((1, 5, 5)))}}
def create_pytorch_module_with_nested_inputs3(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, z:Tuple[torch.Tensor, torch.Tensor], x:torch.Tensor):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
return torch.cat([z1, zeros1], 1) + x, torch.cat([z2, zeros2], 2)
net = PTModel()
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
param3 = ov.opset10.parameter(PartialShape([-1, -1]), dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
add = ov.opset10.add(concat1, param3)
ref_model = Model([concat2, add], [param1, param2, param3], "test")
return net, ref_model, {"example_input": {"x": torch.ones((1, 10)), "z": (torch.zeros((1, 10)), torch.ones((1, 5, 3)))}}
def create_pytorch_module_with_nested_inputs4(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, x:torch.Tensor, z:Tuple[torch.Tensor, torch.Tensor], y:torch.Tensor):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
return torch.cat([z1, zeros1], 1) + x, torch.cat([z2, zeros2], 2) * y
net = PTModel()
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
param3 = ov.opset10.parameter(PartialShape([-1, -1]), dtype=np.float32)
param4 = ov.opset10.parameter(PartialShape([-1]), dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
add = ov.opset10.add(concat1, param3)
mul = ov.opset10.multiply(concat2, param4)
ref_model = Model([mul, add], [param3, param1, param2, param4], "test")
return net, ref_model, {"example_input": {"x": torch.ones((1, 10)), "z": (torch.zeros((1, 10)), torch.ones((1, 5, 10))), "y": torch.ones((1,))}}
def create_pytorch_module_with_nested_inputs5(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, x:torch.Tensor, z:Tuple[torch.Tensor, torch.Tensor], y:torch.Tensor):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
return torch.cat([z1, zeros1], 1) + x, torch.cat([z2, zeros2], 2) * y
net = PTModel()
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
param0 = ov.opset10.parameter(PartialShape([-1, -1]), dtype=np.float32)
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
param4 = ov.opset10.parameter(PartialShape([-1]), dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
add = ov.opset10.add(concat1, param0)
mul = ov.opset10.multiply(concat2, param4)
ref_model = Model([mul, add], [param0, param1, param2, param4], "test")
return net, ref_model, {"example_input": [torch.ones((1, 10)), (torch.zeros((1, 10)), torch.ones((1, 5, 10))), torch.ones((1,))]}
def create_pytorch_module_with_nested_inputs6(tmp_dir):
class PTModel(torch.nn.Module):
def forward(self, x:torch.Tensor, y:torch.Tensor=None, z:Tuple[torch.Tensor, torch.Tensor]=None):
z1, z2 = z
zeros1 = torch.zeros((1, 1))
zeros2 = torch.zeros((1, 5, 1))
if y is not None:
return torch.cat([z1, zeros1], 1) * y, torch.cat([z2, zeros2], 2) * y
return torch.cat([z1, zeros1], 1) + x, torch.cat([z2, zeros2], 2)
net = PTModel()
constant_zeros1 = ov.opset10.constant(np.zeros((1, 1), dtype=np.float32), dtype=np.float32)
constant_zeros2 = ov.opset10.constant(np.zeros((1, 5, 1), dtype=np.float32), dtype=np.float32)
shape1 = PartialShape([1, -1])
shape2 = PartialShape([1, 5, -1])
param0 = ov.opset10.parameter(PartialShape([-1, -1]), dtype=np.float32)
param1 = ov.opset10.parameter(shape1, dtype=np.float32)
param2 = ov.opset10.parameter(shape2, dtype=np.float32)
concat1 = ov.opset10.concat([param1, constant_zeros1], 1)
concat2 = ov.opset10.concat([param2, constant_zeros2], 2)
add1 = ov.opset10.add(concat1, param0)
ref_model = Model([concat2, add1], [param0, param1, param2], "test")
return net, ref_model, {"example_input": {"x": torch.ones((1, 11)), "z": (torch.zeros((1, 10)), torch.ones((1, 5, 10)))}}
class TestMoConvertPyTorch(CommonMOConvertTest):
test_data = [
@@ -770,6 +915,12 @@ class TestMoConvertPyTorch(CommonMOConvertTest):
create_pytorch_module_with_optional_inputs_case5,
create_pytorch_nn_module_with_scalar_input,
create_pytorch_module_with_compressed_int8_constant,
create_pytorch_module_with_nested_inputs,
create_pytorch_module_with_nested_inputs2,
create_pytorch_module_with_nested_inputs3,
create_pytorch_module_with_nested_inputs4,
create_pytorch_module_with_nested_inputs5,
create_pytorch_module_with_nested_inputs6
]
@ pytest.mark.parametrize("create_model", test_data)