File size: 3,399 Bytes
d7228c8 | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """
PyTorch Vision Encoder module for LawSLM Multimodal Vision-Language Architecture.
Converts 3-channel image inputs [B, 3, H, W] into token projection sequence [B, N_patches, d_model].
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional, Dict, Any
from slm.utils.logger import get_logger
logger = get_logger("slm.vision.encoder")
class VisionPatchEmbedding(nn.Module):
"""
Splits image tensor into non-overlapping patches and projects them to target d_model dimension.
"""
def __init__(
self,
image_size: int = 224,
patch_size: int = 16,
in_channels: int = 3,
embed_dim: int = 128
) -> None:
super().__init__()
self.image_size = image_size
self.patch_size = patch_size
self.n_patches = (image_size // patch_size) ** 2
self.embed_dim = embed_dim
# Patch Convolution Projection
self.proj = nn.Conv2d(
in_channels,
embed_dim,
kernel_size=patch_size,
stride=patch_size
)
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embedding = nn.Parameter(torch.randn(1, self.n_patches + 1, embed_dim) * 0.02)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: Input image tensor of shape [B, 3, H, W]
Returns:
Sequence embedding tensor of shape [B, N_patches + 1, d_model]
"""
B, C, H, W = x.shape
if H != self.image_size or W != self.image_size:
x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)
x = self.proj(x) # [B, embed_dim, grid, grid]
x = x.flatten(2).transpose(1, 2) # [B, N_patches, embed_dim]
cls_tokens = self.cls_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, x), dim=1) # [B, N_patches + 1, embed_dim]
x = x + self.pos_embedding
return x
class VisionEncoder(nn.Module):
"""
Lightweight Vision Encoder Transformer Block for image feature extraction.
"""
def __init__(
self,
image_size: int = 224,
patch_size: int = 16,
in_channels: int = 3,
d_model: int = 128,
n_layers: int = 2,
n_heads: int = 4
) -> None:
super().__init__()
self.patch_embed = VisionPatchEmbedding(
image_size=image_size,
patch_size=patch_size,
in_channels=in_channels,
embed_dim=d_model
)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_model * 4,
dropout=0.1,
activation="gelu",
batch_first=True
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
self.ln_post = nn.LayerNorm(d_model)
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
"""
Processes normalized image pixels into sequence embeddings.
Args:
pixel_values: Image tensor [B, 3, H, W]
Returns:
Image token sequence tensor [B, N_patches + 1, d_model]
"""
tokens = self.patch_embed(pixel_values)
out = self.transformer(tokens)
return self.ln_post(out)
|