使用pytorch+cuda实现简单二分类问题

首先我们知道二分类问题的关键是其中的sigmoid函数, sigmoid如何运作的不再赘述.

先创建模型

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

class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression,self).__init__()
        self.linear = nn.Linear(2,1) 
        self.sigmoid = torch.nn.Sigmoid()
    def forward(self, x):
        out = self.sigmoid(self.linear(x))
        return out

由代码可以看到, 在forward前向传播过程中有一个二分类sigmoid.

接下来是创造假数据, 并且将数据转成pytorch可以读取的格式

from sklearn.datasets.samples_generator import make_blobs
X, y = make_blobs(200, centers=2, cluster_std=0.2, random_state=0)

x_train = torch.from_numpy(X).type(torch.FloatTensor)
y_train = torch.from_numpy(y.reshape(len(y), 1)).type(torch.FloatTensor)

开始训练, 这里会根据是否有cuda从而选择是基于GPU训练还是CPU训练.

这里的loss function选择的是BCE(binary cross entropy 二元交叉熵), 并使用随机梯度下降SGD

num_epochs = 1000
model = LogisticRegression()
if torch.cuda.is_available():
    print("GPU1")
    model.cuda()

criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(),lr=1e-3)


for epoch in  range(num_epochs):

    inputs = Variable(x_train)
    target = Variable(y_train)
    if torch.cuda.is_available():

        inputs = inputs.cuda()
        target = target.cuda()
               

    # forward
    out = model(inputs)
    loss = criterion(out,target)
    # backward
    optimizer.zero_grad()  # 梯度归零
    loss.backward()   # 反向传播
    optimizer.step()  # 更新参数
    if epoch%100==0:
      print(loss.item())



model.eval()
inputs = Variable(x_train)
if torch.cuda.is_available():
  inputs = inputs.cuda()
predict = model(inputs)  


predict = predict.ge(0.5).float() #二分类 sigmoid
predict = predict.data.cpu().numpy() #转成这个，不然不能把他用到c上面

plt.scatter(x_train[:,0].numpy(),x_train[:,1].numpy(),c=predict,label='Original data')

plt.show()

每训练100次, 输出loss

结果如下

可以发现是基于GPU训练的,并且loss也一直在减小, 最后的效果看起来不错!