PyTorch实现线性回归

# gradients_pytorch
import torch
import torch.nn as nn

# step0 prepare for data
X = torch.tensor([[1], [2], [3], [4]], dtype=torch.float32)
Y = torch.tensor([[2], [4], [6], [8]], dtype=torch.float32)  # 4,1:4个样本，每个样本一个特征

X_test = torch.tensor([5], dtype=torch.float32)
# Y = torch.tensor([2, 4, 6, 8]) # 4: 一个样本，该样本4个特征
# print(Y.shape) # 打印张量的形状 或者 Y.size()
# torch.Size([4, 1])

# step1 design model:input, output size, forward pass
n_samples, n_features = X.shape
# print(type(X.shape))
# print(type(n_samples))
# <class 'torch.Size'>
# <class 'int'>
input_size = n_features
output_size = n_features  # 注意 这里的输入和输出是大小是特征数而不是样本数

# print(type(nn.Linear(3, 3))) #<class 'torch.nn.modules.linear.Linear'>
class LinearRegression(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegression, self).__init__()
        self.lin = nn.Linear(input_dim, output_dim)

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


model = LinearRegression(input_size, output_size)
# print(model) # LinearRegression(
    #   (lin): Linear(in_features=4, out_features=4, bias=True)
    # )
# print(type(model)) # <class '__main__.LinearRegression'>
Y_test = model(X_test)
# print(type(Y_test)) # <class 'torch.Tensor'>
# print(Y_test.size()) #torch.Size([1])
print(f'Prediction before training:f(5) = {Y_test.item():.3f}')

# step2 loss and optimizer
loss = nn.MSELoss()
# print(type(loss))
# print(loss)
# <class 'torch.nn.modules.loss.MSELoss'>
# MSELoss()
# print(model.parameters())# <generator object Module.parameters at 0x000001CCC5C795B0>
# print(type(model.parameters())) # <class 'generator'>
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# step3 model training
n_iters = 100
for epoch in range(n_iters):
    y_pred = model(X)
    # 计算损失
    l = loss(Y, y_pred)
    # 更新权值准备--计算梯度
    l.backward()
    # 更新权值
    optimizer.step()
    # print(f'{l}')
    # 梯度清零 防止梯度累加
    optimizer.zero_grad()
    # print(f'{l}')
    # if epoch == n_iters - 1:
    #     print(f'l的类型:{type(l)}')
    if (epoch + 1) % 10 == 0:
        [w, b] = model.parameters()
        print(f'epoch{epoch+1} w:{w[0][0]:.3f} loss:{l:.8f}') # 这里的l可以直接打印
# print(type(model(X_test)))
print(f'Prediction after training:{model(X_test).item():.3f}') # 这里是张量要用item方法进行分离