From 6664c6d52fad58e396861946a3bed7d5afc4d44d Mon Sep 17 00:00:00 2001
From: natanielruiz <nataniel777@hotmail.com>
Date: 星期五, 07 七月 2017 10:53:52 +0800
Subject: [PATCH] Training for hopenet works.

---
 code/datasets.py          |   49 +++++++++++-
 code/hopenet.py           |   64 ++++++++++++++++
 code/test_resnet_bins.py  |   26 +++--
 code/train_resnet_bins.py |   94 +++++++++++++++++-----
 4 files changed, 195 insertions(+), 38 deletions(-)

diff --git a/code/datasets.py b/code/datasets.py
index 030059f..3750e71 100644
--- a/code/datasets.py
+++ b/code/datasets.py
@@ -84,10 +84,51 @@
 
         # We get the pose in radians
         pose = utils.get_ypr_from_mat(os.path.join(self.data_dir, self.y_train[index] + self.annot_ext))
-        # And convert to positive degrees.
-        pose = pose * 180 / np.pi + 90
+        # And convert to degrees.
+        pitch = pose[0] * 180 / np.pi
+        yaw = pose[1] * 180 / np.pi
+        roll = pose[2] * 180 / np.pi
+        # Bin values
+        bins = np.array(range(-99, 102, 3))
+        labels = torch.LongTensor(np.digitize([yaw, pitch, roll], bins) - 1)
 
-        label = torch.FloatTensor(pose)
+        if self.transform is not None:
+            img = self.transform(img)
+
+        return img, labels, self.X_train[index]
+
+    def __len__(self):
+        # 122,450
+        return self.length
+
+class AFLW2000_binned(Dataset):
+    def __init__(self, data_dir, filename_path, transform, img_ext='.jpg', annot_ext='.mat'):
+        self.data_dir = data_dir
+        self.transform = transform
+        self.img_ext = img_ext
+        self.annot_ext = annot_ext
+
+        filename_list = get_list_from_filenames(filename_path)
+
+        self.X_train = filename_list
+        self.y_train = filename_list
+        self.length = len(filename_list)
+
+    def __getitem__(self, index):
+        img = Image.open(os.path.join(self.data_dir, self.X_train[index] + self.img_ext))
+        img = img.convert('RGB')
+
+        # We get the pose in radians
+        pose = utils.get_ypr_from_mat(os.path.join(self.data_dir, self.y_train[index] + self.annot_ext))
+        # And convert to degrees.
+        pitch, yaw, roll = pose * 180 / np.pi
+        # Bin values
+        bins = np.array(range(-99, 102, 3))
+        binned_pitch = torch.DoubleTensor(np.digitize(pitch, bins) - 1)
+        binned_yaw = torch.DoubleTensor(np.digitize(yaw, bins) - 1)
+        binned_roll = torch.DoubleTensor(np.digitize(roll, bins) - 1)
+
+        label = binned_yaw, binned_pitch, binned_roll
 
         if self.transform is not None:
             img = self.transform(img)
@@ -95,7 +136,7 @@
         return img, label, self.X_train[index]
 
     def __len__(self):
-        # 122,450
+        # 2,000
         return self.length
 
 def get_list_from_filenames(file_path):
diff --git a/code/hopenet.py b/code/hopenet.py
index e6f8f50..e5c1ed2 100644
--- a/code/hopenet.py
+++ b/code/hopenet.py
@@ -2,6 +2,7 @@
 import torch.nn as nn
 import torchvision.datasets as dsets
 from torch.autograd import Variable
+import math
 
 # CNN Model (2 conv layer)
 class Simple_CNN(nn.Module):
@@ -37,3 +38,66 @@
         out = out.view(out.size(0), -1)
         out = self.fc(out)
         return out
