microsoft / microsoft/onnxruntime

Calculating variance uses ~2x more memory compared to tensorflow (CPU)

Open
#5,860 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

contributions welcome
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.