Skip to content

Commit

Permalink
next
Browse files Browse the repository at this point in the history
  • Loading branch information
counterfactualsimulation committed Sep 26, 2017
1 parent 790c090 commit c13dba8
Show file tree
Hide file tree
Showing 3 changed files with 2 additions and 312 deletions.
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@
*.npy
data/*
output/*
*.jpg
*.jpg
*.png
90 changes: 0 additions & 90 deletions code/hopenet.py
Original file line number Diff line number Diff line change
Expand Up @@ -340,93 +340,3 @@ def forward(self, x):
angles.append(preangles)

return pre_yaw, pre_pitch, pre_roll, angles, sr_output

class Hopenet_LSTM(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_LSTM, 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)

self.softmax = nn.Softmax()
self.fc_finetune = nn.Linear(512 * block.expansion + 3, 3)

self.idx_tensor = Variable(torch.FloatTensor(range(66))).cuda()

self.lstm = nn.LSTM(512 * block.expansion + 3, 256 * block.expansion, 2, batch_first=True)
self.fc_lstm = nn.Linear(256 * block.expansion, 3)

self.block_expansion = block.expansion

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)
pre_yaw = self.fc_yaw(x)
pre_pitch = self.fc_pitch(x)
pre_roll = self.fc_roll(x)

# Yaw, pitch, roll
yaw = self.softmax(pre_yaw)
yaw = Variable(torch.sum(yaw.data * self.idx_tensor.data, 1), requires_grad=True) * 3 - 99
pitch = self.softmax(pre_pitch)
pitch = Variable(torch.sum(pitch.data * self.idx_tensor.data, 1), requires_grad=True) * 3 - 99
roll = self.softmax(pre_roll)
roll = Variable(torch.sum(roll.data * self.idx_tensor.data, 1), requires_grad=True) * 3 - 99
yaw = yaw.view(yaw.size(0), 1)
pitch = pitch.view(pitch.size(0), 1)
roll = roll.view(roll.size(0), 1)
preangles = torch.cat([yaw, pitch, roll], 1)

residuals, _ = self.lstm(torch.cat((preangles, x), 1), (h0, c0))
residuals = self.fc_lstm(residuals[:, -1, :])
final_angles = preangles + residuals

return pre_yaw, pre_pitch, pre_roll, preangles, final_angles
221 changes: 0 additions & 221 deletions code/train_hopenet_lstm.py

This file was deleted.

0 comments on commit c13dba8

Please sign in to comment.