+
+class Hopenet(nn.Module):
+    # This is just Hopenet with 3 output layers for yaw, pitch and roll.
+    def __init__(self, block, layers, num_bins):
+        self.inplanes = 64
+        super(Hopenet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7)
+        self.fc_yaw = nn.Linear(512 * block.expansion, num_bins)
+        self.fc_pitch = nn.Linear(512 * block.expansion, num_bins)
+        self.fc_roll = nn.Linear(512 * block.expansion, num_bins)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        yaw = self.fc_yaw(x)
+        pitch = self.fc_pitch(x)
+        roll = self.fc_roll(x)
+
+        return yaw, pitch, roll
diff --git a/code/test_resnet_bins.py b/code/test_resnet_bins.py
index 34fc8f5..0a093ee 100644
--- a/code/test_resnet_bins.py
+++ b/code/test_resnet_bins.py
@@ -13,7 +13,7 @@
 import os
 import argparse
 
-from datasets import AFLW2000
+import datasets
 import hopenet
 import utils
 
@@ -55,9 +55,10 @@
 
     print 'Loading data.'
 
-    transformations = transforms.Compose([transforms.Scale(224),transforms.RandomCrop(224), transforms.ToTensor()])
+    transformations = transforms.Compose([transforms.Scale(224),
+    transforms.RandomCrop(224), transforms.ToTensor()])
 
