File size: 1,569 Bytes
4c775f9 ac8a72f 4c775f9 bc0e267 4c775f9 6742b9c 4c775f9 | 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 | import torch.nn as nn
from transformers import ViTModel
class SEMViTAutoencoder(nn.Module):
def __init__(self):
super().__init__()
# Load pre-trained ViT-Base as encoder
self.encoder = ViTModel.from_pretrained(
'google/vit-base-patch16-224', # ViT expects 224p, 3 channels (RGB)
add_pooling_layer=False
)
# Custom convolutional decoder to reconstruct the image
# Decoder: Mapping ViT features (768) back to image space
# 512/16 = 32 patches. Latent grid is 32x32.
self.decoder = nn.Sequential(
nn.ConvTranspose2d(768, 256, kernel_size=4, stride=2, padding=1), # 64x64
nn.BatchNorm2d(256),
nn.ReLU(),
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # 128x128
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 256x256
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1), # 512x512
nn.Sigmoid() # Output pixels in [0, 1]
)
def forward(self, x):
# ViT requires position interpolation for resolutions != 224
outputs = self.encoder(x, interpolate_pos_encoding=True)
# Sequence to Grid: [Batch, 1025, 768] -> [Batch, 768, 32, 32]
latent = outputs.last_hidden_state[:, 1:, :].transpose(1, 2).reshape(-1, 768, 32, 32)
return self.decoder(latent) |