def __init__(self,
num_refinement_stages=1,
num_channels=128,
num_heatmaps=17,
to_onnx=False):
super().__init__()
self.model = nn.Sequential(
conv(3, 32, stride=2, bias=False),
conv_dw(32, 64),
conv_dw(64, 128, stride=2),
conv_dw(128, 128),
conv_dw(128, 256, stride=2),
conv_dw(256, 256),
conv_dw(256, 512), # conv4_2
conv_dw(512, 512, dilation=2, padding=2),
conv_dw(512, 512),
conv_dw(512, 512),
conv_dw(512, 512),
conv_dw(512, 512), # conv5_5
)
self.cpm = nn.Sequential(
conv(512, 256),
conv(256, 128),
)
self.initial_stage = InitialStage(num_channels, num_heatmaps)
self.refinement_stages = nn.ModuleList()
for idx in range(num_refinement_stages):
self.refinement_stages.append(
RefinementStage(num_channels + num_heatmaps, num_channels,
num_heatmaps, to_onnx))
def __init__(self, in_channels, out_channels):
super().__init__()
self.initial = conv(in_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False)
self.trunk = nn.Sequential(
conv(out_channels, out_channels),
conv(out_channels, out_channels, dilation=2, padding=2))
def __init__(self, num_channels, num_heatmaps):
super().__init__()
self.trunk = nn.Sequential(conv(num_channels, num_channels),
conv(num_channels, num_channels),
conv(num_channels, num_channels))
self.heatmaps = nn.Sequential(
conv(num_channels, 512, kernel_size=1, padding=0, bn=False),
conv(512,
num_heatmaps,
kernel_size=1,
padding=0,
bn=False,
relu=False))
def __init__(self, in_channels, out_channels, num_heatmaps, to_onnx):
super().__init__()
self.trunk = nn.Sequential(
UShapedContextBlock(in_channels, to_onnx),
RefinementStageBlock(in_channels, out_channels),
RefinementStageBlock(out_channels, out_channels),
RefinementStageBlock(out_channels, out_channels),
)
self.heatmaps = nn.Sequential(
conv(out_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False),
conv(out_channels,
num_heatmaps,
kernel_size=1,
padding=0,
bn=False,
relu=False))
def __init__(self, in_channels, to_onnx=False):
super().__init__()
self.to_onnx = to_onnx
self.encoder1 = nn.Sequential(
conv(in_channels, in_channels * 2, stride=2),
conv(in_channels * 2, in_channels * 2),
)
self.encoder2 = nn.Sequential(
conv(in_channels * 2, in_channels * 2, stride=2),
conv(in_channels * 2, in_channels * 2),
)
self.decoder2 = nn.Sequential(
conv(in_channels * 2 + in_channels * 2, in_channels * 2),
conv(in_channels * 2, in_channels * 2),
)
self.decoder1 = nn.Sequential(conv(in_channels * 3, in_channels * 2),
conv(in_channels * 2, in_channels))