Pytorch网络训练流程的作用原理：源码分析optimizer.zero_grad()loss.backward()optimizer.step()

原创

于 2021-08-25 10:52:53 发布 · 3.9k 阅读

CC 4.0 BY-SA版权

本文详细剖析了PyTorch中optimizer.zero_grad(), loss.backward(), optimizer.step()的核心作用，包括梯度清零、反向传播和参数更新的过程，以及它们在训练模型时的优缺点。从理论到实践，适合理解优化器工作原理的开发者阅读。

常见参数训练流程的作用原理

1 总述

在用pytorch训练模型时，通常会在循环epoch的过程中，不断循环遍历所有训练数据集。

依次用到optimizer.zero_grad()，loss.backward()和optimizer.step()三个函数，如下所示：

model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)  
for epoch in range(1, epochs):
    for i, (inputs, labels) in enumerate(train_loader):
        output= model(inputs)
        loss = criterion(output, labels)
        # compute gradient and do SGD step     
        optimizer.zero_grad()     
        loss.backward()     
    	optimizer.step

最低0.47元/天解锁文章