CNN Encoder Spatial Feature Map
Overview
For image-to-sequence tasks (captioning, OCR, molecular translation), the CNN must output spatial feature maps rather than a single vector. Replace the global pooling and FC head with nn.Identity(), then permute/reshape the output to (batch, H*W, C). Each spatial position becomes an "input token" the decoder can attend to.
Quick Start
pythonimport timm import torch.nn as nn class SpatialEncoder(nn.Module): def __init__(self, model_name="resnet34", pretrained=True): super().__init__() self.cnn = timm.create_model(model_name, pretrained=pretrained) self.n_features = self.cnn.fc.in_features self.cnn.global_pool = nn.Identity() self.cnn.fc = nn.Identity() def forward(self, x): features = self.cnn(x) # (B, C, H, W) features = features.permute(0, 2, 3, 1) # (B, H, W, C) B, H, W, C = features.shape features = features.view(B, H * W, C) # (B, num_pixels, C) return features
Workflow
- Load any pretrained CNN via timm/torchvision
- Replace
global_poolandfcwithnn.Identity() - Forward pass yields (B, C, H, W) spatial tensor
- Permute and reshape to (B, num_pixels, C) for attention input
- Feed spatial features into attention-based decoder
Key Decisions
- Model choice: ResNet-34/50 gives 7x7=49 spatial positions for 224x224 input; EfficientNet varies
- Image size: Larger input = more spatial positions = richer attention but more memory
- Adaptive pooling: Optionally add
nn.AdaptiveAvgPool2d((H, W))before flatten for fixed spatial dims - Fine-tuning: Freeze early layers, train later layers with the decoder
