| """In-repo CNN vision encoder for multimodal training.""" | |
| from typing import Sequence | |
| import torch | |
| import torch.nn as nn | |
| class CNNEncoder(nn.Module): | |
| """A compact convolutional encoder for single-image multimodal inputs.""" | |
| def __init__( | |
| self, | |
| image_size: int, | |
| output_dim: int, | |
| channels: Sequence[int] = (32, 64, 128), | |
| kernel_size: int = 3, | |
| ): | |
| """Initialize the CNN encoder.""" | |
| super().__init__() | |
| if image_size < 8: | |
| raise ValueError("image_size must be at least 8") | |
| if output_dim < 1: | |
| raise ValueError("output_dim must be positive") | |
| if not channels: | |
| raise ValueError("channels must contain at least one stage") | |
| layers: list[nn.Module] = [] | |
| in_channels = 3 | |
| stride = 2 | |
| padding = kernel_size // 2 | |
| for channel_dim in channels: | |
| layers.extend([ | |
| nn.Conv2d(in_channels, channel_dim, kernel_size=kernel_size, stride=stride, padding=padding), | |
| nn.BatchNorm2d(channel_dim), | |
| nn.GELU(), | |
| ]) | |
| in_channels = channel_dim | |
| self.backbone = nn.Sequential(*layers) | |
| self.pool = nn.AdaptiveAvgPool2d((1, 1)) | |
| self.projection = nn.Linear(in_channels, output_dim) | |
| def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: | |
| """Encode pixel values into a single feature vector per image.""" | |
| features = self.backbone(pixel_values) | |
| pooled = self.pool(features).flatten(1) | |
| return self.projection(pooled) | |