-    pose_dataset = AFLW2000(args.data_dir, args.filename_list,
+    pose_dataset = datasets.AFLW2000_binned(args.data_dir, args.filename_list,
                                 transformations)
     test_loader = torch.utils.data.DataLoader(dataset=pose_dataset,
                                                batch_size=batch_size,
@@ -69,7 +70,9 @@
 
     # Test the Model
     model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
-    error = .0
+    yaw_correct = 0
+    pitch_correct = 0
+    roll_correct = 0
     total = 0
     for i, (images, labels, name) in enumerate(test_loader):
         images = Variable(images).cuda(gpu)
@@ -78,13 +81,14 @@
         _, predicted = torch.max(outputs.data, 1)
         total += labels.size(0)
         # TODO: There are more efficient ways.
+        yaw_correct += (outputs[:][0] == labels[:][0])
+        pitch_correct += (outputs[:][])
         for idx in xrange(len(outputs)):
-            # if abs(outputs[idx].data[1] - labels[idx].data[1]) * 180 / np.pi > 30:
-            print name
-            print abs(outputs[idx].data - labels[idx].data) * 180 / np.pi, 180 * outputs[idx].data / np.pi, labels[idx].data * 180 / np.pi
-            # error += utils.mse_loss(outputs[idx], labels[idx])
-            error += abs(outputs[idx].data - labels[idx].data) * 180 / np.pi
+            yaw_correct += (outputs[idx].data[0] == labels[idx].data[0])
+            pitch_correct += (outputs[idx].data[1] == labels[idx].data[1])
+            roll_correct += (outputs[idx].data[2] == labels[idx].data[2])
 
 
-    print('Test MSE error of the model on the ' + str(total) +
-    ' test images: %.4f' % (error / total))
+    print('Test accuracies of the model on the ' + str(total) +
+    ' test images. Yaw: %.4f %%, Pitch: %.4f %%, Roll: %.4f %%' % (yaw_correct / total,
+    pitch_correct / total, roll_correct / total))
diff --git a/code/train_resnet_bins.py b/code/train_resnet_bins.py
index f2ec5f2..1bbf5be 100644
--- a/code/train_resnet_bins.py
+++ b/code/train_resnet_bins.py
@@ -13,8 +13,17 @@
 import os
 import argparse
 
-from datasets import Pose_300W_LP
+import datasets
 import hopenet
+import torch.utils.model_zoo as model_zoo
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
 
 def parse_args():
     """Parse input arguments."""
@@ -36,6 +45,41 @@
 
     return args
 
+def get_ignored_params(model):
+    # Generator function that yields ignored params.
+    b = []
+    b.append(model.conv1)
+    b.append(model.bn1)
+    b.append(model.layer1)
+    b.append(model.layer2)
+    b.append(model.layer3)
+    b.append(model.layer4)
+    for i in range(len(b)):
+        for j in b[i].modules():
+            for k in j.parameters():
+                yield k
+
+def get_non_ignored_params(model):
+    # Generator function that yields params that will be optimized.
+    b = []
+    b.append(model.fc_yaw)
+    b.append(model.fc_pitch)
+    b.append(model.fc_roll)
+    for i in range(len(b)):
+        for j in b[i].modules():
+            for k in j.parameters():
+                    yield k
+
+def load_filtered_state_dict(model, snapshot):
+    # By user apaszke from discuss.pytorch.org
+    model_dict = model.state_dict()
+    # 1. filter out unnecessary keys
+    snapshot = {k: v for k, v in snapshot.items() if k in model_dict}
+    # 2. overwrite entries in the existing state dict
+    model_dict.update(snapshot)
+    # 3. load the new state dict
+    model.load_state_dict(model_dict)
+
 if __name__ == '__main__':
     args = parse_args()
 
@@ -47,21 +91,16 @@
     if not os.path.exists('output/snapshots'):
         os.makedirs('output/snapshots')
 
-    model = torchvision.models.resnet18(pretrained=True)
-    for param in model.parameters():
-        param.requires_grad = False
-    # Parameters of newly constructed modules have requires_grad=True by default
-    num_ftrs = model.fc.in_features
-    model.fc_pitch = nn.Linear(num_ftrs, 3)
-    model.fc_yaw = nn.Linear(num_ftrs, 3)
-    model.fc_roll = nn.Linear(num_ftrs, )
-
+    # ResNet18 with 3 outputs.
+    model = hopenet.Hopenet(torchvision.models.resnet.BasicBlock, [2, 2, 2, 2], 66)
+    load_filtered_state_dict(model, model_zoo.load_url(model_urls['resnet18']))
+    
     print 'Loading data.'
 
-    transformations = transforms.Compose([transforms.Scale(230),transforms.RandomCrop(224),
+    transformations = transforms.Compose([transforms.Scale(224),transforms.RandomCrop(224),
                                           transforms.ToTensor()])
 
-    pose_dataset = Pose_300W_LP(args.data_dir, args.filename_list,
+    pose_dataset = datasets.Pose_300W_LP_binned(args.data_dir, args.filename_list,
                                 transformations)
     train_loader = torch.utils.data.DataLoader(dataset=pose_dataset,
                                                batch_size=batch_size,
@@ -69,31 +108,40 @@
                                                num_workers=2)
 
     model.cuda(gpu)
-    criterion = nn.MSELoss(size_average = True)
-    optimizer = torch.optim.Adam(model.fc.parameters(), lr = args.lr)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': .0},
+                                  {'params': get_non_ignored_params(model), 'lr': args.lr}],
+                                  lr = args.lr)
 
     print 'Ready to train network.'
 
     for epoch in range(num_epochs):
-        for i, (images, labels) in enumerate(train_loader):
+        for i, (images, labels, name) in enumerate(train_loader):
             images = Variable(images).cuda(gpu)
-            labels = Variable(labels).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()
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            loss.backward()
+            yaw, pitch, roll = model(images)
+            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()
 
             if (i+1) % 100 == 0:
-                print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f'
-                       %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss.data[0]))
+                print ('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f'
+                       %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss_yaw.data[0], loss_pitch.data[0], loss_roll.data[0]))
 
         # Save models at even numbered epochs.
         if epoch % 1 == 0 and epoch < num_epochs - 1:
             print 'Taking snapshot...'
             torch.save(model.state_dict(),
-            'output/snapshots/resnet18_epoch_' + str(epoch+1) + '.pkl')
+            'output/snapshots/resnet18_binned_epoch_' + str(epoch+1) + '.pkl')
 
     # Save the final Trained Model
-    torch.save(model.state_dict(), 'output/snapshots/resnet18_epoch_' + str(epoch+1) + '.pkl')
+    torch.save(model.state_dict(), 'output/snapshots/resnet18_binned_epoch_' + str(epoch+1) + '.pkl')

--
Gitblit v1.8.0