From efb75f26b32c42fe500fa8172e411981ceb5ab69 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Mon, 16 Jan 2023 17:01:46 +0800 Subject: [PATCH] Add torch no_grad in inference Signed-off-by: Jael Gu --- test_onnx.py | 2 +- timm_image.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/test_onnx.py b/test_onnx.py index ab7f718..55650a7 100644 --- a/test_onnx.py +++ b/test_onnx.py @@ -45,7 +45,7 @@ for name in models: onnx_path = f'saved/onnx/{saved_name}.onnx' try: - op = TimmImage(model_name=name, device='cpu') + op = ops.image_embedding.timm(model_name=name, device='cpu').get_op() data = torch.rand((1,) + op.config['input_size']) except Exception as e: print(f'***Please re-download model {name}.***') diff --git a/timm_image.py b/timm_image.py index 3d68bb3..ef01ebc 100644 --- a/timm_image.py +++ b/timm_image.py @@ -41,6 +41,13 @@ warnings.filterwarnings('ignore') log = logging.getLogger('timm_op') +def torch_no_grad(f): + def wrap(*args, **kwargs): + with torch.no_grad(): + f(*args, **kwargs) + return wrap + + # @accelerate class Model: def __init__(self, model_name, device, num_classes): @@ -88,6 +95,7 @@ class TimmImage(NNOperator): log.warning('The operator is initialized without specified model.') pass + @torch_no_grad def __call__(self, data: Union[List[towhee._types.Image], towhee._types.Image]): if not isinstance(data, list): imgs = [data] @@ -270,3 +278,12 @@ class TimmImage(NNOperator): return ['onnx'] else: return [] + + def input_schema(self): + return [(Image, (-1, -1, 3))] + + def output_schema(self): + image = Image(numpy.random.randn(480, 480, 3), "RGB") + ret = self(image) + data_type = type(ret.reshape(-1)[0]) + return [(data_type, ret.shape)]