stage = [[3, 3, 3, 1, 1], [7, 7, 7, 7, 7, 1, 1]]branches_cfg = [[[128, 128, 128, 512, 38], [128, 128, 128, 512, 19]], [[128, 128, 128, 128, 128, 128, 38], [128, 128, 128, 128, 128, 128, 19]]]# used for add two branches as well as adapt to certain stagedef add_extra(i, branches_cfg, stage): """ only add CNN of brancdes S & L in stage Ti at the end of net :param in_channels:the input channels & out :param stage: size of filter :param branches_cfg: channels of image :return:list of layers """ in_channels = i layers = [] for k in range(len(stage)): padding = stage[k] // 2 conv2d = nn.Conv2d(in_channels, branches_cfg[k], kernel_size=stage[k], padding=padding) layers += [conv2d, nn.ReLU(inplace=True)] in_channels = branches_cfg[k] return layers
conf_bra_list = []paf_bra_list = []# param for branch networkin_channels = 128for i in range(all_stage): if i > 0: branches = branches_cfg[1] conv_sz = stage[1] else: branches = branches_cfg[0] conv_sz = stage[0] conf_bra_list.append(nn.Sequential(*add_extra(in_channels, branches[0], conv_sz))) paf_bra_list.append(nn.Sequential(*add_extra(in_channels, branches[1], conv_sz))) in_channels = 185
out_0 = x# the base transformfor k in range(len(self.vgg)): out_0 = self.vgg[k](out_0)# local name spacename = locals()confs = []pafs = []outs = []length = len(self.conf_bra)for i in range(length): name['conf_%s' % (i + 1)] = self.conf_bra[i](name['out_%s' % i]) name['paf_%s' % (i + 1)] = self.paf_bra[i](name['out_%s' % i]) name['out_%s' % (i + 1)] = torch.cat([name['conf_%s' % (i + 1)], name['paf_%s' % (i + 1)], out_0], 1) confs.append('conf_%s' % (i + 1)) pafs.append('paf_%s' % (i + 1)) outs.append('out_%s' % (i + 1)) 
