在TensorFlow中实现梯度下降，而不是使用它提供的

import tensorflow as tf

# 定义模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(units=10, input_shape=(784,), activation='softmax')
])

# 定义损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

# 定义优化器
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# 定义训练步骤
@tf.function
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        # 计算损失函数
        logits = model(inputs)
        loss_value = loss_fn(labels, logits)
    
    # 计算梯度
    gradients = tape.gradient(loss_value, model.trainable_variables)
    
    # 更新参数
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# 进行训练
for epoch in range(num_epochs):
    for batch_inputs, batch_labels in train_dataset:
        train_step(batch_inputs, batch_labels)