Graph Execution Mode in TensorFlow 2

尽管 TensorFlow 2 建议以即时执行模式(Eager Execution)作为主要执行模式,然而,图执行模式(Graph Execution)作为 TensorFlow 2 之前的主要执行模式,依旧对于我们理解 TensorFlow 具有重要意义。尤其是当我们需要使用 tf.function 时,对图执行模式的理解更是不可或缺。

图执行模式在 TensorFlow 1.X 和 2.X 版本中的 API 不同:

  • 在 TensorFlow 1.X 中,图执行模式主要通过“直接构建计算图 + tf.Session” 进行操作;

  • 在 TensorFlow 2 中,图执行模式主要通过 tf.function 进行操作。

在本章,我们将在 tf.function:图执行模式 一节的基础上,进一步对图执行模式的这两种 API 进行对比说明,以帮助已熟悉 TensorFlow 1.X 的用户过渡到 TensorFlow 2。

提示

TensorFlow 2 依然支持 TensorFlow 1.X 的 API。为了在 TensorFlow 2 中使用 TensorFlow 1.X 的 API ,我们可以使用 import tensorflow.compat.v1 as tf 导入 TensorFlow,并通过 tf.disable_eager_execution() 禁用默认的即时执行模式。

TensorFlow 1+1

TensorFlow 的图执行模式是一个符号式的(基于计算图的)计算框架。简而言之,如果你需要进行一系列计算,则需要依次进行如下两步:

  • 建立一个“计算图”,这个图描述了如何将输入数据通过一系列计算而得到输出;

  • 建立一个会话,并在会话中与计算图进行交互,即向计算图传入计算所需的数据,并从计算图中获取结果。

使用计算图进行基本运算

这里以计算 1+1 作为 Hello World 的示例。以下代码通过 TensorFlow 1.X 的图执行模式 API 计算 1+1:

import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

# 以下三行定义了一个简单的“计算图”
a = tf.constant(1)  # 定义一个常量张量(Tensor)
b = tf.constant(1)
c = a + b           # 等价于 c = tf.add(a, b),c是张量a和张量b通过 tf.add 这一操作(Operation)所形成的新张量
# 到此为止,计算图定义完毕,然而程序还没有进行任何实质计算。
# 如果此时直接输出张量 c 的值,是无法获得 c = 2 的结果的

sess = tf.Session()     # 实例化一个会话(Session)
c_ = sess.run(c)        # 通过会话的 run() 方法对计算图里的节点(张量)进行实际的计算
print(c_)

输出:

2

而在 TensorFlow 2 中,我们将计算图的建立步骤封装在一个函数中,并使用 @tf.function 修饰符对函数进行修饰。当需要运行此计算图时,只需调用修饰后的函数即可。由此,我们可以将以上代码改写如下:

import tensorflow as tf

# 以下被 @tf.function 修饰的函数定义了一个计算图
@tf.function
def graph():
    a = tf.constant(1)
    b = tf.constant(1)
    c = a + b
    return c
# 到此为止,计算图定义完毕。由于 graph() 是一个函数,在其被调用之前,程序是不会进行任何实质计算的。
# 只有调用函数,才能通过函数返回值,获得 c = 2 的结果

c_ = graph()
print(c_.numpy())

小结

  • 在 TensorFlow 1.X 的 API 中,我们直接在主程序中建立计算图。而在 TensorFlow 2 中,计算图的建立需要被封装在一个被 @tf.function 修饰的函数中;

  • 在 TensorFlow 1.X 的 API 中,我们通过实例化一个 tf.Session ,并使用其 run 方法执行计算图的实际运算。而在 TensorFlow 2 中,我们通过直接调用被 @tf.function 修饰的函数来执行实际运算。

计算图中的占位符与数据输入

上面这个程序只能计算1+1,以下代码通过 TensorFlow 1.X 的图执行模式 API 中的 tf.placeholder() (占位符张量)和 sess.run()feed_dict 参数,展示了如何使用TensorFlow计算任意两个数的和:

import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

a = tf.placeholder(dtype=tf.int32)  # 定义一个占位符Tensor
b = tf.placeholder(dtype=tf.int32)
c = a + b

a_ = int(input("a = "))  # 从终端读入一个整数并放入变量a_
b_ = int(input("b = "))

