帮我用pytorch 写一个全连接网络的训练代码
时间: 2024-05-08 14:15:57 浏览: 7
以下是使用PyTorch编写的全连接神经网络训练代码示例:
```python
import torch
import torch.nn as nn
import torch.optim as optim
# 定义全连接神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 定义训练数据集
train_data = torch.randn(100, 10)
train_target = torch.randn(100, 10)
# 定义网络模型、损失函数和优化器
net = Net()
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)
# 开始训练
for epoch in range(1000):
optimizer.zero_grad()
output = net(train_data)
loss = criterion(output, train_target)
loss.backward()
optimizer.step()
if epoch % 100 == 0:
print('Epoch %d, Loss: %.4f' % (epoch, loss.item()))
# 测试模型
test_data = torch.randn(10, 10)
test_output = net(test_data)
print('Test Output:', test_output)
```
在上述示例代码中,我们首先定义了一个包含两个全连接层的神经网络模型 `Net`,然后定义了训练数据集 `train_data` 和 `train_target`,以及损失函数 `criterion` 和优化器 `optimizer`。
在训练过程中,我们使用了一个循环语句,每次迭代时都将优化器的梯度清零,计算模型输出并计算损失,然后进行反向传播和梯度更新。在每个 epoch 结束时,我们打印出当前 epoch 的损失值。
最后,我们使用测试数据集 `test_data` 进行测试,并打印出模型的输出结果。