Translate legacy source documentation to English without changing executable code

0d8aa44 verified about 14 hours ago

8.07 kB

import torch
import torch.nn as nn

class LTFCore(nn.Module):
    """
    Core module.
    Input:  x (B, N, T, D)
    Output: logits (B, num_classes), token_scores (B, N, T-1, 1), jpeg_score (B, 1)
    """
    def __init__(
        self,
        embed_dim=1024,
        num_heads=8,
        num_classes=2,
        num_layers=1,
        dropout=0.1,
        tile_heads=2,
        tile_layers=2,
        use_global_image=True,
        use_patch_transformer=True,
        linear_classifier=False,
    ):
        super().__init__()
        self.use_global_image = use_global_image
        self.use_patch_transformer = use_patch_transformer
        # Learnable CLS for tile-level aggregation
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # Tile-level encoder (across tiles)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            dropout=dropout,
            nhead=tile_heads,
            batch_first=True  
        )
        self.tile_transformer = nn.TransformerEncoder(encoder_layer, num_layers=tile_layers)
        
        if use_patch_transformer:
            token_layer = nn.TransformerEncoderLayer(
                d_model=embed_dim,
                dropout=dropout,
                nhead=num_heads,
                batch_first=True 
            )
            self.patch_transformer = nn.TransformerEncoder(token_layer, num_layers=num_layers)
        else:
            self.patch_transformer = None
        # Image-level classifier: with global (main tile + aggregated tiles) or only aggregated tiles (ablation)
        classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim
        if linear_classifier:
            self.classifier = nn.Linear(classifier_in_dim, num_classes)
        else:
            self.classifier = nn.Sequential(
                nn.Linear(classifier_in_dim, 512),
                nn.ReLU(),
                nn.Dropout(p=dropout),
                nn.Linear(512, num_classes)
            )
        # Per-token head (for non-CLS tokens)
        self.patch_classifier = nn.Sequential(
            nn.Linear(embed_dim, 512),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(512, 1)
        )
        # Scalar head (e.g., JPEG-related score)
        self.jpeg_estimate = nn.Sequential(
            nn.Linear(embed_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
        )

        nn.init.trunc_normal_(self.cls_token, std=0.02)
    
    def forward(self, x):
        """
        x: (B, N, T, D)
        returns: (B, num_classes), (B, N, T-1, 1), (B, 1)
        """
        B_2, N_2, T_2, D_2 = x.shape

        # Token-level encoding per tile (optional)
        x = x.reshape(B_2*N_2, T_2, D_2)       # [B*N, T, D]
        if self.patch_transformer is not None:
            x = self.patch_transformer(x)
        x = x.reshape(B_2, N_2, T_2, D_2)
        
        # Split CLS vs non-CLS tokens
        new_cls = x[:, :, 0, :]
        new_x = x[:, :, 1:, :]
        # Per-token head on non-CLS tokens
        x_flat = new_x.reshape(-1, D_2)       # [B*N*T, D]
        x_result = self.patch_classifier(x_flat)       # [B*N*(T-1), 1]
        x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1)   # [B, N, T-1, 1]

        main_feat = new_cls[:, 0]              # (B, D)
        tile_feats = new_cls[:, 1:]            # (B, N-1, D)
        # print(main_feat.shape, tile_feats.shape)

        # Tile-level aggregation with a learnable CLS
        cls_token = self.cls_token.expand(B_2, 1, -1)              # (B, 1, D)
        feature_tile = torch.cat([cls_token, tile_feats], dim=1)       # (B, N, D)

        feature_encoded = self.tile_transformer(feature_tile)                     # (B, N, D)
        cls_out = feature_encoded[:, 0]                                # (B, D)

        # Image-level classification: with or without global (main tile) feature (ablation v1_1)
        if self.use_global_image:
            concat_feat = torch.cat([main_feat, cls_out], dim=1)     # (B, 2D)
        else:
            concat_feat = cls_out                                      # (B, D)
        return self.classifier(concat_feat), x_result, torch.sigmoid(self.jpeg_estimate(cls_out))  # (B, num_classes) (B)

