frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add size op in nnapi serializer (#52026)

Browse Source 
Summary:
serializer didn't support aten::size

Pull Request resolved: https://github.com/pytorch/pytorch/pull/52026

Test Plan: Torchvision Mobilenetv2 [script](https://pytorch.org/tutorials/prototype/nnapi_mobilenetv2.html) works. [Test](ecfed07cc5) to be merged after [this PR](https://github.com/pytorch/pytorch/pull/47521/files) is merged

Reviewed By: dreiss

Differential Revision: D26363133

Pulled By: axitkhurana

fbshipit-source-id: 772a6bea62bca69f8bba19c25c582a1734a70eb1
This commit is contained in:
Akshit Khurana
2021-02-10 15:48:35 -08:00
committed by Facebook GitHub Bot
parent a1b8f3d4b6
commit d0fd41dcfe
1 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
torch/backends/_nnapi/serializer.py
View File

@ -593,6 +593,8 @@ class _NnapiSerializer(object):
self.add_tuple_construct(node),
"aten::reshape": lambda self, node:
self.add_reshape(node),
"aten::size": lambda self, node:
self.add_size(node),
"aten::quantize_per_tensor": lambda self, node:
self.add_quantize(node),
"aten::dequantize": lambda self, node:
@ -706,6 +708,16 @@ class _NnapiSerializer(object):
self.add_operation(NNAPI_OperationCode.RESHAPE, inputs, outputs)
def add_size(self, node):
assert node.inputsSize() == 2
assert node.outputsSize() == 1
_, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
_, value = self.constants[node.inputsAt(1)]
res = in_oper.shape[value]
output = node.outputsAt(0)
self.add_constant_value(output, output.type(), res)
def add_quantize(self, node):
assert node.inputsSize() == 4
assert node.outputsSize() == 1