使用梯度下降法或拟牛顿法来最小化损失函数pytorch代码
时间: 2024-04-21 19:22:13 浏览: 110
gradient_descent_机器学习_toodj1_python_descent_梯度下降法_
在PyTorch中,您可以使用梯度下降法或拟牛顿法来最小化损失函数。下面是一个示例代码,其中使用梯度下降法(使用`torch.optim.SGD`优化器)来最小化损失函数:
```python
import torch
import torch.nn as nn
import torch.optim as optim
# 定义模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(1, 1) # 假设输入是1维,输出是1维
def forward(self, x):
return self.linear(x)
# 创建模型实例
model = Model()
# 定义损失函数
criterion = nn.MSELoss()
# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
# 假设输入和目标是一些示例数据
inputs = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
targets = torch.tensor([[2.0], [4.0], [6.0], [8.0]])
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, targets)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 每一轮打印损失值
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
```
在这个例子中,我们首先定义了一个简单的线性模型,并使用均方误差损失函数(`nn.MSELoss()`)来定义损失。然后,我们创建了一个优化器(`optim.SGD`),它将根据梯度更新模型的参数。
在训练循环中,我们使用示例数据进行前向传播,计算损失,并通过反向传播和优化器来更新模型参数。最后,我们打印每一轮的损失值。
如果您想使用拟牛顿法,可以将优化器替换为`torch.optim.LBFGS`,并调整相应的参数。
希望对您有所帮助!如有任何疑问,请随时提问。
阅读全文