Skip to content

GitLab

Commit beec3163 authored by Fan Yang's avatar Fan Yang Committed by A. Unique TensorFlower
Browse files

Internal change. 

PiperOrigin-RevId: 409277517
parent d022a749
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 2 deletions
official/vision/beta/serving/export_saved_model_lib_test.py
View file @ beec3163
...@@ -29,11 +29,16 @@ class WriteModelFlopsAndParamsTest(tf.test.TestCase): ...@@ -29,11 +29,16 @@ class WriteModelFlopsAndParamsTest(tf.test.TestCase):
@mock.patch.object(export_base, 'export', autospec=True, spec_set=True) @mock.patch.object(export_base, 'export', autospec=True, spec_set=True)
def test_retinanet_task(self, unused_export): def test_retinanet_task(self, unused_export):
tempdir = self.create_tempdir() tempdir = self.create_tempdir()
params = configs.retinanet.retinanet_resnetfpn_coco()
print(params.task.model.backbone)
params.task.model.backbone.resnet.model_id = 18
params.task.model.num_classes = 2
params.task.model.max_level = 6
export_saved_model_lib.export_inference_graph( export_saved_model_lib.export_inference_graph(
input_type='image_tensor', input_type='image_tensor',
batch_size=1, batch_size=1,
input_image_size=[128, 128], input_image_size=[64, 64],
params=configs.retinanet.retinanet_resnetfpn_coco(), params=params,
checkpoint_path=os.path.join(tempdir, 'unused-ckpt'), checkpoint_path=os.path.join(tempdir, 'unused-ckpt'),
export_dir=tempdir, export_dir=tempdir,
log_model_flops_and_params=True) log_model_flops_and_params=True)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment