tf.GradientTape() 示例

tf.GradientTape() 示例

简单写写，代码请自行补全

# 定义模型
model = keras.Sequential([ 
    layers.Dense(256, activation='relu'),
    layers.Dense(10)])

# 定义优化器
optimizer = optimizers.SGD(learning_rate=0.001)

# 定义损失函数
def loss(x,y):
	loss = tf.reduce_sum(tf.square(out - y)) / x.shape[0]
	return loss

# 定义梯度函数
@tf.function # 加速函数
def grads(inputs,label):
    with tf.GradientTape() as tape:  # 自动求导
        predictions = model(inputs)  # 预测值
        loss = loss(
            labels,  # 标签
            predictions  # 预测值
        )  # 损失值
    grads = tape.gradient(  # 梯度
        loss,  # 损失值
        model.trainable_variables  # 模型变量（自动更新） 
    )
    return grads

# 训练模型
for i in epoch:
    # 计算梯度
    grads = grads(
        inputs,  # 输入数据
        label  # 标签数据
    )
    #  反向传播，调整参数
    optimizer.apply_gradients(zip(
        grads,   # 梯度值
        model.trainable_variables  # 模型变量（自动更新） 
    ))