【pytorch】分类

# 分类
n_data=torch.ones(100,2)
# print(n_data)
x0=torch.normal(2*n_data,1) #class 0
y0=torch.zeros(100)
x1=torch.normal(-2*n_data,1)    #class 1
y1=torch.ones(100)
x=torch.cat((x0,x1)).type(torch.float32)
y=torch.cat((y0,y1)).type(torch.long)
x,y=Variable(x),Variable(y)

plt.scatter(x.data.numpy()[:,0],x.data.numpy()[:,1],c=y.data.numpy(),cmap='RdYlGn')
# plt.show()

# plt.scatter(n_data.numpy())
# plt.show()

class Net(torch.nn.Module):
    def __init__(self,n_feature,n_hidden,n_output):
        super(Net,self).__init__()
        self.hidden=torch.nn.Linear(n_feature,n_hidden)
        self.out=torch.nn.Linear(n_hidden,n_output)

    def forward(self,x):
        x=torch.relu(self.hidden(x))
        x=F.softmax(self.out(x))
        return x

net=Net(2,10,2)
# print(net)
# print(list(net.parameters()))
plt.ion()
optimizer=torch.optim.SGD(net.parameters(),lr=0.02)
loss_func=torch.nn.CrossEntropyLoss()


for t in range(100):
    out=net(x)
    loss=loss_func(out,y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if t % 2 == 0:
        plt.cla()
        prediction=torch.argmax(out,1)#每行最大值
        pred_y=prediction.data.numpy()
        target_y=y.data.numpy()
        plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1], c=pred_y, cmap='RdYlGn')
        acc=np.mean(pred_y == target_y)
        plt.text(1.4,-4,"Accuracy=%.2f"%acc,fontdict={'size':20,'color':'red'})
        plt.pause(0.2)

plt.ioff()
plt.show()