def __init__(self):
super(TestNet, self).__init__()
self.pre = Hourglass(
3, 1, increase=10) #downsample the feature map for 5 times
self.location = Hourglass(5, 1, increase=15)
self.xVector = Hourglass(
5, 3, increase=10
) #predict the offset from corners to center along x axis
self.yVector = Hourglass(
5, 3, increase=10
) # predict the offset from corners to center along y axis
self.conv = Conv(3, 1, kernel_size=1, stride=1)
self.conv2 = Conv(3, 1, kernel_size=1, stride=1, relu=False)
self.conv3 = Conv(3, 1, kernel_size=1, stride=1, relu=False)
def __init__(self,
nstack,
inp_dim,
oup_dim,
bn=False,
increase=128,
**kwargs):
super(PoseNet, self).__init__()
self.pre = nn.Sequential(Conv(3, 64, 7, 2, bn=bn), Conv(64, 128,
bn=bn),
Pool(2, 2), Conv(128, 128, bn=bn),
Conv(128, inp_dim, bn=bn))
self.features = nn.ModuleList([
nn.Sequential(Hourglass(4, inp_dim, bn, increase),
Conv(inp_dim, inp_dim, 3, bn=False),
Conv(inp_dim, inp_dim, 3, bn=False))
for i in range(nstack)
])
self.outs = nn.ModuleList([
Conv(inp_dim, oup_dim, 1, relu=False, bn=False)
for i in range(nstack)
])
self.merge_features = nn.ModuleList(
[Merge(inp_dim, inp_dim) for i in range(nstack - 1)])
self.merge_preds = nn.ModuleList(
[Merge(oup_dim, inp_dim) for i in range(nstack - 1)])
self.nstack = nstack
self.myAEloss = tagLoss
self.heatmapLoss = HeatmapLoss()
def __init__(self,
nstack,
inp_dim,
oup_dim,
bn=False,
increase=128,
**kwargs):
super(EdgeNet, self).__init__()
self.pre = nn.Sequential(
#Conv(3, 64, 7, 2, bn=bn),
Conv(3, 64, 7, 1, bn=bn),
Conv(64, 128, bn=bn),
#Pool(2, 2),
Conv(128, 128, bn=bn),
Conv(128, inp_dim, bn=bn))
self.features = nn.ModuleList([
nn.Sequential(
Hourglass(5, inp_dim, bn, increase), # Orig 4
Conv(inp_dim, inp_dim, 3, bn=False),
Conv(inp_dim, inp_dim, 3, bn=False)) for i in range(nstack)
])
self.outs = nn.ModuleList([
Conv(inp_dim, oup_dim, 1, relu=False, bn=False)
for i in range(nstack)
])
self.merge_features = nn.ModuleList(
[Merge(inp_dim, inp_dim) for i in range(nstack - 1)])
self.merge_preds = nn.ModuleList(
[Merge(oup_dim, inp_dim) for i in range(nstack - 1)])
self.nstack = nstack
def __init__(self,
nstack,
nfeatures,
nlandmarks,
bn=False,
increase=0,
**kwargs):
super(EyeNet, self).__init__()
self.img_w = 160
self.img_h = 96
self.nstack = nstack
self.nfeatures = nfeatures
self.nlandmarks = nlandmarks
self.heatmap_w = self.img_w / 2
self.heatmap_h = self.img_h / 2
self.nstack = nstack
self.pre = nn.Sequential(Conv(1, 64, 7, 1, bn=True, relu=True),
Residual(64, 128), Pool(2, 2),
Residual(128, 128), Residual(128, nfeatures))
self.pre2 = nn.Sequential(
Conv(nfeatures, 64, 7, 2, bn=True, relu=True), Residual(64, 128),
Pool(2, 2), Residual(128, 128), Residual(128, nfeatures))
self.hgs = nn.ModuleList([
nn.Sequential(Hourglass(4, nfeatures, bn, increase), )
for i in range(nstack)
])
self.features = nn.ModuleList([
nn.Sequential(Residual(nfeatures, nfeatures),
Conv(nfeatures, nfeatures, 1, bn=True, relu=True))
for i in range(nstack)
])
self.outs = nn.ModuleList([
Conv(nfeatures, nlandmarks, 1, relu=False, bn=False)
for i in range(nstack)
])
self.merge_features = nn.ModuleList(
[Merge(nfeatures, nfeatures) for i in range(nstack - 1)])
self.merge_preds = nn.ModuleList(
[Merge(nlandmarks, nfeatures) for i in range(nstack - 1)])
self.gaze_fc1 = nn.Linear(
in_features=int(nfeatures * self.img_w * self.img_h / 64 +
nlandmarks * 2),
out_features=256)
self.gaze_fc2 = nn.Linear(in_features=256, out_features=2)
self.nstack = nstack
self.heatmapLoss = HeatmapLoss()
self.landmarks_loss = nn.MSELoss()
self.gaze_loss = nn.MSELoss()
def main():
print('test')
train = SurfaceNormalsDataset('../data/train', test=False)
trainloader = torch.utils.data.DataLoader(train,
batch_size=25,
shuffle=True)
model = torch.nn.Sequential(
Conv(3, 10),
Hourglass(4, 10, bn=None, increase=20),
Conv(10, 10),
Conv(10, 3),
)
criterion = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
for epoch in range(3): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs
inputs = data['image']
labels = data['normal']
# wrap them in Variable
inputs, labels = Variable(inputs), Variable(labels)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs)
#print(outputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.data[0]
if i % 10 == 9: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss /
(10 * 25 * 128 * 128 * 3)))
running_loss = 0.0
if i == 100:
pass
#break
torch.save(model, '../models/hg_test')
print('Finished Training')
def __init__(self,
nstack,
inp_dim,
oup_dim,
bn=False,
increase=128,
init_weights=True,
**kwargs):
"""
Pack or initialize the trainable parameters of the network
:param nstack: number of stack
:param inp_dim: input tensor channels fed into the hourglass block
:param oup_dim: channels of regressed feature maps
:param bn:
:param increase: increased channels once down-sampling
:param kwargs:
"""
super(PoseNet, self).__init__()
self.pre = nn.Sequential(Conv(3, 64, 7, 2, bn=bn), Conv(64, 128,
bn=bn),
nn.MaxPool2d(2, 2), Conv(128, 128, bn=bn),
Conv(128, inp_dim, bn=bn))
self.hourglass = nn.ModuleList(
[Hourglass(4, inp_dim, increase, bn=bn) for _ in range(nstack)])
self.features = nn.ModuleList([
Features(inp_dim, increase=increase, bn=bn) for _ in range(nstack)
])
# predict 5 different scales of heatmpas per stack, keep in mind to pack the list using ModuleList.
# Notice: nn.ModuleList can only identify Module subclass! Thus, we must pack the inner layers in ModuleList.
self.outs = nn.ModuleList([
nn.ModuleList([
Conv(inp_dim + j * increase, oup_dim, 1, relu=False, bn=False)
for j in range(5)
]) for i in range(nstack)
])
self.merge_features = nn.ModuleList([
nn.ModuleList([
Merge(inp_dim + j * increase, inp_dim + j * increase)
for j in range(5)
]) for i in range(nstack - 1)
])
self.merge_preds = nn.ModuleList([
nn.ModuleList(
[Merge(oup_dim, inp_dim + j * increase) for j in range(5)])
for i in range(nstack - 1)
])
self.nstack = nstack
if init_weights:
self._initialize_weights()
def __init__(self,
nstack,
inp_dim,
oup_dim,
num_class=9,
bn=False,
increase=0,
**kwargs):
super(PoseNet, self).__init__()
self.nstack = nstack
self.pre = nn.Sequential(Conv(1, 64, 7, 2, bn=True, relu=True),
Residual(64, 128), Pool(2, 2),
Residual(128, 128), Residual(128, inp_dim))
self.hgs = nn.ModuleList([
nn.Sequential(Hourglass(4, inp_dim, bn, increase), )
for i in range(nstack)
])
self.features = nn.ModuleList([
nn.Sequential(
Residual(inp_dim, inp_dim),
# Residual(inp_dim, inp_dim),
Conv(inp_dim, inp_dim, 1, bn=True, relu=True))
for i in range(nstack)
])
self.outs_heatmap = nn.ModuleList([
Conv(inp_dim, oup_dim, 1, relu=True, bn=True)
for i in range(nstack)
])
self.outs_label = nn.ModuleList([
nn.Sequential(
Residual(inp_dim + oup_dim, inp_dim + oup_dim),
Conv(inp_dim + oup_dim, num_class, 1, relu=True, bn=True))
for i in range(nstack)
])
self.merge_features = nn.ModuleList(
[Merge(inp_dim, inp_dim) for i in range(nstack - 1)])
self.merge_preds = nn.ModuleList(
[Merge(oup_dim, inp_dim) for i in range(nstack - 1)])
self.nstack = nstack
def __init__(self,
nstack,
inp_dim,
oup_dim,
bn=False,
increase=128,
**kwargs):
super(PoseNet, self).__init__()
# 'nn.Sequential' is a Container that contains each Module to construct the layer sequence ofCNN
self.pre = nn.Sequential(
# Conv() parameter is 'inp_dim, out_dim, kernel_size, stride' in layers.py
# uses module 'nn.Conv2d'
Conv(3, 64, 7, 2, bn=bn),
Conv(64, 128, bn=bn),
Pool(2, 2),
Conv(128, 128, bn=bn),
Conv(128, inp_dim, bn=bn))
self.features = nn.ModuleList([
nn.Sequential(Hourglass(4, inp_dim, bn, increase),
Conv(inp_dim, inp_dim, 3, bn=False),
Conv(inp_dim, inp_dim, 3, bn=False))
for i in range(nstack)
])
# build for number of the nstack layers
# the 'nn.ModuleList' is to store nn.Modules, similar to python list, to pass as input
self.outs = nn.ModuleList([
Conv(inp_dim, oup_dim, 1, relu=False, bn=False)
for i in range(nstack)
])
self.merge_features = nn.ModuleList(
[Merge(inp_dim, inp_dim) for i in range(nstack - 1)])
self.merge_preds = nn.ModuleList(
[Merge(oup_dim, inp_dim) for i in range(nstack - 1)])
self.nstack = nstack
self.myAEloss = AEloss()
self.heatmapLoss = HeatmapLoss()
def __init__(self, feature_dim, predicate_dim, object_dim, pretrained,
dropout, num_layers, backbone, two_stream, rgbd, only_2d,
only_appr):
super(DRNet, self).__init__()
self.num_layers = num_layers
self.predicate_dim = predicate_dim
self.object_dim = object_dim
self.dropout = dropout
# if both two_stream and rgdb are false, we use only rgb input
self.two_stream = two_stream
self.rgbd = rgbd
# if both only_2d and only_appr are false, then both streams are used
self.only_2d = only_2d
self.only_appr = only_appr
if self.only_2d:
assert not (self.rgbd or self.two_stream)
assert not (self.rgbd and self.two_stream)
output_size = 0
if not self.only_appr:
self.pos_module = nn.Sequential(
OrderedDict([
('conv1_p', nn.Conv2d(2, 32, 5, 2, 2)),
('normalize1_p', nn.BatchNorm2d(32)),
('relu1_p', nn.ReLU()),
('conv2_p', nn.Conv2d(32, 64, 3, 1, 1)),
('normalize2_p', nn.BatchNorm2d(64)),
('relu2_p', nn.ReLU()),
('maxpool2_p', nn.MaxPool2d(2)),
('hg', Hourglass(8, 64)),
('normalize_p', nn.BatchNorm2d(64)),
('relu_p', nn.ReLU()),
('maxpool_p', nn.MaxPool2d(2)),
('conv3_p', nn.Conv2d(64, 256, 4)),
('normalize3_p', nn.BatchNorm2d(256)),
]))
output_size += 256
if not self.only_2d:
if self.rgbd:
if pretrained:
print(
"Using pretrained model with RGBD input, this will work but might not be the correct strategy"
)
_module = torchvision.models.__dict__[backbone](
pretrained=pretrained)
_module.fc = nn.Linear(512, 256)
self.appr_module = nn.Sequential(
OrderedDict([('conv1', nn.Conv2d(4, 4, 3, 1, 1)),
('normalize1', nn.BatchNorm2d(4)),
('relu1', nn.ReLU()),
('conv2', nn.Conv2d(4, 4, 3, 1, 1)),
('normalize2', nn.BatchNorm2d(4)),
('relu2', nn.ReLU()),
('conv3', nn.Conv2d(4, 3, 3, 1, 1)),
('normalize3', nn.BatchNorm2d(3)),
('relu3', nn.ReLU()), ('model', _module)]))
output_size += 256
elif self.two_stream:
self.appr_module = torchvision.models.__dict__[backbone](
pretrained=pretrained)
self.appr_module.fc = nn.Linear(512, 256)
self.depth_module = torchvision.models.__dict__[backbone](
pretrained=False)
self.depth_module.fc = nn.Linear(512, 256)
output_size += 512
else:
self.appr_module = torchvision.models.__dict__[backbone](
pretrained=pretrained)
self.appr_module.fc = nn.Linear(512, 256)
output_size += 256
self.PhiR_0 = nn.Linear(output_size, feature_dim)
self.normalize = nn.BatchNorm1d(feature_dim)
self.PhiA = nn.Linear(object_dim, feature_dim)
self.PhiB = nn.Linear(object_dim, feature_dim)
self.PhiR = nn.Linear(feature_dim, feature_dim)
self.fc = nn.Linear(feature_dim, predicate_dim)