tensorboard2.py

import os
from datetime import datetime
import coloredlogs, logging

import tensorflow as tf
import tensorflow.python.util.deprecation as deprecation

tf.debugging.set_log_device_placement(True)

coloredlogs.install()

deprecation._PRINT_DEPRECATION_WARNINGS = False

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

os.system('rm -rf ./logs/ ')

log_dir = "logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")
writer = tf.summary.create_file_writer(log_dir)
tf.summary.trace_on(graph=True, profiler=True)


@tf.function
def power_with_sig(a, b):
    print('Tracing "power_with_sig"\n')
    return a ** b


p = power_with_sig.get_concrete_function(tf.TensorSpec(shape=[None], dtype=tf.float32),
                                         tf.TensorSpec(shape=[None], dtype=tf.float32))

value = p(tf.constant([5.]), tf.constant([2.]))

with writer.as_default():
    tf.summary.trace_export(
        name="my_func_trace",
        step=0,
        profiler_outdir=log_dir)

print(value.numpy())

t = power_with_sig.get_concrete_function(tf.TensorSpec(shape=[None], dtype=tf.int32),
                                         tf.TensorSpec(shape=[None], dtype=tf.int32))
with tf.device('/GPU:0'):
    value2 = t(tf.constant([[5, 2], [2, 4]]), tf.constant([2, 3]))
    print(value2.numpy())

!tensorboard --logdir logs/fit --host 192.168.68.102