PyTorch - 線性回歸 - Linear Regression - 2

假定一個簡單的一元二次函數 y=a*x^2+b 利用此產生亂數個點，然後來試著擬合這些點。

1. 建立 DATA

import torch
import matplotlib.pyplot as plt
X = torch.unsqueeze(torch.linspace(-1, 1, 200), dim=1)  # x data (tensor), shape=(100, 1)
Y = 2*X.pow(2) + 0.3*torch.rand(X.size())                 # noisy y data (tensor), shape=(100, 1)
plt.scatter(X.data.numpy(), Y.data.numpy())
plt.show()

將X,Y 丟進 Variable 做梯度計算用。

X=Variable(torch.Tensor(X.reshape(200,1)))
Y=Variable(torch.Tensor(Y.reshape(200,1)))

2. 建立簡單模型(快速搭建)

這邊使用快速搭建簡易的 Sequential model，且使用兩層 linear neural network，且第一層使用Relu激活函數，優化器使用 RMSprop，lr=0.01 ，Loss function 一樣使用 MSE。

model = torch.nn.Sequential(
    torch.nn.Linear(1, 15),
    torch.nn.ReLU(),
    torch.nn.Linear(15, 1)
)
optimizer = torch.optim.RMSprop(model.parameters(), lr=0.01)
loss_function = torch.nn.MSELoss()

3. 開始訓練

迭代500次訓練。

epochs = 500
for epoch in range(epochs):
    prediction = model(X)
    loss = loss_function(prediction, Y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

4. 訓練結果可視化

plt.figure(1, figsize=(25, 5))
plt.subplot(131)
plt.title('model')
plt.scatter(X.data.numpy(), Y.data.numpy())
plt.plot(X.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
plt.text(0.5, 0, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color':  'red'})
plt.show()

5. 動態可視化

在訓練同時，將訓練過程畫出來，可以更了解訓練過程，優化器在對模型做什麼事情。

epochs = 500
for epoch in range(epochs):
    prediction = model(X)
    loss = loss_function(prediction, Y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 5 == 0:
        # plot and show learning process
        plt.cla()
        plt.scatter(X.data.numpy(), Y.data.numpy())
        plt.plot(X.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
        plt.text(0.5, 0, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color':  'red'})
        plt.pause(0.1)