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(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()