Build a simple linear regression model with pytorch

This article builds a simple linear regression model through pytorch, which is mainly used to understand the basic training process of pytorch.

For example, the following linear regression problem:

import numpy as np 
import matplotlib.pyplot as plt 

X_train = np.arange(10, dtype='float32').reshape((10, 1))
y_train = np.array([1.0, 1.5, 2.8, 4, 3, 5.2, 7.2, 6.9, 8.5, 9.5], dtype='float32')

plt.plot(X_train, y_train, 'o', markersize=8)
plt.xlabel('x')
plt.ylabel('y')
plt.show()

First, you need to convert the data into a way that pytorch can read:

from torch.utils.data import TensorDataset, DataLoader

#标准化
X_train_norm = (X_train- np.mean(X_train)) / np.std(X_train)
X_train_norm = torch.from_numpy(X_train_norm)
y_train = torch.from_numpy(y_train)
train_ds = TensorDataset(X_train_norm, y_train)
batch_size = 1 
train_dl= DataLoader(train_ds, batch_size, shuffle=True)

Then you need to define the loss function, model, and optimizer

import torch.nn as nn 
import torch 

loss_fn = nn.MSELoss(reduction='mean')
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

The loss function uses the mean square error under torch.nn

Linear model for model: inputsize outputsize is 1

SGD used by the optimizer, learning rate=0.001

Finally train the model

num_epochs = 200
for epoch in range(num_epochs):
    for x_batch, y_batch in train_dl:
        #1.产生 输出
        pred = model(x_batch)[:, 0]
        #2. 计算 损失
        loss = loss_fn(pred, y_batch)
        #3. 计算梯度
        loss.backward()
        #4. 更新参数
        optimizer.step()
        #5.梯度reset to 0 
        optimizer.zero_grad()

Let's take a look at the parameters obtained by training:

print(model.weight.item(), model.bias.item())

2.690967321395874 4.879624843597412

Visualization:

X_test = np.linspace(0, 9, num=100, dtype='float32').reshape(-1, 1)
X_test_norm = (X_test-np.mean(X_train))/np.std(X_train)
X_test_norm = torch.from_numpy(X_test_norm) 
y_pred = model(X_test_norm).detach().numpy()
fig = plt.figure(figsize=(12, 5))
ax = fig.add_subplot(1, 1, 1)
plt.plot(X_train_norm, y_train, 'o', markersize= 8)
plt.plot(X_test_norm, y_pred, '--', lw=2)
plt.legend(['Train examples', 'Linear regress'], fontsize=12)
ax.set_xlabel('x', size=12)
ax.set_ylabel('y', size=12)
plt.show()

参考自: Machine Learning with PyTorch and Scikit-Learn Book by Sebastian Raschka

Build a simple linear regression model with pytorch

Guess you like