From 4b67b5c8ed5566ec3030d537536282e830d87e40 Mon Sep 17 00:00:00 2001
From: natanielruiz <nruiz9@gatech.edu>
Date: 星期一, 30 十月 2017 07:15:49 +0800
Subject: [PATCH] next

---
 code/train_preangles.py |   14 +++++++++++---
 1 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/code/train_preangles.py b/code/train_preangles.py
index 1fe626c..600a9ae 100644
--- a/code/train_preangles.py
+++ b/code/train_preangles.py
@@ -124,6 +124,10 @@
     # Regression loss coefficient
     alpha = args.alpha
 
+    softmax = nn.Softmax().cuda(gpu)
+    idx_tensor = [idx for idx in xrange(66)]
+    idx_tensor = Variable(torch.FloatTensor(idx_tensor)).cuda(gpu)
+
     optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': 0},
                                   {'params': get_non_ignored_params(model), 'lr': args.lr},
                                   {'params': get_fc_params(model), 'lr': args.lr * 5}],
@@ -153,9 +157,13 @@
             loss_roll = criterion(roll, label_roll)
 
             # MSE loss
-            yaw_predicted = angles[:,0]
-            pitch_predicted = angles[:,1]
-            roll_predicted = angles[:,2]
+            yaw_predicted = softmax(yaw)
+            pitch_predicted = softmax(pitch)
+            roll_predicted = softmax(roll)
+
+            yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1) * 3 - 99
+            pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1) * 3 - 99
+            roll_predicted = torch.sum(roll_predicted * idx_tensor, 1) * 3 - 99
 
             loss_reg_yaw = reg_criterion(yaw_predicted, label_yaw_cont)
             loss_reg_pitch = reg_criterion(pitch_predicted, label_pitch_cont)

--
Gitblit v1.8.0