《PyTorch深度学习实践》第2讲—线性模型

知识点

1. 使用matplotlib进行3D曲面绘图
2. numpy.meshgrid函数

使用matplotlib进行3D曲面绘图

matplotlib包本身不具有绘制3D图形功能，需要借助它依赖的Axes3D工具实现绘制3D图形

同时需要注意，plot_surface函数接收的数据参数都是numpy.ndarray类型。

numpy.meshgrid函数

numpy.meshgrid函数的功能是对两个向量进行扩展，扩展成为一个n*n的矩阵

作业代码

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

xData = [1.0, 2.0, 3.0]
yData = [2.0, 4.0, 6.0]

# predict function
def predict(x : float, w : np.ndarray, b : np.ndarray) -> np.ndarray:
    return x * w + b

# loss function
def loss(predictY : np.ndarray, realY : float) -> float:
    return (predictY - realY) ** 2

wList = np.arange(0.0, 4.1, 0.1)
bList = np.arange(-2.0, 2.0, 0.1)
w, b = np.meshgrid(wList, bList)
mse = np.zeros(w.shape)

# enum w and b to predict
for x, y in zip(xData, yData):
    predictY = predict(x, w, b)
    mse += loss(predictY, y)

mse /= len(xData)

# draw
# add axes
fig = plt.figure()
ax = Axes3D(fig)

# After matplotlib 3.4, Axes3D will need to be explicitly added to the figure
fig.add_axes(ax)
plt.xlabel(r'w', fontsize=20, color='cyan')
plt.ylabel(r'b', fontsize=20, color='cyan')
ax.plot_surface(w, b, mse, rstride = 1, cstride = 1, cmap = 'rainbow')
plt.show()