output = lenet_model(batch_img) #进行前向计算
# print output.data.size() #output输出为batch_size*10
# print batch_label.data.size()
loss = loss_func(output,
batch_label) #计算loss 此loss已经在一个batch_size上做了平均
#print loss.data.size() loss输出为一个标量
train_loss += loss.item()
max, max_index = torch.max(output, 1) #返回1维度上的最大值 记忆下标
#print max.size()
train_correct = torch.sum(
(max_index.data == batch_label.data)) #统计一个batch中预测正确的数量
# print "train_correct:",train_correct
train_acc += train_correct
# print "train_acc:",train_acc
#反向传播
lenet_model.zero_grad() #因为梯度是累计的 所以先将梯度清0
loss.backward() #将自动计算所有可学习参数的梯度
optimizer.step() #调用优化器来更新参数
#显示设置
if (num_iter % iter_display == 0):
print "iter_num:", num_iter
if num_iter == 0:
train_loss = train_loss
train_acc = float(train_acc) / batch_size
else:
train_loss = train_loss / iter_display
train_acc = float(train_acc) / (iter_display * batch_size)
recoder["num_iter"].append(num_iter)
recoder["train_loss"].append(train_loss)
recoder["train_acc"].append(train_acc)
