diffusers_lite/utils/network.py

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from diffusers.models.normalization import FP32LayerNorm

class QueryAttention(nn.Module):
    """
    Query-based attention pooling module using PyTorch's built-in MultiheadAttention.
    Uses learnable query vectors to attend to sequence features.
    """
    def __init__(self, feature_dim, num_queries=1, num_heads=8, dropout=0.1, layer_norm=False, return_type=None, product_text=False, text_dim=768):
        super(QueryAttention, self).__init__()
        self.feature_dim = feature_dim
        self.num_queries = num_queries
        self.num_heads = num_heads
        self.layer_norm = layer_norm
        self.return_type = return_type
        self.product_text = product_text

        # Use PyTorch's built-in MultiheadAttention
        self.multihead_attn = nn.MultiheadAttention(
            embed_dim=feature_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True  # Use batch_first=True for easier handling
        )
        # Learnable query vectors
        self.queries = nn.Parameter(torch.randn(num_queries, feature_dim))
        # Initialize query parameters
        nn.init.xavier_uniform_(self.queries)

        if self.layer_norm:
            self.norm = FP32LayerNorm(feature_dim, eps=1e-6, elementwise_affine=False)
        
        if self.product_text:
            self.text_proj = nn.Linear(text_dim, feature_dim)
            nn.init.xavier_uniform_(self.text_proj.weight)
            if self.text_proj.bias is not None:
                nn.init.zeros_(self.text_proj.bias)

        
    def forward(self, x, e = None, text = None):
        """
        Args:
            x: Input tensor of shape [batch_size, seq_len, feature_dim] or [batch_size, feature_dim]
        Returns:
            Pooled features of shape [batch_size, feature_dim]
        """

        if self.layer_norm:
            x = self.norm(x)

        batch_size = x.shape[0]
        original_shape = x.shape
        
        # Handle different input shapes
        if len(x.shape) == 2:  # [batch_size, feature_dim]
            # Add sequence dimension
            x = x.unsqueeze(1)  # [batch_size, 1, feature_dim]
            seq_len = 1
        elif len(x.shape) == 3:  # [batch_size, seq_len, feature_dim]
            seq_len = x.shape[1]
        elif len(x.shape) == 4:  # [sp_size, batch_size, seq_len, feature_dim]
            # Handle sequence parallel case
            sp_size, batch_size, seq_len, feature_dim = x.shape
            x = x.view(sp_size * batch_size, seq_len, feature_dim)
            batch_size = sp_size * batch_size
        else:
            raise ValueError(f"Unsupported input shape: {x.shape}")

        # Expand queries to batch size
        queries = self.queries.unsqueeze(0).expand(batch_size, -1, -1)  # [batch_size, num_queries, feature_dim]
        if e is not None:
            queries = queries + e.unsqueeze(0).expand(batch_size, -1, -1)
        # Use PyTorch's MultiheadAttention
        # query: [batch_size, num_queries, feature_dim]
        # key, value: [batch_size, seq_len, feature_dim]
        attended, attention_weights = self.multihead_attn(
            query=queries,
            key=x,
            value=x,
            need_weights=False  # We don't need attention weights for pooling
        )
        
        # attended: [batch_size, num_queries, feature_dim]
        
        # If multiple queries, average them
        if self.num_queries > 1:
            output = attended.mean(dim=1)  # [batch_size, feature_dim]
        else:
            output = attended.squeeze(1)  # [batch_size, feature_dim]
        
        # Handle sequence parallel case
        if len(original_shape) == 4:
            output = output.view(sp_size, batch_size // sp_size, -1)
            output = output.mean(dim=0)  # Average across SP devices
        
        if self.layer_norm:
            output = self.norm(output)

        if self.return_type == 'query':
            output = output + queries

        if self.product_text and text is not None:
            output_product_text = torch.mul(self.text_proj(text), output)
            return output_product_text
        else:
            return output

class MLP(nn.Module):
    def __init__(self, input_dim):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_dim, 1024)  # First hidden layer
        self.fc2 = nn.Linear(1024, 512)        # Second hidden layer
        self.fc3 = nn.Linear(512, 1)           # Output layer (binary classification)
        
        # 初始化权重，避免梯度消失
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # 使用Xavier初始化
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)  # 注意：这里不应用sigmoid，因为forward_siamese会处理
        return x

class MultiHead(nn.Module):
    def __init__(self, input_dim, num_heads = 3):
        super().__init__()
        self.num_heads = num_heads
        self.mlps = torch.nn.ModuleList(
            [MLP(input_dim) for _ in range(num_heads)]
        )

    def forward_mlp(self, head_idx, x):
        return torch.sigmoid(self.mlps[head_idx](x))

    def forward(self, x):
        out = [self.forward_mlp(h, x) for h in range(self.num_heads)]
        return torch.stack(out)

def forward_mlp(model, input):
    return torch.sigmoid(model(input))

def forward_siamese(model, input1, input2):
    # Pass both inputs through the same model (weight sharing)
    reward1 = model(input1)
    reward2 = model(input2)
    # Compute the difference between the two embeddings
    diff = reward1 - reward2

    # Use this difference for binary prediction (preference/ranking)
    return torch.sigmoid(diff)

def train_model(model, device, model_mode, X_train, y_train, X_test, y_test, epochs=3, lr=0.001, batch_size = 512, verbose=False, ealry_stopping_patience=3):
    model = model.to(device)  # Move model to GPU
    criterion = nn.BCELoss()  # Binary cross-entropy loss with logits for numerical stability
    optimizer = optim.Adam(model.parameters(), lr=lr)
    batch_size = min(batch_size, X_train.shape[0])
    val_losses = []
    for epoch in range(epochs):
        for n_batch in range(0, X_train.shape[0], batch_size):

            # randomly selecte batch
            batch_idx = torch.randperm(X_train.shape[0])[:batch_size]
            X_batch = X_train[batch_idx]
            y_batch = y_train[batch_idx]

            model.train()
            optimizer.zero_grad()

            # Forward pass
            if model_mode == 'clf':
                outputs = forward_mlp(model, X_batch)
            elif model_mode == 'siamese':
                outputs = forward_siamese(model, X_batch[:, 0], X_batch[:, 1])
            loss = criterion(outputs, y_batch)

            # Backward pass and optimization
            loss.backward()
            optimizer.step()

        # Evaluate on validation set
        model.eval()
        with torch.no_grad():
            if model_mode == 'clf':
                val_outputs = forward_mlp(model, X_test)
            elif model_mode == 'siamese':
                val_outputs = forward_siamese(model, X_test[:, 0], X_test[:, 1])
        # early stopping?
        val_loss = criterion(val_outputs, y_test)
        val_losses.append(val_loss.cpu().detach().item())
        if len(val_losses) > ealry_stopping_patience:
            if all(val_losses[-1] > x for x in val_losses[-(ealry_stopping_patience+1):-1]):
                if verbose:
                    print(f"Early stopping at epoch {epoch+1}")
                break
        if verbose:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.cpu().detach().item()}, Val Loss: {val_loss.cpu().detach().item()}")
        # accuracy
        val_outputs = val_outputs.cpu().detach().numpy()
        val_pred = (val_outputs > 0.5).astype(int)
        accuracy = (val_pred == y_test.cpu().detach().numpy()).mean()
        if verbose:
            print(f"Accuracy: {accuracy}")

def save_model(model, path):
    torch.save(model.state_dict(), path)