File size: 909 Bytes
3275441 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """Xavante - image_encoder.py - Encoder de imagem (Teorema 16.1/16.2)."""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
logger = logging.getLogger(__name__)
class ImageEncoder(nn.Module):
"""CNN leve -> features em d_model."""
def __init__(self, d_model: int = 512, in_channels: int = 3):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, 32, 4, 2, 1),
nn.GELU(),
nn.Conv2d(32, 64, 4, 2, 1),
nn.GELU(),
nn.Conv2d(64, 128, 4, 2, 1),
nn.GELU(),
nn.AdaptiveAvgPool2d((1, 1)),
)
self.proj = nn.Linear(128, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, C, H, W]
feat = self.conv(x).flatten(1) # [B, 128]
return self.proj(feat) # [B, d_model]
__all__ = ["ImageEncoder"]
|