sess = tf.Session()
c_ = sess.run(c, feed_dict={a: a_, b: b_})  # feed_dict参数传入为了计算c所需要的张量的值
print("a + b = %d" % c_)

运行程序:

>>> a = 2
>>> b = 3
a + b = 5

而在 TensorFlow 2 中,我们可以通过为函数指定参数来实现与占位符张量相同的功能。为了在计算图运行时送入占位符数据,只需在调用被修饰后的函数时,将数据作为参数传入即可。由此,我们可以将以上代码改写如下:

import tensorflow as tf

@tf.function
def graph(a, b):
    c = a + b
    return c

a_ = int(input("a = "))
b_ = int(input("b = "))
c_ = graph(a_, b_)
print("a + b = %d" % c_)

小结

在 TensorFlow 1.X 的 API 中,我们使用 tf.placeholder() 在计算图中声明占位符张量,并通过 sess.run()feed_dict 参数向计算图中的占位符传入实际数据。而在 TensorFlow 2 中,我们使用 tf.function 的函数参数作为占位符张量,通过向被 @tf.function 修饰的函数传递参数,来为计算图中的占位符张量提供实际数据。

计算图中的变量

变量的声明

变量 (Variable)是一种特殊类型的张量,在 TensorFlow 1.X 的图执行模式 API 中使用 tf.get_variable() 建立。与编程语言中的变量很相似。使用变量前需要先初始化,变量内存储的值可以在计算图的计算过程中被修改。以下示例代码展示了如何建立一个变量,将其值初始化为0,并逐次累加1。

import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

a = tf.get_variable(name='a', shape=[])
initializer = tf.assign(a, 0.0)   # tf.assign(x, y)返回一个“将张量y的值赋给变量x”的操作
plus_one_op = tf.assign(a, a + 1.0)

sess = tf.Session()
sess.run(initializer)
for i in range(5):
    sess.run(plus_one_op)       # 对变量a执行加一操作
    print(sess.run(a))          # 输出此时变量a在当前会话的计算图中的值

输出:

1.0
2.0
3.0
4.0
5.0

提示

为了初始化变量,也可以在声明变量时指定初始化器(initializer),并通过 tf.global_variables_initializer() 一次性初始化所有变量,在实际工程中更常用:

import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

a = tf.get_variable(name='a', shape=[], 
    initializer=tf.zeros_initializer)   # 指定初始化器为全0初始化
plus_one_op = tf.assign(a, a + 1.0)

sess = tf.Session()
sess.run(tf.global_variables_initializer()) # 初始化所有变量
for i in range(5):
    sess.run(plus_one_op)
    print(sess.run(a))

在 TensorFlow 2 中,我们通过实例化 tf.Variable 类来声明变量。由此,我们可以将以上代码改写如下:

import tensorflow as tf

a = tf.Variable(0.0)

@tf.function
def plus_one_op():
    a.assign(a + 1.0)
    return a

for i in range(5):
    plus_one_op()
    print(a.numpy())

小结

在 TensorFlow 1.X 的 API 中,我们使用 tf.get_variable() 在计算图中声明变量节点。而在 TensorFlow 2 中,我们直接通过 tf.Variable 实例化变量对象,并在计算图中使用这一变量对象。

变量的作用域与重用

在 TensorFlow 1.X 中,我们建立模型时经常需要指定变量的作用域,以及复用变量。此时,TensorFlow 1.X 的图执行模式 API 为我们提供了 tf.variable_scope()reuse 参数来实现变量作用域和复用变量的功能。以下的例子使用了 TensorFlow 1.X 的图执行模式 API 建立了一个三层的全连接神经网络,其中第三层复用了第二层的变量。

import tensorflow.compat.v1 as tf
import numpy as np
tf.disable_eager_execution()

def dense(inputs, num_units):
    weight = tf.get_variable(name='weight', shape=[inputs.shape[1], num_units])
    bias = tf.get_variable(name='bias', shape=[num_units])
    return tf.nn.relu(tf.matmul(inputs, weight) + bias)

