pytorch 自动求导 【初始化梯度】

print('开始训练')
for epoch in range(3):
    runing_loss = 0.0

    for i,data in enumerate(trainloader,0):
    inputs,label = data             #1.数据加载
        if device == 'gpu':
            inputs = inputs.cuda()
            label = label.cuda()
        optimizer.zero_grad()           #2.初始化梯度
        output = Net(inputs)            #3.计算前馈
        loss = criterion(output,label)  #4.计算损失
        loss.backward()                 #5.计算梯度
        optimizer.step()                #6.更新权值

        runing_loss += loss.item()
        if i % 20 == 19:
            print('epoch:',epoch,'loss',runing_loss/20)
            runing_loss = 0.0

print('训练完成')

Reference

《Pytorch深度学习实战》 5.5.4 自动求导更新及关闭 Page 120

posted @ 2024-03-19 00:13  光辉233  阅读(4)  评论(0编辑  收藏  举报