这篇文章主要是为了展示matplotlib 3D图可视化梯度下降路径的效果。

模型选择了最简单的y=w*x+b模型，创建一组数据，迭代训练500次后，画出平面曲线图和原始数据散点图，看看我们训练处的模型拟合原始数据的表现如何。再画一个3D图，x轴为w，y轴为b，z轴为每次训练的loss值，看看我们的梯度下降路径的变化。

代码如下：

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

lr = 0.1
iteration = 500
# 真正的w和b
real_params = [1.2, 2.5]

X = tf.placeholder(tf.float32, [None, 1])
Y = tf.placeholder(tf.float32, [None, 1])
# 要训练的w和b初始化为5和4
weight = tf.Variable(initial_value=[[5]], dtype=tf.float32)
bias = tf.Variable(initial_value=[[4]], dtype=tf.float32)
y = tf.matmul(X, weight) + bias

loss = tf.losses.mean_squared_error(Y, y)
train_op = tf.train.GradientDescentOptimizer(lr).minimize(loss)

x_data = np.linspace(-1, 1, 200)[:, np.newaxis]
noise = np.random.normal(0, 0.1, x_data.shape)
y_data = x_data * real_params[0] + real_params[1] + noise

with tf.Session() as sess:
	sess.run(tf.global_variables_initializer())
	# 记录训练过程中的w、b和loss
	w_record = []
	b_record = []
	loss_record = []
	for i in range(iteration):
		w, b, cost, _ = sess.run([weight, bias, loss, train_op], feed_dict={X: x_data, Y: y_data})
		w_record.append(w)
		b_record.append(b)
		loss_record.append(cost)
	result = sess.run(y, feed_dict={X: x_data, Y: y_data})

plt.figure(1)
plt.scatter(x_data, y_data, s=1, color='r', alpha=0.5)
plt.plot(x_data, result, lw=1)
plt.show()
plt.close()

fig = plt.figure(2)
ax_3d = Axes3D(fig)
w_3d, b_3d = np.meshgrid(np.linspace(-4.2, 6.2, 50), np.linspace(-2.5, 7.5, 50))
loss_3d = np.array(
	[np.mean(np.square((x_data * w_ + b_) - y_data)) for w_, b_ in zip(w_3d.ravel(), b_3d.ravel())]).reshape(w_3d.shape)
ax_3d.plot_surface(w_3d, b_3d, loss_3d, cmap=plt.get_cmap('rainbow'))
# .ravel()函数会直接修改原始矩阵，也是将矩阵降成一维
w_record = np.array(w_record).ravel()
b_record = np.array(b_record).ravel()
# 梯度下降曲线
ax_3d.plot(w_record, b_record, loss_record, lw=6, c='blue')
# 起始点
ax_3d.scatter(w_record[0], b_record[0], loss_record[0], s=10, color='black')
ax_3d.set_xlabel('w')
ax_3d.set_ylabel('b')
plt.show()
plt.close()

首先将lr=0.1，看看效果如何：

这个时候曲线拟合的很好。

观察3D图，可以发现梯度下降直接到了最低的loss点。

现在我们将lr设为1.0，看看效果如何：

可以看到由于学习率过大，这个时候曲线很显然不能拟合原始数据。

扫描二维码关注公众号，回复： 4634172 查看本文章

再看3D图，我们可以发现loss值由于学习率过大，一直在"山谷"两侧震荡，无法到达最低点。