print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
print('Learning Rate for this epoch: {}'.format(learning_rate))
total_loss = 0.
# train by iter on every epoch
for i, (images, target) in enumerate(train_loader):
images = Variable(images)
target = Variable(target)
if USE_GPU:
images,target = images.cuda(),target.cuda()
# forward
pred = net(images) # [1,7,7,30]
loss = criterion(pred,target)
total_loss += loss.item()
vis.plot_train_val(loss_train=loss.item()) # train plot
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# logout
if (i+1) % print_every_step == 0:
print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'%(epoch+1, num_epochs, i+1, len(train_loader), loss.item(), total_loss / (i+1)))
num_iter += 1
# save checkpoint model
torch.save(net.state_dict(), './output/%s.pth' % net.name)
# validate
images, target = images.cuda(), target.cuda()
pred = net(images)
loss = criterion(pred, target)
total_loss += loss.data[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 5 == 0:
print(
'Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f' %
(epoch + 1, num_epochs, i + 1, len(train_loader), loss.data[0],
total_loss / (i + 1)))
num_iter += 1
vis.plot_train_val(loss_train=total_loss / (i + 1))
#validation
validation_loss = 0.0
net.eval()
for i, (images, target) in enumerate(test_loader):
images = Variable(images, volatile=True)
target = Variable(target, volatile=True)
if use_gpu:
images, target = images.cuda(), target.cuda()
pred = net(images)
loss = criterion(pred, target)
validation_loss += loss.data[0]
validation_loss /= len(test_loader)
vis.plot_train_val(loss_val=validation_loss)
print('Learning Rate for this epoch: {}'.format(learning_rate))
total_loss = 0.
for i, (images, target) in enumerate(train_loader):
images = Variable(images)
target = Variable(target)
if use_gpu:
images, target = images.cuda(), target.cuda()
pred = net(images)
print(pred.size(), target.size())
loss = criterion(pred, target)
total_loss += float(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 10 == 0:
print(
'Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f' %
(epoch + 1, num_epochs, i + 1, len(train_loader), loss.item(),
total_loss / (i + 1)))
num_iter += 1
vis.plot_train_val(loss_train=total_loss / (i + 1))
torch.save(
net.state_dict(),
'/home/duanyajun/文档/目标识别项目/自己改写代码/mobilenet_yolov1/yolo_test_boat.pth')
