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()
