pytorch 逻辑回归

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


dot_num = 100
x_p = np.random.normal(3., 1, dot_num)
y_p = np.random.normal(6., 1, dot_num)
label_p = np.ones(dot_num)
C1 = np.array([x_p, y_p, label_p]).T      # 标签为1的数据

x_n = np.random.normal(6., 1, dot_num)
y_n = np.random.normal(3., 1, dot_num)
label_n = np.zeros(dot_num)
C2 = np.array([x_n, y_n, label_n]).T     # 标签为0的数据

# plt.scatter(C1[:, 0], C1[:, 1], c='b', marker='+')
# plt.scatter(C2[:, 0], C2[:, 1], c='g', marker='o')
# plt.show()

data_set = np.concatenate((C1, C2), axis=0)
np.random.shuffle(data_set)


class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.fc = nn.Linear(2, 1, bias=True)

    def forward(self, x):
        logits = self.fc(x)
        logits = torch.sigmoid(logits)
        return logits


model = LogisticRegression()
if torch.cuda.is_available():
    model.cuda()

# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)


x1,x2,y = list(zip(*data_set))
x = list(zip(x1, x2))

# 开始训练
for epoch in range(50):
    if torch.cuda.is_available():
        x_data = torch.FloatTensor(x).cuda()
        y_data = torch.FloatTensor(y).cuda()
    else:
        x_data = torch.FloatTensor(x)
        y_data = torch.FloatTensor(y)

    out = model(x_data)
    loss = criterion(out, y_data)
    print_loss = loss.data.item()
    mask = out.ge(0.5).squeeze().float()  # 以0.5为阈值进行分类
    correct = (mask == y_data).sum()  # 计算正确预测的样本个数
    acc = correct.item() / x_data.size(0)  # 计算精度

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # 每隔20轮打印一下当前的误差和精度
    if (epoch + 1) % 5 == 0:
        print('*' * 10)
        print('epoch {}'.format(epoch + 1))  # 训练轮数
        print('loss is {:.4f}'.format(print_loss))  # 误差
        print('acc is {:.4f}'.format(acc))  # 精度


# 结果可视化
w0, w1 = model.fc.weight[0]
w0 = float(w0.item())
w1 = float(w1.item())
b = float(model.fc.bias.item())
plot_x = np.arange(-7, 7, 0.1)
plot_y = (-w0 * plot_x - b) / w1
plt.scatter(data_set[:, 0], data_set[:, 1], c=data_set[:,2], s=100, lw=0, cmap='RdYlGn')
plt.plot(plot_x, plot_y)
plt.show()