def model(inputs):
    with tf.variable_scope('dense1'):   # 限定变量的作用域为 dense1
        x = dense(inputs, 10)           # 声明了 dense1/weight 和 dense1/bias 两个变量
    with tf.variable_scope('dense2'):   # 限定变量的作用域为 dense2
        x = dense(x, 10)                # 声明了 dense2/weight 和 dense2/bias 两个变量
    with tf.variable_scope('dense2', reuse=True):   # 第三层复用第二层的变量
        x = dense(x, 10)
    return x

inputs = tf.placeholder(shape=[10, 32], dtype=tf.float32)
outputs = model(inputs)
print(tf.global_variables())    # 输出当前计算图中的所有变量节点
sess = tf.Session()
sess.run(tf.global_variables_initializer())
outputs_ = sess.run(outputs, feed_dict={inputs: np.random.rand(10, 32)})
print(outputs_)

在上例中,计算图的所有变量节点为:

[<tf.Variable 'dense1/weight:0' shape=(32, 10) dtype=float32>,
 <tf.Variable 'dense1/bias:0' shape=(10,) dtype=float32>,
 <tf.Variable 'dense2/weight:0' shape=(10, 10) dtype=float32>,
 <tf.Variable 'dense2/bias:0' shape=(10,) dtype=float32>]

可见, tf.variable_scope() 为在其上下文中的,以 tf.get_variable 建立的变量的名称添加了“前缀”或“作用域”,使得变量在计算图中的层次结构更为清晰,不同“作用域”下的同名变量各司其职,不会冲突。同时,虽然我们在上例中调用了3次 dense 函数,即调用了6次 tf.get_variable 函数,但实际建立的变量节点只有4个。这即是 tf.variable_scope()reuse 参数所起到的作用。当 reuse=True 时, tf.get_variable 遇到重名变量时将会自动获取先前建立的同名变量,而不会新建变量,从而达到了变量重用的目的。

而在 TensorFlow 2 的图执行模式 API 中,不再鼓励使用 tf.variable_scope() ,而应当使用 tf.keras.layers.Layertf.keras.Model 来封装代码和指定作用域,具体可参考 本手册第三章。上面的例子与下面基于 tf.kerastf.function 的代码等价。

import tensorflow as tf
import numpy as np

class Dense(tf.keras.layers.Layer):
    def __init__(self, num_units, **kwargs):
        super().__init__(**kwargs)
        self.num_units = num_units

    def build(self, input_shape):
        self.weight = self.add_variable(name='weight', shape=[input_shape[-1], self.num_units])
        self.bias = self.add_variable(name='bias', shape=[self.num_units])

    def call(self, inputs):
        y_pred = tf.matmul(inputs, self.weight) + self.bias
        return y_pred

class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = Dense(num_units=10, name='dense1')
        self.dense2 = Dense(num_units=10, name='dense2')

    @tf.function
    def call(self, inputs):
        x = self.dense1(inputs)
        x = self.dense2(inputs)
        x = self.dense2(inputs)
        return x

model = Model()
print(model(np.random.rand(10, 32)))

我们可以注意到,在 TensorFlow 2 中,变量的作用域以及复用变量的问题自然地淡化了。基于Python类的模型建立方式自然地为变量指定了作用域,而变量的重用也可以通过简单地多次调用同一个层来实现。

为了详细了解上面的代码对变量作用域的处理方式,我们使用 get_concrete_function 导出计算图,并输出计算图中的所有变量节点:

graph = model.call.get_concrete_function(np.random.rand(10, 32))
print(graph.variables)

输出如下:

(<tf.Variable 'dense1/weight:0' shape=(32, 10) dtype=float32, numpy=...>,
 <tf.Variable 'dense1/bias:0' shape=(10,) dtype=float32, numpy=...>,
 <tf.Variable 'dense2/weight:0' shape=(32, 10) dtype=float32, numpy=...>,
 <tf.Variable 'dense2/bias:0' shape=(10,) dtype=float32, numpy=...)

可见,TensorFlow 2 的图执行模式在变量的作用域上与 TensorFlow 1.X 实际保持了一致。我们通过 name 参数为每个层指定的名称将成为层内变量的作用域。

小结

在 TensorFlow 1.X 的 API 中,使用 tf.variable_scope()reuse 参数来实现变量作用域和复用变量的功能。在 TensorFlow 2 中,使用 tf.keras.layers.Layertf.keras.Model 来封装代码和指定作用域,从而使变量的作用域以及复用变量的问题自然淡化。两者的实质是一样的。