class LTFNet(nn.Module):
    """
    Wrapper with a frozen visual backbone (e.g., CLIP visual) + LTFCore.
    Input:  images x (B, N, 3, H, W)
    Output: same as LTFCore
    """
    def __init__(
        self,
        clip_visual,
        num_heads=8,
        num_classes=2,
        num_layers=1,
        dropout=0.2,
        tile_heads=2,
        tile_layers=2,
        use_global_image=True,
        use_patch_transformer=True,
        linear_classifier=False,
    ):
        super().__init__()
        self.visual = clip_visual 
        self.dtype = self.visual.conv1.weight.dtype

        for p in self.visual.parameters():
            p.requires_grad = False
        self.visual.eval()

        self.input_resolution = self.visual.input_resolution
        # Token feature dim = transformer hidden size (ViT-L -> 1024). Use conv1 out_channels;
        # do NOT use output_dim which is the projection-head dim (often 768) and would cause
        # "expecting 768 but got 1024" in patch_transformer.
        self.embed_dim = self.visual.conv1.weight.shape[0]

        self.core = LTFCore(
            embed_dim=self.embed_dim,
            num_heads=num_heads,
            num_classes=num_classes,
            num_layers=num_layers,
            dropout=dropout,
            tile_heads=tile_heads,
            tile_layers=tile_layers,
            use_global_image=use_global_image,
            use_patch_transformer=use_patch_transformer,
            linear_classifier=linear_classifier,
        )

    def get_tokens(self, x) -> torch.Tensor:
        """
        Extract token sequence (with CLS) from the visual backbone.
        Assumes CLIP ViT structure: conv1 -> patch embed -> class_embedding + pos_embed -> ln_pre -> transformer.
        x: (B, 3, H, W) -> tokens: (B, T, D)
        """
        x = x.type(self.dtype)
        x = self.visual.conv1(x)  # [B, C, H, W]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # [B, C, HW]
        x = x.permute(0, 2, 1)  # [B, HW, C]
        cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device)
        x = torch.cat([cls, x], dim=1)  # [B, 1+HW, C]
        x = x + self.visual.positional_embedding.to(x.dtype)
        x = self.visual.ln_pre(x)
        x = x.permute(1, 0, 2)  # [sequence_len, B, C]
        x = self.visual.transformer(x)
        x = x.permute(1, 0, 2)  # [B, sequence_len, C]
        return x 

    def encode_image(self, image) -> torch.Tensor:
        # image: (B, 3, H, W) -> feature: (B, D)
        return self.visual(image.type(self.dtype))

    def forward(self, x: torch.Tensor, chunk_size=256):
        # Tensor[B, 5, 3, 224, 224]
        """
        x: Tensor of shape (B, N, 3, H, W)
        chunk_size: number of tiles to pass through the visual backbone per batch. Same tiles are
                    computed and concatenated; only speed/memory change, no effect on outputs.
                    Eval can set model.eval_chunk_size (e.g. 512) in config for faster run.
        """
        chunk_size = getattr(self, "eval_chunk_size", chunk_size)
        B, N, C, H, W = x.shape
        imgs = x.reshape(B * N, C, H, W)
        feats_chunks = []
        with torch.no_grad():
            for start in range(0, imgs.size(0), chunk_size):
                end = min(start + chunk_size, imgs.size(0))
                feats_chunk = self.get_tokens(imgs[start:end])   # (chunk, T, D)
                feats_chunks.append(feats_chunk)
        feats_token = torch.cat(feats_chunks, dim=0)
        feats = feats_token.reshape(B, N, feats_token.shape[-2], -1)  # (B, N, T, D)
        feats = feats.float()
        return self.core(feats)          # logits: (B, num_classes)  (B)

	import torch
	import torch.nn as nn

	class LTFCore(nn.Module):
	"""
	Core module.
	Input: x (B, N, T, D)
	Output: logits (B, num_classes), token_scores (B, N, T-1, 1), jpeg_score (B, 1)
	"""
	def __init__(
	self,
	embed_dim=1024,
	num_heads=8,
	num_classes=2,
	num_layers=1,
	dropout=0.1,
	tile_heads=2,
	tile_layers=2,
	use_global_image=True,
	use_patch_transformer=True,
	linear_classifier=False,
	):
	super().__init__()
	self.use_global_image = use_global_image
	self.use_patch_transformer = use_patch_transformer
	# Learnable CLS for tile-level aggregation
	self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
	# Tile-level encoder (across tiles)
	encoder_layer = nn.TransformerEncoderLayer(
	d_model=embed_dim,
	dropout=dropout,
	nhead=tile_heads,
	batch_first=True
	)
	self.tile_transformer = nn.TransformerEncoder(encoder_layer, num_layers=tile_layers)

	if use_patch_transformer:
	token_layer = nn.TransformerEncoderLayer(
	d_model=embed_dim,
	dropout=dropout,
	nhead=num_heads,
	batch_first=True
	)
	self.patch_transformer = nn.TransformerEncoder(token_layer, num_layers=num_layers)
	else:
	self.patch_transformer = None
	# Image-level classifier: with global (main tile + aggregated tiles) or only aggregated tiles (ablation)
	classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim
	if linear_classifier:
	self.classifier = nn.Linear(classifier_in_dim, num_classes)
	else:
	self.classifier = nn.Sequential(
	nn.Linear(classifier_in_dim, 512),
	nn.ReLU(),
	nn.Dropout(p=dropout),
	nn.Linear(512, num_classes)
	)
	# Per-token head (for non-CLS tokens)
	self.patch_classifier = nn.Sequential(
	nn.Linear(embed_dim, 512),
	nn.ReLU(),
	nn.Dropout(p=dropout),
	nn.Linear(512, 1)
	)
	# Scalar head (e.g., JPEG-related score)
	self.jpeg_estimate = nn.Sequential(
	nn.Linear(embed_dim, 256),
	nn.ReLU(),
	nn.Linear(256, 128),
	nn.ReLU(),
	nn.Linear(128, 128),
	nn.ReLU(),
	nn.Linear(128, 1),
	)

	nn.init.trunc_normal_(self.cls_token, std=0.02)

	def forward(self, x):
	"""
	x: (B, N, T, D)
	returns: (B, num_classes), (B, N, T-1, 1), (B, 1)
	"""
	B_2, N_2, T_2, D_2 = x.shape

	# Token-level encoding per tile (optional)
	x = x.reshape(B_2N_2, T_2, D_2) # [BN, T, D]
	if self.patch_transformer is not None:
	x = self.patch_transformer(x)
	x = x.reshape(B_2, N_2, T_2, D_2)

	# Split CLS vs non-CLS tokens
	new_cls = x[:, :, 0, :]
	new_x = x[:, :, 1:, :]
	# Per-token head on non-CLS tokens
	x_flat = new_x.reshape(-1, D_2) # [BNT, D]
	x_result = self.patch_classifier(x_flat) # [BN(T-1), 1]
	x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1) # [B, N, T-1, 1]

	main_feat = new_cls[:, 0] # (B, D)
	tile_feats = new_cls[:, 1:] # (B, N-1, D)
	# print(main_feat.shape, tile_feats.shape)

	# Tile-level aggregation with a learnable CLS
	cls_token = self.cls_token.expand(B_2, 1, -1) # (B, 1, D)
	feature_tile = torch.cat([cls_token, tile_feats], dim=1) # (B, N, D)

	feature_encoded = self.tile_transformer(feature_tile) # (B, N, D)
	cls_out = feature_encoded[:, 0] # (B, D)

	# Image-level classification: with or without global (main tile) feature (ablation v1_1)
	if self.use_global_image:
	concat_feat = torch.cat([main_feat, cls_out], dim=1) # (B, 2D)
	else:
	concat_feat = cls_out # (B, D)
	return self.classifier(concat_feat), x_result, torch.sigmoid(self.jpeg_estimate(cls_out)) # (B, num_classes) (B)

	class LTFNet(nn.Module):
	"""
	Wrapper with a frozen visual backbone (e.g., CLIP visual) + LTFCore.
	Input: images x (B, N, 3, H, W)
	Output: same as LTFCore
	"""
	def __init__(
	self,
	clip_visual,
	num_heads=8,
	num_classes=2,
	num_layers=1,
	dropout=0.2,
	tile_heads=2,
	tile_layers=2,
	use_global_image=True,
	use_patch_transformer=True,
	linear_classifier=False,
	):
	super().__init__()
	self.visual = clip_visual
	self.dtype = self.visual.conv1.weight.dtype

	for p in self.visual.parameters():
	p.requires_grad = False
	self.visual.eval()

	self.input_resolution = self.visual.input_resolution
	# Token feature dim = transformer hidden size (ViT-L -> 1024). Use conv1 out_channels;
	# do NOT use output_dim which is the projection-head dim (often 768) and would cause
	# "expecting 768 but got 1024" in patch_transformer.
	self.embed_dim = self.visual.conv1.weight.shape[0]

	self.core = LTFCore(
	embed_dim=self.embed_dim,
	num_heads=num_heads,
	num_classes=num_classes,
	num_layers=num_layers,
	dropout=dropout,
	tile_heads=tile_heads,
	tile_layers=tile_layers,
	use_global_image=use_global_image,
	use_patch_transformer=use_patch_transformer,
	linear_classifier=linear_classifier,
	)

	def get_tokens(self, x) -> torch.Tensor:
	"""
	Extract token sequence (with CLS) from the visual backbone.
	Assumes CLIP ViT structure: conv1 -> patch embed -> class_embedding + pos_embed -> ln_pre -> transformer.
	x: (B, 3, H, W) -> tokens: (B, T, D)
	"""
	x = x.type(self.dtype)
	x = self.visual.conv1(x) # [B, C, H, W]
	x = x.reshape(x.shape[0], x.shape[1], -1) # [B, C, HW]
	x = x.permute(0, 2, 1) # [B, HW, C]
	cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device)
	x = torch.cat([cls, x], dim=1) # [B, 1+HW, C]
	x = x + self.visual.positional_embedding.to(x.dtype)
	x = self.visual.ln_pre(x)
	x = x.permute(1, 0, 2) # [sequence_len, B, C]
	x = self.visual.transformer(x)
	x = x.permute(1, 0, 2) # [B, sequence_len, C]
	return x

	def encode_image(self, image) -> torch.Tensor:
	# image: (B, 3, H, W) -> feature: (B, D)
	return self.visual(image.type(self.dtype))

	def forward(self, x: torch.Tensor, chunk_size=256):
	# Tensor[B, 5, 3, 224, 224]
	"""
	x: Tensor of shape (B, N, 3, H, W)
	chunk_size: number of tiles to pass through the visual backbone per batch. Same tiles are
	computed and concatenated; only speed/memory change, no effect on outputs.
	Eval can set model.eval_chunk_size (e.g. 512) in config for faster run.
	"""
	chunk_size = getattr(self, "eval_chunk_size", chunk_size)
	B, N, C, H, W = x.shape
	imgs = x.reshape(B * N, C, H, W)
	feats_chunks = []
	with torch.no_grad():
	for start in range(0, imgs.size(0), chunk_size):
	end = min(start + chunk_size, imgs.size(0))
	feats_chunk = self.get_tokens(imgs[start:end]) # (chunk, T, D)
	feats_chunks.append(feats_chunk)
	feats_token = torch.cat(feats_chunks, dim=0)
	feats = feats_token.reshape(B, N, feats_token.shape[-2], -1) # (B, N, T, D)
	feats = feats.float()
	return self.core(feats) # logits: (B, num_classes) (B)