microsoft / microsoft/onnxruntime
Calculating variance uses ~2x more memory compared to tensorflow (CPU)
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
**Describe the bug**
I noticed this when comparing the memory profiles of a TF1 (1.13.1) model that uses `layer_norm`. After converting the graph to onnx, memory usage seems to double.
I was able to pin it down to the call to `tf.squared_difference` when computing the variance (`tf.nn.moments`). It seems that this is a primitive op in tensorflow, but in onnx it is broken into the primitives (Sub, Square). How come TF's is much more memory efficient -- maybe because of the implementation as a single op instead of two primitives (e.g. extra copies somewhere)?
**Urgency**
If there are particular important use cases blocked by this or strict project-related timelines, please share more information and dates. If there are no hard deadlines, please specify none.
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
- ONNX Runtime installed from (source or binary): binary
- ONNX Runtime version: 1.5.1
- Python version: 3.6.10
- Visual Studio version (if applicable):
- GCC/Compiler version (if compiling from source):
- CUDA/cuDNN version: n/a
- GPU model and memory: n/a
**To Reproduce**
- Describe steps/code to reproduce the behavior.
`pip install tensorflow==1.13.1 onnxruntime tf2onnx memory_profiler`
```
import timeit
import numpy as np
import onnxruntime as ort
import warnings
warnings.filterwarnings('ignore',category=FutureWarning)
import tensorflow as tf
from tf2onnx import tfonnx, optimizer
from memory_profiler import profile
BATCH_SIZE = 32
IMG_SIZE = 128
FILTERS = 200
CHANNELS = 500
NUMBER = 1
x = tf.placeholder(tf.float32, [BATCH_SIZE, IMG_SIZE, IMG_SIZE, CHANNELS])
y = tf.identity(tf.squared_difference(x, 1), name="y")
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)
rand_in = np.random.random(x.shape).astype(np.float32)
@profile
def run_tf(sess):
return sess.run(y, feed_dict={x: rand_in})
print(timeit.timeit(lambda: run_tf(sess), number=NUMBER))
# convert to onnx
frozen_graph = tf.graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), [y.name.split(':')[0]])
tf.reset_default_graph()
tf.import_graph_def(frozen_graph, name="") # by default, this will prefix 'import' to all the names
onnx_graph = tfonnx.process_tf_graph(tf.get_default_graph(), opset=12, input_names=[x.name], output_names=[y.name])
onnx_graph = optimizer.optimize_graph(onnx_graph)
model_proto = onnx_graph.make_model("onnx_model")
with open("/tmp/model.onnx", "wb") as f:
f.write(model_proto.SerializeToString())
ort_sess = ort.InferenceSession(model_proto.SerializeToString())
@profile
def run_onnx(ort_sess):
# needs to be np.float32 otherwise ort gives this nice message:
# 'Actual: (N11onnxruntime17PrimitiveDataTypeIdEE) , expected: (N11onnxruntime17PrimitiveDataTypeIfEE)'
return ort_sess.run([y.name], {x.name: rand_in})
print(timeit.timeit(lambda: run_onnx(ort_sess), number=NUMBER))
```
- Attach the ONNX model to the issue (where applicable) to expedite investigation.
**Expected behavior**
A clear and concise description of what you expected to happen.
**Screenshots**
If applicable, add screenshots to help explain your problem.
**Additional context**
Add any other context about the problem here. If the issue is about a particular model, please share the model details as well to facilitate debugging.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Run the supplied Python reproduction comparing TensorFlow's tf.squared_difference with the converted ONNX Sub and Square operations, using the stated CPU and memory-profiler setup. Start by tracing tf.nn.moments and the generated graph through ONNX Runtime's execution path; done means explaining the excess memory and demonstrating that the CPU memory usage is reduced or the limitation is documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, tensorflow
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100