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