Caffe2 - (十五) 简单的回归示例 Toy Regression

# 导入 python packages
from caffe2.python import core, cnn, net_drawer, workspace, visualize
import numpy as np
import matplotlib.pyplot as plt

init_net = core.Net("init")
# The ground truth parameters.
W_gt = init_net.GivenTensorFill([], "W_gt", shape=[1, 2], values=[2.0, 1.5])
B_gt = init_net.GivenTensorFill([], "B_gt", shape=[1], values=[0.5])
# Constant value ONE is used in weighted sum when updating parameters.
ONE = init_net.ConstantFill([], "ONE", shape=[1], value=1.)
# ITER is the iterator count.
ITER = init_net.ConstantFill([], "ITER", shape=[1], value=0, dtype=core.DataType.INT32)

# 待学习参数
# 随机初始化权重为区间[-1, 1]的值，初始化 bias=0.0
W = init_net.UniformFill([], "W", shape=[1, 2], min=-1., max=1.)
B = init_net.ConstantFill([], "B", shape=[1], value=0.0)
print('Created init net.')

train_net = core.Net("train")
# 首先，随机生成数据样本 X；并创建 ground truth
X = train_net.GaussianFill([], "X", shape=[64, 2], mean=0.0, std=1.0, run_once=0)
Y_gt = X.FC([W_gt, B_gt], "Y_gt")
# 添加 Gaussian noise 到 ground truth
noise = train_net.GaussianFill([], "noise", shape=[64, 1], mean=0.0, std=1.0, run_once=0)
Y_noise = Y_gt.Add(noise, "Y_noise")
# 由于不需要计算 Y_noise 的梯度，故设置 StopGradient 来忽略其自动微分计算.
Y_noise = Y_noise.StopGradient([], "Y_noise")

# 线性回归预测
Y_pred = X.FC([W, B], "Y_pred")

# Loss 函数是 squared L2 distance，并对 minibatch 求平均
dist = train_net.SquaredL2Distance([Y_noise, Y_pred], "dist")
loss = dist.AveragedLoss([], ["loss"])

graph = net_drawer.GetPydotGraph(train_net.Proto().op, "train", rankdir="LR")
graph.write_png('graph.png')

# Get gradients for all the computations above.
gradient_map = train_net.AddGradientOperators([loss])
graph = net_drawer.GetPydotGraph(train_net.Proto().op, "train", rankdir="LR")
graph.write_png('graph_gradient.png')

# 增量迭代，+1
train_net.Iter(ITER, ITER)
# 计算迭代后的学习率
LR = train_net.LearningRate(ITER, "LR", base_lr=-0.1,
                            policy="step", stepsize=20, gamma=0.9)

# 权重累加
train_net.WeightedSum([W, ONE, gradient_map[W], LR], W)
train_net.WeightedSum([B, ONE, gradient_map[B], LR], B)

# 在此可视化 graph
graph = net_drawer.GetPydotGraph(train_net.Proto().op, "train", rankdir="LR")
graph.write_png('graph2.png')

workspace.RunNetOnce(init_net)
workspace.CreateNet(train_net)

# 查看参数 W 和 B 
print("Before training, W is: {}".format(workspace.FetchBlob("W")))
print("Before training, B is: {}".format(workspace.FetchBlob("B")))

# 迭代训练 100 次
for i in range(100):
    workspace.RunNet(train_net.Proto().name)

# 训练后的参数
print("After training, W is: {}".format(workspace.FetchBlob("W")))
print("After training, B is: {}".format(workspace.FetchBlob("B")))

print("Ground truth W is: {}".format(workspace.FetchBlob("W_gt")))
print("Ground truth B is: {}".format(workspace.FetchBlob("B_gt")))

workspace.RunNetOnce(init_net)
w_history = []
b_history = []
for i in range(50):
    workspace.RunNet(train_net.Proto().name)
    w_history.append(workspace.FetchBlob("W"))
    b_history.append(workspace.FetchBlob("B"))
w_history = np.vstack(w_history)
b_history = np.vstack(b_history)
plt.plot(w_history[:, 0], w_history[:, 1], 'r')
plt.axis('equal')
plt.xlabel('w_0')
plt.ylabel('w_1')
plt.grid(True)
plt.figure()
plt.plot(b_history)
plt.xlabel('iter')
plt.ylabel('b')
plt.grid(True)
plt.show()