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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
