tensorflow2线性拟合教程

import os

import tensorflow as tf
import numpy as np

os.environ['CUDA_VISIBLE_DEVICES'] = "0"  # Specify visible gpus.
tf.debugging.set_log_device_placement(True)  # Show the devices when calculating.

x0 = np.array([i * 1.0 for i in range(100)], dtype=np.float64)
y0 = 2.0 * x0 + 3.0


initializer = tf.initializers.GlorotUniform()
W = tf.Variable(initializer(shape=(1, ), dtype=tf.float64), name="W")
b = tf.Variable(initializer(shape=(1, ), dtype=tf.float64), name="b")

optimizer = tf.optimizers.Adam()


def train_step(x_batch, y_batch, epoch, batch_i):
    """
    Note that here we should not call:

    y_predict = W * x_batch + b
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean(tf.math.pow(y_predict - y_batch, 2))

    which will cause "No gradients provided for any variable:" error.
    I don't know the reason but it's better to define tensors within the "tf.GradientTape" block.
    """
    with tf.GradientTape() as tape:
        y_predict = W * x_batch + b
        loss = tf.reduce_mean(tf.math.pow(y_predict - y_batch, 2))
    # print("loss", loss)
    train_variables = [W, b]
    gradients = tape.gradient(loss, train_variables)
    # print("grads", gradients)
    optimizer.apply_gradients(zip(gradients, train_variables))
    print(epoch, batch_i, W, b)


def fit():
    for epoch in range(10000):
        for i in range(0, 100, 10):
            start = i
            end = i + 10
            x_batch = x0[start:end]
            y_batch = y0[start:end]
            train_step(x_batch, y_batch, epoch, i)

fit()