PyTorch框架代码
时间: 2024-01-11 21:02:03 浏览: 16
下面是一个简单的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, 1)
def forward(self, x):
x = torch.sigmoid(self.fc1(x))
x = self.fc2(x)
return x
# 定义训练数据集
train_data = torch.randn(100, 10)
train_labels = torch.randn(100, 1)
# 定义模型、损失函数、优化器
net = Net()
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)
# 训练模型
for epoch in range(100):
optimizer.zero_grad()
output = net(train_data)
loss = criterion(output, train_labels)
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {loss.item()}')
```
这个代码中定义了一个简单的神经网络模型 `Net`,使用MSE损失函数 `nn.MSELoss()` 和随机梯度下降优化器 `optim.SGD()` 进行模型训练。每隔10个epoch打印一次训练损失。