在keras变量发生变化时执行预计算

import tensorflow as tf

@tf.function
def compute(x, y):
    z = tf.matmul(x, y)
    return z

x = tf.constant([[1, 2], [3, 4]])
y = tf.constant([[5, 6], [7, 8]])

result = compute(x, y)
print(result)

import tensorflow as tf
from tensorflow import keras

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

# 定义模型、损失函数和优化器
model = keras.Sequential([...])
loss_fn = keras.losses.SparseCategoricalCrossentropy()
optimizer = keras.optimizers.Adam()

# 训练循环
for epoch in range(num_epochs):
    for batch in dataset:
        inputs, labels = batch
        loss = train_step(inputs, labels, model, loss_fn, optimizer)
        # 执行预计算，更新模型参数