Method forward

src/modules/pixie/models/hrnet.py:425–497

(self, x)

        return nn.Sequential(*modules), num_inchannels

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.layer1(x)

        x_list = []
        for i in range(self.stage2_cfg['num_branches']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)
        y_list = self.stage2(x_list)

        x_list = []
        for i in range(self.stage3_cfg['num_branches']):
            if self.transition2[i] is not None:
                if i < self.stage2_cfg['num_branches']:
                    x_list.append(self.transition2[i](y_list[i]))
                else:
                    x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage3(x_list)

        x_list = []
        for i in range(self.stage4_cfg['num_branches']):
            if self.transition3[i] is not None:
                if i < self.stage3_cfg['num_branches']:
                    x_list.append(self.transition3[i](y_list[i]))
                else:
                    x_list.append(self.transition3[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        if not self.use_old_impl:
            y_list = self.stage4(x_list)

        output = {}
        for idx, x in enumerate(y_list):
            output[f'layer{idx + 1}'] = x

        feat_list = []
        if self.use_old_impl:
            x3 = self.subsample_3(x_list[1])
            x2 = self.subsample_2(x_list[2])
            x1 = x_list[3]
            feat_list = [x3, x2, x1]
        else:
            x4 = self.subsample_4(y_list[0])
            x3 = self.subsample_3(y_list[1])
            x2 = self.subsample_2(y_list[2])
            x1 = y_list[3]
            feat_list = [x4, x3, x2, x1]

        xf = self.conv_layers(torch.cat(feat_list, dim=1))