pytorch的ANN
时间: 2024-01-27 22:14:48 浏览: 154
以下是使用PyTorch实现的ANN(人工神经网络)的示例代码:
```python
import torch
import torch.nn as nn
# 定义ANN模型
class ANN(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(ANN, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size) # 输入层到隐藏层的全连接层
self.relu = nn.ReLU() # 激活函数
self.fc2 = nn.Linear(hidden_size, output_size) # 隐藏层到输出层的全连接层
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
# 创建ANN模型实例
input_size = 10 # 输入特征的维度
hidden_size = 20 # 隐藏层的大小
output_size = 2 # 输出类别的数量
model = ANN(input_size, hidden_size, output_size)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 随机梯度下降优化器
# 训练ANN模型
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('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# 使用ANN模型进行预测
with torch.no_grad():
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
# 输出预测结果
print('Predicted classes:', predicted)
```
这段代码展示了如何使用PyTorch库来实现一个简单的ANN模型。首先,我们定义了一个继承自nn.Module的ANN类,其中包含了输入层、隐藏层和输出层。然后,我们定义了损失函数和优化器,并进行了模型的训练。最后,我们使用训练好的模型进行预测,并输出预测结果。
阅读全文