pytorch（二）：分类

import torch
import torch.nn.functional as f
from torch.autograd import Variable
import matplotlib.pyplot as plt


# 建造数据集
data = torch.ones((100, 2))
x0 = torch.normal(2*data, 1)
y0 = torch.zeros(100)  # y0是标签  shape(100,),是一维
x1 = torch.normal(-2*data, 1)
y1 = torch.ones(100)  # y1也是标签 shape(100,)，是一维
x = torch.cat((x0, x1), 0).type(torch.FloatTensor)  # 参数0表示维度，在纵向方向将x0,x1合并，合并后shape(200, 2）)
y = torch.cat((y0, y1), 0).type(torch.LongTensor)  # 标签是0或1，类型为整数，LongTensor = 64-bit integer,
x, y = Variable(x), Variable(y)  # 训练神经网络只能接受变量输入，故要把x, y转化为变量
plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1],  # 这两个参数分别代表x,y轴坐标
            c=y.data.numpy(), s=100, cmap='RdYlGn')  # c为color,y有两种标签，代表两种颜色的点，'RdYlGn'红色和绿色
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 = f.relu(self.hidden(x))
        y = self.out(x)
        return y


# 定义神经网络
net = Net(n_feature=2, n_hidden=10, n_output=2) 
# n_output=2,因为它返回一个元素为2的列表。[0, 1]表示学习到的内容为标签1，[1, 0]表示学习到的内容为标签0。
print(net)


# 训练神经网络模型并将训练过程可视化
optimizer = torch.optim.SGD(net.parameters(), lr=0.02)
loss_func = torch.nn.CrossEntropyLoss()
plt.ion()
for i in range(100):
    out = net(x)
    loss = loss_func(out, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # 绘图
    if i % 2 == 0:
        plt.cla()
        # torch.max(a,1) 返回每一行中最大值的那个元素，且返回其索引（返回最大元素在这一行的列索引
        # f.softmax(out）是将out的内容以概率表示。
        # torch.max()返回的是两个Variable，第一个Variable存的是最大值，第二个存的是其对应的位置索引index。这里我们想要得到的是索引，所以后面用[1]。
        prediction = torch.max(f.softmax(out), 1)[1]
        pred_y = prediction.data.numpy().squeeze()
        target_y = y.data.numpy()
        plt.scatter(x.data.numpy()[:, 0], x.data.numpy()[:, 1], c=pred_y, s=100, cmap='RdYlGn')
        accuracy = sum(pred_y == target_y)/200
        plt.text(1.5, -4, 'accuracy=%.2f'%accuracy, fontdict={'size':10, 'color':'red'})
        plt.pause(0.1)
plt.ioff()
plt.show()