onnx/onnx-tensorflow

after .export_graph, using c version, load savemodel error.

zuixiaosanlang opened this issue · 0 comments

class JITWrapper_tf_2(nn.Module):
def init(self):
super(JITWrapper_tf_2, self).init()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=5, stride=1, padding=2)

def forward(self, image):
    image_t = image.permute(0, 3, 1, 2)
    out = self.conv1(image_t)

    out = out.permute(0, 2, 3, 1)

    return out

def save_model_lama():
jITWrapper_tf_2 = JITWrapper_tf_2()

input_np = np.random.uniform(0, 1, (1, 512, 512, 3)).astype(np.float32)
img_tf = torch.from_numpy(input_np)

img_tf = Variable(img_tf)
out = jITWrapper_tf_2(img_tf)

onnx_name = "lama_12.onnx"
torch.onnx.export(jITWrapper_tf_2,  # model being run
                  img_tf,
                  # model input (or a tuple for multiple inputs)
                  onnx_name,  # where to save the model (can be a file or file-like object)
                  export_params=True,  # store the trained parameter weights inside the model file
                  opset_version=11,  # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names=["image"],  # the model's input names
                  output_names=["out"],  # the model's output names
                  dynamic_axes={'image': [0, 1, 2], 'out': [0, 1, 2]},
                  verbose=True
                  )

onnx_model = onnx.load(onnx_name)  # load onnx model
tf_model = prepare(onnx_model)

tf_model.export_graph("tf_saved_model_lama")  # export the model

save_model_lama()

and:
image

when call by tensorflow_c:
image

i donnot know what is the 'saver_filename'.

env:
tensorflow==2.3.0
torch==1.7.1+cu101
onnx==1.10.2
onnx-tf==1.10.0