自动求导机制与优化器

在本节中,我们对 TensorFlow 1.X 和 TensorFlow 2 在图执行模式下的自动求导机制进行较深入的比较说明。

自动求导机制

我们首先回顾 TensorFlow 1.X 中的自动求导机制。在 TensorFlow 1.X 的图执行模式 API 中,可以使用 tf.gradients(y, x) 计算计算图中的张量节点 y 相对于变量 x 的导数。以下示例展示了在 TensorFlow 1.X 的图执行模式 API 中计算 y = x^2x = 3 时的导数。

x = tf.get_variable('x', dtype=tf.float32, shape=[], initializer=tf.constant_initializer(3.))
y = tf.square(x)    # y = x ^ 2
y_grad = tf.gradients(y, x)

以上代码中,计算图中的节点 y_grad 即为 y 相对于 x 的导数。

而在 TensorFlow 2 的图执行模式 API 中,我们使用 tf.GradientTape 这一上下文管理器封装需要求导的计算步骤,并使用其 gradient 方法求导,代码示例如下:

x = tf.Variable(3.)
@tf.function
def grad():
    with tf.GradientTape() as tape:
        y = tf.square(x)
    y_grad = tape.gradient(y, x)
    return y_grad

小结

在 TensorFlow 1.X 中,我们使用 tf.gradients() 求导。而在 TensorFlow 2 中,我们使用使用 tf.GradientTape 这一上下文管理器封装需要求导的计算步骤,并使用其 gradient 方法求导。

优化器

由于机器学习中的求导往往伴随着优化,所以 TensorFlow 中更常用的是优化器(Optimizer)。在 TensorFlow 1.X 的图执行模式 API 中,我们往往使用 tf.train 中的各种优化器,将求导和调整变量值的步骤合二为一。例如,以下代码片段在计算图构建过程中,使用 tf.train.GradientDescentOptimizer 这一梯度下降优化器优化损失函数 loss

y_pred = model(data_placeholder)    # 模型构建
loss = ...                          # 计算模型的损失函数 loss
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
train_one_step = optimizer.minimize(loss)
# 上面一步也可拆分为
# grad = optimizer.compute_gradients(loss)
# train_one_step = optimizer.apply_gradients(grad)

以上代码中, train_one_step 即为一个将求导和变量值更新合二为一的计算图节点(操作),也就是训练过程中的“一步”。特别需要注意的是,对于优化器的 minimize 方法而言,只需要指定待优化的损失函数张量节点 loss 即可,求导的变量可以自动从计算图中获得(即 tf.trainable_variables )。在计算图构建完成后,只需启动会话,使用 sess.run 方法运行 train_one_step 这一计算图节点,并通过 feed_dict 参数送入训练数据,即可完成一步训练。代码片段如下:

for data in dataset:
    data_dict = ... # 将训练所需数据放入字典 data 内
    sess.run(train_one_step, feed_dict=data_dict)

而在 TensorFlow 2 的 API 中,无论是图执行模式还是即时执行模式,均先使用 tf.GradientTape 进行求导操作,然后再使用优化器的 apply_gradients 方法应用已求得的导数,进行变量值的更新。也就是说,和 TensorFlow 1.X 中优化器的 compute_gradients + apply_gradients 十分类似。同时,在 TensorFlow 2 中,无论是求导还是使用导数更新变量值,都需要显式地指定变量。计算图的构建代码结构如下:

optimizer = tf.keras.optimizer.SGD(learning_rate=...)

@tf.function
def train_one_step(data):
    with tf.GradientTape() as tape:
        y_pred = model(data)    # 模型构建
        loss = ...              # 计算模型的损失函数 loss
    grad = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))

在计算图构建完成后,我们直接调用 train_one_step 函数并送入训练数据即可:

for data in dataset:
    train_one_step(data)

小结

在 TensorFlow 1.X 中,我们多使用优化器的 minimize 方法,将求导和变量值更新合二为一。而在 TensorFlow 2 中,我们需要先使用 tf.GradientTape 进行求导操作,然后再使用优化器的 apply_gradients 方法应用已求得的导数,进行变量值的更新。而且在这两步中,都需要显式指定待求导和待更新的变量。

