源代码：
import tensorflow as tf
from B_MLP_CNN import MLP
from B_MLP_CNN import MNISTLoader
# 设置仅在需要时申请显存空间
gpus = tf.config.list_physical_devices (device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth (device=gpu, enable=True)
num_batches = 1000
batch_size = 50
learning_rate = 0.001
log_dir = 'tensorboard'
model = MLP ()
data_loader = MNISTLoader ()
optimizer = tf.keras.optimizers.Adam (learning_rate=learning_rate)
summary_writer = tf.summary.create_file_writer (log_dir) # 实例化记录器
tf.summary.trace_on (graph = True, profiler=True) # 开启 Trace
for batch_index in range (num_batches):
X, y = data_loader.get_batch (batch_size)
with tf.GradientTape () as tape:
y_pred = model (X)
loss = tf.keras.losses.sparse_categorical_crossentropy (y_true=y, y_pred=y_pred)
loss = tf.reduce_mean (loss)
print ("batch %d: loss %f" % (batch_index, loss.numpy ()))
with summary_writer.as_default (): # 指定记录器
tf.summary.scalar ("loss", loss, step=batch_index) # 将当前损失函数的值写入记录器
grads = tape.gradient (loss, model.variables)
optimizer.apply_gradients (grads_and_vars=zip (grads, model.variables))
with summary_writer.as_default ():
tf.summary.trace_export (name="model_trace", step=0, profiler_outdir=log_dir) # 保存 Trace 信息到文件
报错提示：
tensorflow-gpu：2.3.0，tensorboard：2.3.0