[torch] xor 分类问题训练

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt# 设置随机种子以确保结果的可重复性
torch.manual_seed(1)# 生成正确的XOR数据集x_data= np.float32(np.random.rand(500,2))
y_data = np.float32( np.sign(x_data[:,0]-0.5)* np.sign(x_data[:,1]-0.5) +1)/2# 定义模型，一个具有1个隐藏层的多层感知器[6,7](@ref)
class MyModel(nn.Module):def __init__(self):super(MyModel, self).__init__()self.hidden1 = nn.Linear(in_features=2, out_features=6)  # 隐藏层4个神经元self.hidden2 = nn.Linear(in_features=6, out_features=4)  # 隐藏层4个神经元self.output = nn.Linear(in_features=4, out_features=1)   # 输出层def forward(self, x):x = torch.tanh(self.hidden1(x))  # 使用tanh激活函数x = torch.tanh(self.hidden2(x))  # 使用tanh激活函数x = self.output(x)return xmodel = MyModel()# 定义损失函数 - 使用二元交叉熵[4,8](@ref)
loss_fn = nn.BCEWithLogitsLoss()  # 包含sigmoid和交叉熵计算# 设置优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # 学习率可以适当调高# 准备训练数据
x_tensor = torch.from_numpy(x_data)
y_tensor = torch.from_numpy(y_data)# 训练模型
epochs = 5000  # XOR问题需要更多迭代次数
losses = []predictions = torch.sigmoid(model(x_tensor))
pre_arr=predictions[:,0]
pre_arr=pre_arr.detach().numpy()plt.subplot(131)
plt.scatter(x_data[:,0],x_data[:,1],c=pre_arr,cmap='viridis',label='Untrained')
plt.legend()for epoch in range(1, epochs + 1):# 前向传播y_pred = model(x_tensor)loss = loss_fn(y_pred.reshape(500), y_tensor)# 反向传播和优化optimizer.zero_grad()  # 清空梯度loss.backward()        # 反向传播optimizer.step()       # 更新参数losses.append(loss.item())if epoch % 500 == 0:print(f'Epoch {epoch}: Loss = {loss.item()}')predictions = torch.sigmoid(model(x_tensor))
pre_arr=predictions[:,0]
pre_arr=pre_arr.detach().numpy()plt.subplot(132)
plt.scatter(x_data[:,0],x_data[:,1],c=pre_arr,cmap='viridis',label='Trained')
plt.legend()
plt.subplot(133)
plt.scatter(x_data[:,0],x_data[:,1],c=y_data,cmap='viridis',label='Real Data')
plt.legend()
plt.show()