自动求导机制的计算图对比 *

在本节,为了帮助读者更深刻地理解 TensorFlow 的自动求导机制,我们以前节的“计算 y = x^2x = 3 时的导数”为例,展示 TensorFlow 1.X 和 TensorFlow 2 在图执行模式下,为这一求导过程所建立的计算图,并进行详细讲解。

在 TensorFlow 1.X 的图执行模式 API 中,将生成的计算图使用 TensorBoard 进行展示:

../../_images/grad_v1.png

在计算图中,灰色的块为节点的命名空间(Namespace,后文简称“块”),椭圆形代表操作节点(OpNode),圆形代表常量,灰色的箭头代表数据流。为了弄清计算图节点 xyy_grad 与计算图中节点的对应关系,我们将这些变量节点输出,可见:

  • x : <tf.Variable 'x:0' shape=() dtype=float32>

  • y : Tensor("Square:0", shape=(), dtype=float32)

  • y_grad : [<tf.Tensor 'gradients/Square_grad/Mul_1:0' shape=() dtype=float32>]

在 TensorBoard 中,我们也可以通过点击节点获得节点名称。通过比较我们可以得知,变量 x 对应计算图最下方的x,节点 y 对应计算图“Square”块的“ (Square) ”,节点 y_grad 对应计算图上方“Square_grad”的 Mul_1 节点。同时我们还可以通过点击节点发现,“Square_grad”块里的const节点值为2,“gradients”块里的 grad_ys_0 值为1, Shape 值为空,以及“x”块的const节点值为3。

接下来,我们开始具体分析这个计算图的结构。我们可以注意到,这个计算图的结构是比较清晰的,“x”块负责变量的读取和初始化,“Square”块负责求平方 y = x ^ 2 ,而“gradients”块则负责对“Square”块的操作求导,即计算 y_grad = 2 * x。由此我们可以看出, tf.gradients 是一个相对比较“庞大”的操作,并非如一般的操作一样往计算图中添加了一个或几个节点,而是建立了一个庞大的子图,以应用链式法则求计算图中特定节点的导数。

在 TensorFlow 2 的图执行模式 API 中,将生成的计算图使用 TensorBoard 进行展示:

../../_images/grad_v2.png

我们可以注意到,除了求导过程没有封装在“gradients”块内,以及变量的处理简化以外,其他的区别并不大。由此,我们可以看出,在图执行模式下, tf.GradientTape 这一上下文管理器的 gradient 方法和 TensorFlow 1.X 的 tf.gradients 是基本等价的。

小结

TensorFlow 1.X 中的 tf.gradients 和 TensorFlow 2 图执行模式下的 tf.GradientTape 上下文管理器尽管在 API 层面的调用方法略有不同,但最终生成的计算图是基本一致的。

基础示例:线性回归

在本节,我们为 第一章的线性回归示例 提供一个基于 TensorFlow 1.X 的图执行模式 API 的版本,供有需要的读者对比参考。

与第一章的NumPy和即时执行模式不同,TensorFlow的图执行模式使用 符号式编程 来进行数值运算。首先,我们需要将待计算的过程抽象为计算图,将输入、运算和输出都用符号化的节点来表达。然后,我们将数据不断地送入输入节点,让数据沿着计算图进行计算和流动,最终到达我们需要的特定输出节点。

以下代码展示了如何基于TensorFlow的符号式编程方法完成与前节相同的任务。其中, tf.placeholder() 即可以视为一种“符号化的输入节点”,使用 tf.get_variable() 定义模型的参数(Variable类型的张量可以使用 tf.assign() 操作进行赋值),而 sess.run(output_node, feed_dict={input_node: data}) 可以视作将数据送入输入节点,沿着计算图计算并到达输出节点并返回值的过程。

import tensorflow.compat.v1 as tf
tf.disable_eager_execution()

# 定义数据流图
learning_rate_ = tf.placeholder(dtype=tf.float32)
X_ = tf.placeholder(dtype=tf.float32, shape=[5])
y_ = tf.placeholder(dtype=tf.float32, shape=[5])
a = tf.get_variable('a', dtype=tf.float32, shape=[], initializer=tf.zeros_initializer)
b = tf.get_variable('b', dtype=tf.float32, shape=[], initializer=tf.zeros_initializer)

