Rayan-Q1 / model.py

Update model.py

b689ca4 verified almost 2 years ago

3.44 kB

#model.py
import torch
import torch.nn as nn
import open_clip

class CONFIG:
    # Model parameters
    clip_model_name = "ViT-B-32"
    pretrained_dataset = "laion400m_e31"
    embedding_dim = 512
    fusion_type = 'sum'  # 'sum', 'concat', or 'attention'
    device = "cuda" if torch.cuda.is_available() else "cpu"
    image_size = (224, 224)
    batch_size = 16
    device = "cuda" if torch.cuda.is_available() else "cpu"


class FusionModel(nn.Module):
    def __init__(self, clip_model, fusion_type='sum'):
        super(FusionModel, self).__init__()
        self.clip_model = clip_model
        self.fusion_type = fusion_type

        # Depending on the fusion type, define the combined embedding size
        if fusion_type == 'sum':
            self.combined_dim = CONFIG.embedding_dim
        elif fusion_type == 'concat':
            self.combined_dim = CONFIG.embedding_dim * 2
        elif fusion_type == 'attention':
            self.combined_dim = CONFIG.embedding_dim
            # Define the attention layer
            self.attention = nn.MultiheadAttention(embed_dim=CONFIG.embedding_dim, num_heads=8)
        else:
            raise ValueError("Invalid fusion type")

        # MLP to project the combined embedding to the final embedding space
        self.mlp = nn.Sequential(
            nn.Linear(self.combined_dim, CONFIG.embedding_dim),
            nn.ReLU(),
            nn.Dropout(p=0.1),
            nn.Linear(CONFIG.embedding_dim, CONFIG.embedding_dim)
        )

    def forward(self, image, text):
        # Encode image and text using CLIP encoders
        image_features = self.clip_model.encode_image(image)  # Shape: [batch_size, embedding_dim]
        text_tokens = open_clip.tokenize(text).to(CONFIG.device)  # Tokenize text
        text_features = self.clip_model.encode_text(text_tokens)  # Shape: [batch_size, embedding_dim]

        # Normalize the features
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        # Fusion strategies
        if self.fusion_type == 'sum':
            combined_features = image_features + text_features
        elif self.fusion_type == 'concat':
            combined_features = torch.cat([image_features, text_features], dim=-1)
        elif self.fusion_type == 'attention':
            combined_features = self.attention_fusion(image_features, text_features)
        else:
            raise ValueError("Invalid fusion type")

        # Pass through MLP
        output = self.mlp(combined_features)
        return output

    def attention_fusion(self, image_features, text_features):
        # Reshape features to (sequence_length, batch_size, embedding_dim)
        # Since we have two features per sample, sequence_length = 2
        features = torch.stack([image_features, text_features], dim=0)  # Shape: [2, batch_size, embedding_dim]

        # Multi-head attention expects inputs in shape (sequence_length, batch_size, embedding_dim)
        # Query, Key, Value are all the features in this case
        attn_output, _ = self.attention(features, features, features)    # Shape: [2, batch_size, embedding_dim]

        # Aggregate the attention outputs (e.g., take the mean over the sequence length)
        combined_features = attn_output.mean(dim=0)                      # Shape: [batch_size, embedding_dim]

        return combined_features

	#model.py
	import torch
	import torch.nn as nn
	import open_clip

	class CONFIG:
	# Model parameters
	clip_model_name = "ViT-B-32"
	pretrained_dataset = "laion400m_e31"
	embedding_dim = 512
	fusion_type = 'sum' # 'sum', 'concat', or 'attention'
	device = "cuda" if torch.cuda.is_available() else "cpu"
	image_size = (224, 224)
	batch_size = 16
	device = "cuda" if torch.cuda.is_available() else "cpu"


	class FusionModel(nn.Module):
	def __init__(self, clip_model, fusion_type='sum'):
	super(FusionModel, self).__init__()
	self.clip_model = clip_model
	self.fusion_type = fusion_type

	# Depending on the fusion type, define the combined embedding size
	if fusion_type == 'sum':
	self.combined_dim = CONFIG.embedding_dim
	elif fusion_type == 'concat':
	self.combined_dim = CONFIG.embedding_dim * 2
	elif fusion_type == 'attention':
	self.combined_dim = CONFIG.embedding_dim
	# Define the attention layer
	self.attention = nn.MultiheadAttention(embed_dim=CONFIG.embedding_dim, num_heads=8)
	else:
	raise ValueError("Invalid fusion type")

	# MLP to project the combined embedding to the final embedding space
	self.mlp = nn.Sequential(
	nn.Linear(self.combined_dim, CONFIG.embedding_dim),
	nn.ReLU(),
	nn.Dropout(p=0.1),
	nn.Linear(CONFIG.embedding_dim, CONFIG.embedding_dim)
	)

	def forward(self, image, text):
	# Encode image and text using CLIP encoders
	image_features = self.clip_model.encode_image(image) # Shape: [batch_size, embedding_dim]
	text_tokens = open_clip.tokenize(text).to(CONFIG.device) # Tokenize text
	text_features = self.clip_model.encode_text(text_tokens) # Shape: [batch_size, embedding_dim]

	# Normalize the features
	image_features = image_features / image_features.norm(dim=-1, keepdim=True)
	text_features = text_features / text_features.norm(dim=-1, keepdim=True)

	# Fusion strategies
	if self.fusion_type == 'sum':
	combined_features = image_features + text_features
	elif self.fusion_type == 'concat':
	combined_features = torch.cat([image_features, text_features], dim=-1)
	elif self.fusion_type == 'attention':
	combined_features = self.attention_fusion(image_features, text_features)
	else:
	raise ValueError("Invalid fusion type")

	# Pass through MLP
	output = self.mlp(combined_features)
	return output

	def attention_fusion(self, image_features, text_features):
	# Reshape features to (sequence_length, batch_size, embedding_dim)
	# Since we have two features per sample, sequence_length = 2
	features = torch.stack([image_features, text_features], dim=0) # Shape: [2, batch_size, embedding_dim]

	# Multi-head attention expects inputs in shape (sequence_length, batch_size, embedding_dim)
	# Query, Key, Value are all the features in this case
	attn_output, _ = self.attention(features, features, features) # Shape: [2, batch_size, embedding_dim]

	# Aggregate the attention outputs (e.g., take the mean over the sequence length)
	combined_features = attn_output.mean(dim=0) # Shape: [batch_size, embedding_dim]

	return combined_features