pytorch 写一个ANN神经网络数据分类模型
时间: 2024-04-17 16:26:03 浏览: 225
下面是使用 PyTorch 构建一个基本的人工神经网络(Artificial Neural Network,简称 ANN)进行数据分类的示例代码:
```python
import torch
import torch.nn as nn
import torch.optim as optim
# 定义神经网络模型
class ANN(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(ANN, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
# 准备数据
input_size = 10 # 输入特征的维度
hidden_size = 20 # 隐藏层大小
num_classes = 2 # 类别数量
# 创建模型实例
model = ANN(input_size, hidden_size, num_classes)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 每训练一轮输出一次损失值
if (epoch+1) % 1 == 0:
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
# 使用模型进行预测
with torch.no_grad():
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
# 输出预测结果
print("Predicted classes:", predicted)
```
在上面的代码中,我们首先定义了一个名为 `ANN` 的神经网络模型,它包含一个输入层、一个隐藏层和一个输出层。然后,我们定义了损失函数和优化器,使用交叉熵损失函数和随机梯度下降(SGD)优化器。接下来,我们进行模型训练,通过多次迭代更新模型的参数。最后,我们使用训练好的模型进行预测,并输出预测结果。
请注意,上述代码只是一个示例,实际应用中可能需要根据具体问题和数据集的特点进行适当的调整和修改。
阅读全文