zhaoyang-star commented on PR #11208: URL: https://github.com/apache/tvm/pull/11208#issuecomment-1276903736
Hi @lhutton1 , thanks for your contributition. After running FuseOps pass, I want to get the memory usage per op or per primitive func by `AnnotateUsedMemory` pass for furture optimization. I get a resnet18 ir model, then put it as the input IRModule of `AnnotateUsedMemory` pass. The output IRModule has no `used_memory` attr. Test code as follow: ```python import pytest from collections import OrderedDict import numpy as np import tvm from tvm import relay from tvm.relay import testing def AnnotateUsedMemory(): return relay.transform._ffi_api.AnnotateUsedMemory() def _get_data(in_data_shapes, dtype="float32"): in_data = OrderedDict() for name, shape in in_data_shapes.items(): in_data[name] = np.random.uniform(size=shape).astype(dtype) return in_data def _run_relay(mod, params, in_data, pass_enabled): target = "llvm" dev = tvm.device("llvm", 0) in_data = [tvm.nd.array(value) for value in in_data.values()] if pass_enabled: mod = relay.transform.InferType()(mod) mod = relay.transform.ToANormalForm()(mod) mod = relay.transform.InferType()(mod) mod = AnnotateUsedMemory()(mod) # create primitive functions mod = relay.transform.FuseOps()(mod) print(f'\nmod when AnnotateUsedMemory is {pass_enabled}:\n {mod}') out_data = relay.create_executor( "graph", mod, device=dev, target=target).evaluate()(*in_data, **params) return out_data.numpy() def _verify_results(mod, params, in_data, rtol=1e-5, atol=1e-5): before = _run_relay(mod, params, in_data, False) after = _run_relay(mod, params, in_data, True) np.testing.assert_allclose(before, after, rtol, atol) def test_resnet(): num_class = 1000 in_data_shapes = OrderedDict({"data": (1, 3, 224, 224)}) in_data = _get_data(in_data_shapes, dtype="float32") for n in [18]: # 18, 34, 50, 101 mod, params = tvm.relay.testing.resnet.get_workload( batch_size=1, num_classes=num_class, num_layers=n) _verify_results(mod, params, in_data) if __name__ == "__main__": pytest.main([__file__]) ``` I am not familar with `AnnotateUsedMemory` pass. Could memory usage per op or per primitive func be gotten by your pass? If not, how to get it based on your pass? Thanks in advance ^_^ -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org