From fdf1fedb0d3b4beb672464a438c22b94b9cb7d0f Mon Sep 17 00:00:00 2001 From: natanielruiz <nataniel777@hotmail.com> Date: 星期六, 12 八月 2017 10:59:08 +0800 Subject: [PATCH] Cleanup --- code/train.py | 35 ++++++++++++----------------------- 1 files changed, 12 insertions(+), 23 deletions(-) diff --git a/code/train_resnet_bins_comb.py b/code/train.py similarity index 85% rename from code/train_resnet_bins_comb.py rename to code/train.py index eb23590..ec9e63f 100644 --- a/code/train_resnet_bins_comb.py +++ b/code/train.py @@ -113,11 +113,10 @@ num_workers=2) model.cuda(gpu) - criterion = nn.CrossEntropyLoss() - reg_criterion = nn.MSELoss() + criterion = nn.CrossEntropyLoss().cuda() + reg_criterion = nn.MSELoss().cuda() # Regression loss coefficient - alpha = 0.1 - lsm = nn.Softmax() + alpha = 0.01 idx_tensor = [idx for idx in xrange(66)] idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu) @@ -125,33 +124,25 @@ optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': args.lr}, {'params': get_non_ignored_params(model), 'lr': args.lr * 10}], lr = args.lr) - # optimizer = torch.optim.SGD([{'params': get_ignored_params(model), 'lr': args.lr}, - # {'params': get_non_ignored_params(model), 'lr': args.lr}], - # lr = args.lr, momentum=0.9, weight_decay=5e-4) - # optimizer = torch.optim.RMSprop([{'params': get_ignored_params(model), 'lr': args.lr}, - # {'params': get_non_ignored_params(model), 'lr': args.lr * 10}], - # lr = args.lr) print 'Ready to train network.' for epoch in range(num_epochs): for i, (images, labels, name) in enumerate(train_loader): - images = Variable(images).cuda(gpu) - label_yaw = Variable(labels[:,0]).cuda(gpu) - label_pitch = Variable(labels[:,1]).cuda(gpu) - label_roll = Variable(labels[:,2]).cuda(gpu) + images = Variable(images.cuda(gpu)) + label_yaw = Variable(labels[:,0].cuda(gpu)) + label_pitch = Variable(labels[:,1].cuda(gpu)) + label_roll = Variable(labels[:,2].cuda(gpu)) optimizer.zero_grad() + model.zero_grad() + yaw, pitch, roll = model(images) + # Cross entropy loss loss_yaw = criterion(yaw, label_yaw) loss_pitch = criterion(pitch, label_pitch) loss_roll = criterion(roll, label_roll) - - # loss_seq = [loss_yaw, loss_pitch, loss_roll] - # grad_seq = [torch.Tensor(1).cuda(gpu) for _ in range(len(loss_seq))] - # torch.autograd.backward(loss_seq, grad_seq) - # optimizer.step() # MSE loss yaw_predicted = F.softmax(yaw) @@ -166,15 +157,13 @@ loss_reg_pitch = reg_criterion(pitch_predicted, label_pitch.float()) loss_reg_roll = reg_criterion(roll_predicted, label_roll.float()) - # print yaw_predicted[0], label_yaw.data[0] - + # Total loss loss_yaw += alpha * loss_reg_yaw loss_pitch += alpha * loss_reg_pitch loss_roll += alpha * loss_reg_roll loss_seq = [loss_yaw, loss_pitch, loss_roll] grad_seq = [torch.Tensor(1).cuda(gpu) for _ in range(len(loss_seq))] - model.zero_grad() torch.autograd.backward(loss_seq, grad_seq) optimizer.step() @@ -195,4 +184,4 @@ 'output/snapshots/resnet50_epoch_'+ str(epoch+1) + '.pkl') # Save the final Trained Model - torch.save(model.state_dict(), 'output/snapshots/resnet50_epoch_' + str(epoch+1) + '.pkl') + torch.save(model.state_dict(), 'output/snapshots/resnet50_epoch' + str(epoch+1) + '.pkl') -- Gitblit v1.8.0