代码中的学习率

发布时间 2023-11-13 21:48:21作者: 黑逍逍
import torch
import torch.nn as nn
import torch.optim as optim

# 假设有一个简单的神经网络模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, 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):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型实例
input_size = 20
hidden_size = 64
output_size = 10
model = SimpleModel(input_size, hidden_size, output_size)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 设置学习率和优化器
learning_rate = 0.001
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 在训练过程中使用设置的学习率进行参数更新
for epoch in range(num_epochs):
    # 前向传播、计算损失等步骤

    # 反向传播和优化步骤
    optimizer.zero_grad()  # 清零梯度
    loss.backward()  # 反向传播
    optimizer.step()  # 参数更新

 

使用optimizer.zero_grad()来清零梯度,然后进行反向传播和参数更新