Create data with op graphs and profile data. The `profile` run has tags `prof_f` with both profile and op graph data (`graph_run_metadata` plugin), and `prof_g` with profile data only (`graph_run_metadata_graph` plugin).
()
| 104 | |
| 105 | |
| 106 | def profile(): |
| 107 | """Create data with op graphs and profile data. |
| 108 | |
| 109 | The `profile` run has tags `prof_f` with both profile and op graph data |
| 110 | (`graph_run_metadata` plugin), and `prof_g` with profile data only |
| 111 | (`graph_run_metadata_graph` plugin). |
| 112 | """ |
| 113 | |
| 114 | logdir = os.path.join(LOGDIR, "profile") |
| 115 | |
| 116 | @tf.function |
| 117 | def f(i): |
| 118 | return i + i |
| 119 | |
| 120 | @tf.function |
| 121 | def g(i): |
| 122 | return i * i |
| 123 | |
| 124 | with tf.summary.create_file_writer(logdir).as_default(): |
| 125 | for step in range(3): |
| 126 | # Suppress the profiler deprecation warnings from tf.summary.trace_*. |
| 127 | with _silence_deprecation_warnings(): |
| 128 | tf.summary.trace_on(profiler=True, profiler_outdir=logdir) |
| 129 | print(f(tf.constant(step)).numpy()) |
| 130 | tf.summary.trace_export("prof_f", step=step) |
| 131 | |
| 132 | tf.summary.trace_on(profiler=False) |
| 133 | print(g(tf.constant(step)).numpy()) |
| 134 | tf.summary.trace_export("prof_g", step=step) |
| 135 | |
| 136 | |
| 137 | def main(): |