pytorch实现逻辑回归分类

学习b站刘二大人视频所写
视频传送门

import torch
import matplotlib.pyplot as plt
import numpy as np

x_data = torch.tensor([[1.0], [2.0], [3.0]])
y_data = torch.tensor([[0.0], [0.0], [1.0]])


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

    def forward(self, x):
        y_pred = torch.sigmoid(self.linear(x))
        return y_pred


model = LogisticModel()

# construct loss and optimizer
# criterion = torch.nn.MSELoss(size_average = False)
criterion = torch.nn.BCELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(), lr=0.02)  # model.parameters()自动完成参数的初始化操作
loss_sum = []
# training cycle forward, backward, update
for epoch in range(1000):
    y_pred = model(x_data)  # forward 算预测值
    loss = criterion(y_pred, y_data)  # forward: 算损失值
    print(epoch, loss.item(), model.linear.weight.item())
    loss_sum.append(loss)
    optimizer.zero_grad()  # 清除上一轮的梯度，防止累积
    loss.backward()  # backward: autograd，自动计算梯度，反向传播
    optimizer.step()  # update 参数，即更新w和b的值


print('w = ', model.linear.weight.item())
print('b = ', model.linear.bias.item())

x_test = torch.Tensor([[4.0]])
y_test = model(x_test)
print('y_pred = ', y_test.data)

x = range(1000)
y = loss_sum
plt.plot(x, y)
plt.xlabel('Epoch')
plt.ylabel('loss')
# plt.gcf().set_facecolor(np.ones(3) * 240 / 255)   # 生成画布的大小
plt.grid() # 生成网格
plt.show()

x = np.linspace(0, 10, 200)
x_t = torch.Tensor(x).view(200, 1)
y_t = model(x_t)
y = y_t.data.numpy()
plt.plot(x, y)
plt.plot([0, 10], [0.5, 0.5], c='r')
plt.xlabel('Hours')
plt.ylabel('Probability of pass')
plt.grid()
plt.show()