未验证 提交 1b4bcd61 编写于 作者: L lzzyzlbb 提交者: GitHub

modify export_model (#545)

上级 541e71bc
...@@ -214,7 +214,7 @@ class FirstOrderModel(BaseModel): ...@@ -214,7 +214,7 @@ class FirstOrderModel(BaseModel):
kp_driving=kp_norm) kp_driving=kp_norm)
return out['prediction'] return out['prediction']
def export_model(self, export_model=None, output_dir=None, inputs_size=[]): def export_model(self, export_model=None, output_dir=None, inputs_size=[], export_serving_model=False):
source = paddle.rand(shape=inputs_size[0], dtype='float32') source = paddle.rand(shape=inputs_size[0], dtype='float32')
driving = paddle.rand(shape=inputs_size[1], dtype='float32') driving = paddle.rand(shape=inputs_size[1], dtype='float32')
......
...@@ -309,7 +309,8 @@ class StyleGAN2Model(BaseModel): ...@@ -309,7 +309,8 @@ class StyleGAN2Model(BaseModel):
def export_model(self, def export_model(self,
export_model=None, export_model=None,
output_dir=None, output_dir=None,
inputs_size=[[1, 1, 512], [1, 1]]): inputs_size=[[1, 1, 512], [1, 1]],
export_serving_model=False):
infer_generator = self.InferGenerator() infer_generator = self.InferGenerator()
infer_generator.set_generator(self.nets['gen']) infer_generator.set_generator(self.nets['gen'])
style = paddle.rand(shape=inputs_size[0], dtype='float32') style = paddle.rand(shape=inputs_size[0], dtype='float32')
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册