Tensorflow 2.3 profiler only have one step timeline on 'trace_viewer'

[Jianbing-D]

Consider Stack Overflow for getting support using TensorBoard—they have
a larger community with better searchability:

https://stackoverflow.com/questions/tagged/tensorboard

Do not use this template for for setup, installation, or configuration
issues. Instead, use the “installation problem” issue template:

https://github.com/tensorflow/tensorboard/issues/new?template=installation_problem.md

To report a problem with TensorBoard itself, please fill out the
remainder of this template.

Environment information (required)

Please run diagnose_tensorboard.py (link below) in the same
environment from which you normally run TensorFlow/TensorBoard, and
paste the output here:

https://github.com/raw/tensorflow/tensorboard/master/tensorboard/tools/diagnose_tensorboard.py

Diagnostics

Diagnostics output

--- check: autoidentify
INFO: diagnose_tensorboard.py version 724b56cee52e7d8eb89bbeec1f0d5ce3e38c9682

--- check: general
INFO: sys.version_info: sys.version_info(major=3, minor=6, micro=9, releaselevel='final', serial=0)
INFO: os.name: posix
INFO: os.uname(): posix.uname_result(sysname='Linux', nodename='b706f63f5739', release='4.15.0-117-generic', version='#118-Ubuntu SMP Fri Sep 4 20:02:41 UTC 2020', machine='x86_64')
INFO: sys.getwindowsversion(): N/A

--- check: package_management
INFO: has conda-meta: False
INFO: $VIRTUAL_ENV: None

--- check: installed_packages
INFO: installed: tensorboard==2.3.0
INFO: installed: tensorflow-gpu==2.3.1
INFO: installed: tensorflow-estimator==2.3.0

--- check: tensorboard_python_version
INFO: tensorboard.version.VERSION: '2.3.0'

--- check: tensorflow_python_version
2020-10-30 06:55:52.304929: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
INFO: tensorflow.__version__: '2.3.1'
INFO: tensorflow.__git_version__: 'v2.3.0-54-gfcc4b966f1'

--- check: tensorboard_binary_path
INFO: which tensorboard: b'/usr/local/bin/tensorboard\n'

--- check: addrinfos
socket.has_ipv6 = True
socket.AF_UNSPEC = <AddressFamily.AF_UNSPEC: 0>
socket.SOCK_STREAM = <SocketKind.SOCK_STREAM: 1>
socket.AI_ADDRCONFIG = <AddressInfo.AI_ADDRCONFIG: 32>
socket.AI_PASSIVE = <AddressInfo.AI_PASSIVE: 1>
Loopback flags: <AddressInfo.AI_ADDRCONFIG: 32>
Loopback infos: [(<AddressFamily.AF_INET: 2>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('127.0.0.1', 0))]
Wildcard flags: <AddressInfo.AI_PASSIVE: 1>
Wildcard infos: [(<AddressFamily.AF_INET: 2>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('0.0.0.0', 0)), (<AddressFamily.AF_INET6: 10>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('::', 0, 0, 0))]

--- check: readable_fqdn
INFO: socket.getfqdn(): 'b706f63f5739'

--- check: stat_tensorboardinfo
INFO: directory: /tmp/.tensorboard-info
INFO: .tensorboard-info directory does not exist

--- check: source_trees_without_genfiles
INFO: tensorboard_roots (1): ['/usr/local/lib/python3.6/dist-packages']; bad_roots (0): []