y_pred = a * X_ + b
loss = tf.constant(0.5) * tf.reduce_sum(tf.square(y_pred - y_))

# 反向传播,手动计算变量(模型参数)的梯度
grad_a = tf.reduce_sum((y_pred - y_) * X_)
grad_b = tf.reduce_sum(y_pred - y_)

# 梯度下降法,手动更新参数
new_a = a - learning_rate_ * grad_a
new_b = b - learning_rate_ * grad_b
update_a = tf.assign(a, new_a)
update_b = tf.assign(b, new_b)

train_op = [update_a, update_b] 
# 数据流图定义到此结束
# 注意,直到目前,我们都没有进行任何实质的数据计算,仅仅是定义了一个数据图

num_epoch = 10000
learning_rate = 1e-3
with tf.Session() as sess:
    # 初始化变量a和b
    tf.global_variables_initializer().run()
    # 循环将数据送入上面建立的数据流图中进行计算和更新变量
    for e in range(num_epoch):
        sess.run(train_op, feed_dict={X_: X, y_: y, learning_rate_: learning_rate})
    print(sess.run([a, b]))

自动求导机制

在上面的两个示例中,我们都是手工计算获得损失函数关于各参数的偏导数。但当模型和损失函数都变得十分复杂时(尤其是深度学习模型),这种手动求导的工程量就难以接受了。因此,在图执行模式中,TensorFlow同样提供了 自动求导机制 。类似于即时执行模式下的 tape.grad(ys, xs) ,可以利用TensorFlow的求导操作 tf.gradients(ys, xs) 求出损失函数 loss 关于 ab 的偏导数。由此,我们可以将上节中的两行手工计算导数的代码


# 反向传播,手动计算变量(模型参数)的梯度
grad_a = tf.reduce_sum((y_pred - y_) * X_)

替换为

grad_a, grad_b = tf.gradients(loss, [a, b])

计算结果将不会改变。

优化器

TensorFlow在图执行模式下也附带有多种 优化器 (optimizer),可以将求导和梯度更新一并完成。我们可以将上节的代码


# 反向传播,手动计算变量(模型参数)的梯度
grad_a = tf.reduce_sum((y_pred - y_) * X_)
grad_b = tf.reduce_sum(y_pred - y_)

# 梯度下降法,手动更新参数
new_a = a - learning_rate_ * grad_a
new_b = b - learning_rate_ * grad_b
update_a = tf.assign(a, new_a)
update_b = tf.assign(b, new_b)

整体替换为

optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate_)
grad = optimizer.compute_gradients(loss)
train_op = optimizer.apply_gradients(grad)

这里,我们先实例化了一个TensorFlow中的梯度下降优化器 tf.train.GradientDescentOptimizer() 并设置学习率。然后利用其 compute_gradients(loss) 方法求出 loss 对所有变量(参数)的梯度。最后通过 apply_gradients(grad) 方法,根据前面算出的梯度来梯度下降更新变量(参数)。

以上三行代码等价于下面一行代码:

train_op = tf.train.GradientDescentOptimizer(learning_rate=learning_rate_).minimize(loss)

使用自动求导机制和优化器简化后的代码如下:

import tensorflow as tf

learning_rate_ = tf.placeholder(dtype=tf.float32)
X_ = tf.placeholder(dtype=tf.float32, shape=[5])
y_ = tf.placeholder(dtype=tf.float32, shape=[5])
a = tf.get_variable('a', dtype=tf.float32, shape=[], initializer=tf.zeros_initializer)
b = tf.get_variable('b', dtype=tf.float32, shape=[], initializer=tf.zeros_initializer)

y_pred = a * X_ + b
loss = tf.constant(0.5) * tf.reduce_sum(tf.square(y_pred - y_))

# 反向传播,利用TensorFlow的梯度下降优化器自动计算并更新变量(模型参数)的梯度
train_op = tf.train.GradientDescentOptimizer(learning_rate=learning_rate_).minimize(loss)

num_epoch = 10000
learning_rate = 1e-3
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    for e in range(num_epoch):
        sess.run(train_op, feed_dict={X_: X, y_: y, learning_rate_: learning_rate})
    print(sess.run([a, b]))