El módulo extractor de características de P2PNet se construye sobre la conocida arquitectura VGG16. A continuación se detalla la implementación del componente troncal, el cual se encarga de procesar los tensores de entrada y extraer representaciones jerárquicas.

Implementación del backbone

El diseño se compone de una clase base que encapsula las capas convolucionales de VGG16, con soporte para retornar capas intermedias cuando se requiere una arquitectura tipo FPN. La distinción principal entre vgg16 y vgg16_bn radica en la inclusión de capas de normalización por lotes (BatchNorm) después de cada operación de convolución.

class VGGFeatureExtractor(nn.Module):
    def __init__(self, vgg_model: nn.Module, output_channels: int, variant: str, extract_intermediate: bool):
        super().__init__()
        conv_blocks = list(vgg_model.features.children())
        
        self.extract_intermediate = extract_intermediate
        self.output_channels = output_channels
        
        if extract_intermediate:
            if variant == 'vgg16_bn':
                self.stage1 = nn.Sequential(*conv_blocks[:13])
                self.stage2 = nn.Sequential(*conv_blocks[13:23])
                self.stage3 = nn.Sequential(*conv_blocks[23:33])
                self.stage4 = nn.Sequential(*conv_blocks[33:43])
            else:
                self.stage1 = nn.Sequential(*conv_blocks[:9])
                self.stage2 = nn.Sequential(*conv_blocks[9:16])
                self.stage3 = nn.Sequential(*conv_blocks[16:23])
                self.stage4 = nn.Sequential(*conv_blocks[23:30])
        else:
            cutoff = 44 if variant == 'vgg16_bn' else 30
            self.feature_pipe = nn.Sequential(*conv_blocks[:cutoff])

    def forward(self, input_tensor):
        feature_maps = []
        
        if self.extract_intermediate:
            current = input_tensor
            for block in [self.stage1, self.stage2, self.stage3, self.stage4]:
                current = block(current)
                feature_maps.append(current)
        else:
            current = self.feature_pipe(input_tensor)
            feature_maps.append(current)
            
        return feature_maps


class VGGBackbone(VGGFeatureExtractor):
    def __init__(self, variant: str, extract_intermediate: bool):
        if variant == 'vgg16_bn':
            pretrained_vgg = models.vgg16_bn(pretrained=True)
        elif variant == 'vgg16':
            pretrained_vgg = models.vgg16(pretrained=True)
        else:
            raise ValueError(f"Variante no soportada: {variant}")
            
        channel_dim = 256
        super().__init__(pretrained_vgg, channel_dim, variant, extract_intermediate)

Estructura de VGG16 sin BatchNorm

En esta variante, cada bloque convolucional aplica una convolución 3x3 seguida de una activación ReLU, sin normalización intermedia. La reducción espacial se logra mediante capas MaxPool2d que reducen la resolución a la mitad. El modelo completo realiza un submuestreo de 16x.

VGG16Backbone(
  (feature_pipe): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
  )
)

Estrcutura de VGG16 con BatchNorm

La variante vgg16_bn inserta una capa BatchNorm2d entre cada convolución y su activación ReLU correspondiente. Esto estabiliza el entrenamiento al normalizar las activaciones de cada canal, reduciendo la sensibilidad a la inicialización de pesos y acelerando la convergencia.

VGG16BNBackbone(
  (feature_pipe): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
    (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): ReLU(inplace=True)
    (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    ...
  )
)

Ambas variantes realizan un submuestreo total de 16x respecto a la resolución de entrada, lo que significa que un tensor de imagen de 256x256 se reduce a 16x16 en la salida del backbone. El número de filtros aumenta progresivamente: 64, 128, 256 y 512 canales en cada uno de los cuatro bloques principales.