--- check: full_pip_freeze
INFO: pip freeze --all:
absl-py==0.10.0
astunparse==1.6.3
cachetools==4.1.1
certifi==2020.6.20
chardet==3.0.4
gast==0.3.3
google-auth==1.22.1
google-auth-oauthlib==0.4.1
google-pasta==0.2.0
grpcio==1.33.1
gviz-api==1.9.0
h5py==2.10.0
idna==2.10
importlib-metadata==2.0.0
joblib==0.17.0
Keras-Preprocessing==1.1.2
Markdown==3.3.3
mpi4py==3.0.3
numpy==1.18.5
nvtx-plugins==0.1.8
oauthlib==3.1.0
opt-einsum==3.3.0
ortools==8.0.8283
pandas==1.1.3
pip==20.2.4
protobuf==3.13.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
python-dateutil==2.8.1
pytz==2020.1
requests==2.24.0
requests-oauthlib==1.3.0
rsa==4.6
scikit-learn==0.23.2
scipy==1.5.3
setuptools==50.3.2
six==1.15.0
sklearn==0.0
tensorboard==2.3.0
tensorboard-plugin-profile==2.3.0
tensorboard-plugin-wit==1.7.0
tensorflow-estimator==2.3.0
tensorflow-gpu==2.3.1
termcolor==1.1.0
threadpoolctl==2.1.0
urllib3==1.25.11
Werkzeug==1.0.1
wheel==0.30.0
wrapt==1.12.1
zipp==3.4.0

Next steps

No action items identified. Please copy ALL of the above output,
including the lines containing only backticks, into your GitHub issue
or comment. Be sure to redact any sensitive information.

For browser-related issues, please additionally specify:

• Browser type and version (e.g., Chrome 64.0.3282.140):
• Screenshot, if it’s a visual issue:

Issue description

I build a model using TF 2.3, and use

with tf.profiler.experimental.Profile("./logdir"):
    for step, datas in enumerate(dataset):
        # some train steps
        if step >= 50:
            break

But when I open the profiler: trace_viwer, it only have one step information rather than 50 steps.

[rmothukuru]

@Jianbing-D,
In order to expedite the trouble-shooting process, please provide a code snippet to reproduce the issue reported here. Thanks!

[Jianbing-D]

Hi, @rmothukuru
You can use following codes to reproduce this issue.
(Actually, my issue is that profiler step number less than my training steps, but following code will produce far more profiler steps than my training steps.)

import tensorflow as tf
import numpy as np

def create_dataset(batch_size):
    dataset = tf.data.Dataset.from_tensor_slices(tf.random.normal([batch_size * 100, 32], dtype=tf.float32))
    dataset = dataset.batch(batch_size)
    dataset = dataset.repeat(2)
    dataset = dataset.map(lambda batch: batch + 1)
    dataset = dataset.prefetch(buffer_size=8)
    return dataset

class Model(tf.keras.models.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.input_layer = tf.keras.layers.Dense(units=1024, activation='relu', use_bias=True)
        self.internel_layer = tf.keras.layers.Dense(units=1024, activation="relu", use_bias=True)
        self.output_layer = tf.keras.layers.Dense(units=1, activation="sigmoid", use_bias=True)
    
    @tf.function
    def call(self, inputs, training=True):
        internel = self.input_layer(inputs)
        for i in tf.range(1024, dtype=tf.int32):
            internel = self.internel_layer(internel)
        return self.output_layer(internel)

@tf.function
def train_step(model, inputs, labels, loss_fn, optimizer):
    with tf.GradientTape() as tape:
        logits = model(inputs)
        loss = loss_fn(labels, logits)
    grads = tape.gradient(loss, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    return loss

if __name__ == "__main__":

    batch_size = 32

    dataset = create_dataset(batch_size)
    
    model = Model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False)


    with tf.profiler.experimental.Profile("./profiler_test"):
        for step, datas in enumerate(dataset):
            labels = np.random.randint(low=0, high=2, size=(batch_size, 1), dtype='l')
            loss = train_step(model, datas, labels, loss_fn, optimizer)
            if step % 10 == 0 and step != 0:
                tf.print("step %d, loss = %.5f" %(step, loss))

            if step >= 50:
                break

[qiuminxu]

Hi, for custom training loops, the step number needs to be annotated so that profiler can pick it up. Can you try the following example and see if it works?
https://www.tensorflow.org/api_docs/python/tf/profiler/experimental/Trace#example_usage