after .export_graph, using c version, load savemodel error.
zuixiaosanlang opened this issue · 0 comments
zuixiaosanlang commented
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()
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