能不能帮我写一个pytorch环境的神经网络代码
时间: 2024-05-11 21:14:09 浏览: 102
当然可以!以下是一个简单的神经网络代码示例,使用PyTorch框架:
```
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(784, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = torch.flatten(x, 1)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
output = nn.functional.log_softmax(x, dim=1)
return output
# 定义损失函数和优化器
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 训练模型
for epoch in range(10):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 1000 == 999:
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 1000))
running_loss = 0.0
print('Finished Training')
```
这个示例代码将创建一个包含两个线性层的神经网络,用于对手写数字进行分类。训练过程将使用SGD优化器和交叉熵损失函数。你可以根据你的需求修改模型结构、优化器和损失函数等。
阅读全文