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