用GradientTape实现TensorFlow 2.0线性回归中的爆炸损失

import tensorflow as tf
import numpy as np

# 生成随机数据
x = np.random.rand(100).astype(np.float32)
y = 2 * x + 1

# 定义模型
class LinearRegression(tf.keras.Model):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.W = tf.Variable(0.0)
        self.b = tf.Variable(0.0)

    def call(self, inputs):
        return self.W * inputs + self.b

# 定义损失函数
def loss_fn(model, inputs, targets):
    predictions = model(inputs)
    return tf.reduce_mean(tf.square(predictions - targets))

# 定义训练步骤
def train_step(model, inputs, targets, optimizer):
    with tf.GradientTape() as tape:
        loss_value = loss_fn(model, inputs, targets)
    gradients = tape.gradient(loss_value, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# 创建模型和优化器
model = LinearRegression()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# 执行训练
for epoch in range(100):
    train_step(model, x, y, optimizer)