code by PyTorch simply

import torch
import torch.nn as nn
from matplotlib import pyplot as plt

data

x_num, feature_num = 100, 1

true_w, true_b = 1., -3.

train_x = torch.normal(0, 1, [x_num, feature_num])

true_data = true_w * train_x + true_b
noise = torch.normal(0, 0.25, true_data.size())
train_y = true_data + noise

plt.scatter(train_x, train_y)
plt.plot([-4., 4.], [true_w*-4.+true_b, true_w*4.+true_b], linewidth = '3', label = "true", color='red')
plt.legend(["true line","data"])
plt.show()
 

torch data loader

import torch.utils.data as Data

batch_size = 10
data_set = Data.TensorDataset(train_x, train_y)
data_iter = torch.utils.data.DataLoader(data_set, batch_size, shuffle=True)
print(data_iter)
for X, y in data_iter:
    print(X.size(), y.size())
    break

<torch.utils.data.dataloader.DataLoader object at 0x000001AD2D8BF370>
torch.Size([10, 1]) torch.Size([10, 1])
 

linear net

class LinearNet(nn.Module):
    def __init__(self, num_feature):
        super(LinearNet, self).__init__()
        self.linear = nn.Linear(num_feature, 1)

    def forward(self, x):
        return self.linear(x)

net = LinearNet(feature_num)

# initialize parameter
nn.init.normal_(net.linear.weight, mean=0, std=0.01)
nn.init.constant_(net.linear.bias, val=0)

# cost function
loss = nn.MSELoss()

# optimizer
optimizer = torch.optim.SGD(net.parameters(), lr=0.03)

# train
epochs = 5
for epoch in range(epochs):
    cost = 0.0
    for X, y in data_iter:
        output = net(X)
        cost = loss(output, y.view(-1, 1))

        optimizer.zero_grad()
        cost.backward()
        optimizer.step()
    print("epoch %d, loss: %f" % (epoch, cost.item()))
print('\n', true_w, true_b,'\n')
print(net.linear.weight, net.linear.bias)

pred_w = net.linear.weight
pred_b = net.linear.bias
plt.scatter(train_x, train_y)
plt.plot([-4., 4.], [true_w*-4.+true_b, true_w*4.+true_b], linewidth = '3', label = "true", color='red')
plt.plot([-4., 4.], [pred_w*-4.+pred_b, pred_w*4.+pred_b], linewidth = '3', label = "true", color='green')
plt.legend(["true line","predict", "data"])
plt.show()
 

epoch 0, loss: 2.461334
epoch 1, loss: 1.148570
epoch 2, loss: 0.322340
epoch 3, loss: 0.122028
epoch 4, loss: 0.200466

 1.0 -3.0 

Parameter containing:
tensor([[0.9121]], requires_grad=True) Parameter containing:
tensor([-2.8788], requires_grad=True)