projector.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class PointwiseConv(nn.Module):
    """点卷积用于时序维度的投影"""

    def __init__(self, in_channels, out_channels):
        super(PointwiseConv, self).__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=1)
        # 初始化为接近恒等映射
        nn.init.xavier_uniform_(self.conv.weight)
        nn.init.zeros_(self.conv.bias)

    def forward(self, x):
        # 输入x形状: [batch_size, seq_len, dim]
        x = self.conv(x)
        return x


class ResidualCNNBlock(nn.Module):
    """带残差连接的CNN块"""

    def __init__(self, dim, kernel_size=3):
        super(ResidualCNNBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2),
            nn.BatchNorm1d(dim),
            nn.ReLU(),
            nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2),
            nn.BatchNorm1d(dim)
        )
        self.activation = nn.ReLU()

    def forward(self, x):
        # 输入x形状: [batch_size, seq_len, dim]
        identity = x

        # 转置为卷积所需的形状
        x = x.transpose(1, 2)  # [batch_size, dim, seq_len]
        x = self.conv_block(x)
        x = x.transpose(1, 2)  # [batch_size, seq_len, dim]

        # 残差连接
        x = x + identity
        x = self.activation(x)

        return x


class CNNFeatureExtractor(nn.Module):
    """CNN特征提取器，带残差连接"""

    def __init__(self, dim, num_layers=3):
        super(CNNFeatureExtractor, self).__init__()
        self.cnn_blocks = nn.ModuleList([
            ResidualCNNBlock(dim)
            for _ in range(num_layers)
        ])

    def forward(self, x):
        # 输入x形状: [batch_size, seq_len, dim]
        for block in self.cnn_blocks:
            x = block(x)
        return x


class TimeSeriesProjectionModule(nn.Module):
    """针对时序维度的投影模块"""

    def __init__(self, in_seq_len, out_seq_len, dim):
        super(TimeSeriesProjectionModule, self).__init__()
        self.point_conv = PointwiseConv(in_seq_len, out_seq_len)
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        # x形状: [batch_size, in_seq_len, dim]
        x = self.point_conv(x)  # [batch_size, out_seq_len, dim]
        x = self.norm(x)
        return x


class FeatureProjector(nn.Module):
    def __init__(self,
                 input_seq_len=320, output_seq_len=208, input_dim=256, output_dim=768,
                 hidden_dim=512, num_cnn_layers=3, num_transformer_layers=8):
        super(FeatureProjector, self).__init__()

        # 1. 初始特征维度投影到hidden_dim
        self.initial_projection = nn.Linear(input_dim, hidden_dim)

        # 2. CNN特征提取器（使用hidden_dim）
        self.cnn_extractor = CNNFeatureExtractor(hidden_dim, num_layers=num_cnn_layers)

        # 3. 时序维度投影（维持hidden_dim）
        self.time_projection = TimeSeriesProjectionModule(input_seq_len, output_seq_len, hidden_dim)

        # 4. 位置编码
        self.pos_encoder = nn.Parameter(torch.zeros(1, output_seq_len, hidden_dim))
        nn.init.trunc_normal_(self.pos_encoder, std=0.02)

        # 5. Transformer编码器（使用PyTorch内置组件）
        # TransformerEncoderLayer内部已包含残差连接
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=8,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            activation="gelu",
            batch_first=True,
            norm_first=True
        )
        # TransformerEncoder的层之间没有额外的残差连接，但每层内部有
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_transformer_layers
        )

        # 6. 最终特征维度投影到out_dim
        self.final_projection = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        # 输入x形状: [batch_size, in_seq_len, in_dim]

        # 1. 初始特征维度投影
        x = self.initial_projection(x)  # [batch_size, in_seq_len, hidden_dim]

        # 2. CNN特征提取（带残差）
        x = self.cnn_extractor(x)  # [batch_size, in_seq_len, hidden_dim]

        # 3. 时序维度投影
        x = self.time_projection(x)  # [batch_size, out_seq_len, hidden_dim]

        # 4. 添加位置编码
        x = x + self.pos_encoder

        # 5. Transformer编码（已包含内部残差）
        x = self.transformer_encoder(x)  # [batch_size, out_seq_len, hidden_dim]

        # 6. 最终特征维度投影
        x = self.final_projection(x)  # [batch_size, out_seq_len, out_dim]

        return x


class ProjectorLoss(nn.Module):
    def __init__(self, time_bins=26, freq_bins=8, feature_dim=768):
        super().__init__()
        self.time_bins = time_bins
        self.freq_bins = freq_bins
        self.feature_dim = feature_dim
        self.mse = nn.MSELoss()
        self.cos = nn.CosineSimilarity(dim=-1)

    def forward(self, student_feat, teacher_feat):
        """
        Args:
            student_feat: [B, 208, 768]
            teacher_feat: [B, 208, 768]
        """
        B = student_feat.shape[0]

        # 1. 基础重建损失
        mse_loss = self.mse(student_feat, teacher_feat)

        # 2. 重塑特征为时间-频带形式
        # [B, 208, 768] -> [B, 26, 8, 768]
        s_feat = student_feat.reshape(B, self.time_bins, self.freq_bins, -1)
        t_feat = teacher_feat.reshape(B, self.time_bins, self.freq_bins, -1)

        # 3. 时间维度的连续性损失
        time_continuity_loss = self.temporal_continuity_loss(s_feat, t_feat)

        # 4. 频带内部相关性损失
        freq_correlation_loss = self.frequency_correlation_loss(s_feat, t_feat)

        # 5. 时频联合分布损失
        # joint_distribution_loss = self.time_frequency_joint_loss(s_feat, t_feat)

        # 6. 特征动态范围损失
        # dynamic_range_loss = self.dynamic_range_loss(s_feat, t_feat)

        # 7. 局部结构保持损失
        local_structure_loss = self.local_structure_loss(s_feat, t_feat)

        total_loss = (mse_loss +
                      0.2 * time_continuity_loss +
                      0.2 * freq_correlation_loss +
                      # 0.1 * joint_distribution_loss +
                      # 0.1 * dynamic_range_loss +
                      0.2 * local_structure_loss)

        loss_dict = {
            'mse_loss': mse_loss,
            'time_continuity_loss': time_continuity_loss,
            'freq_correlation_loss': freq_correlation_loss,
            # 'joint_distribution_loss': joint_distribution_loss,
            # 'dynamic_range_loss': dynamic_range_loss,
            'local_structure_loss': local_structure_loss
        }
        return total_loss, loss_dict

    def temporal_continuity_loss(self, s_feat, t_feat):
        """时间维度连续性损失：确保时间维度上的平滑变化"""
        # 计算时间维度上的梯度
        s_temp_grad = s_feat[:, 1:] - s_feat[:, :-1]  # [B, 25, 8, 768]
        t_temp_grad = t_feat[:, 1:] - t_feat[:, :-1]  # [B, 25, 8, 768]

        # 梯度相似性损失
        grad_loss = self.mse(s_temp_grad, t_temp_grad)

        # 时序相关性损失
        s_time_corr = self.cos(s_feat[:, 1:].reshape(-1, self.feature_dim),
                               s_feat[:, :-1].reshape(-1, self.feature_dim)).mean()
        t_time_corr = self.cos(t_feat[:, 1:].reshape(-1, self.feature_dim),
                               t_feat[:, :-1].reshape(-1, self.feature_dim)).mean()
        corr_loss = torch.abs(s_time_corr - t_time_corr)

        return grad_loss + 0.5 * corr_loss

    def frequency_correlation_loss(self, s_feat, t_feat):
        """频带内部相关性损失：确保频带特征的正确分布"""

        # 计算不同频带间的相关性矩阵
        def compute_freq_corr(x):
            x = x.mean(dim=1)  # [B, 8, 768]
            x = F.normalize(x, dim=-1)
            return torch.bmm(x, x.transpose(1, 2))  # [B, 8, 8]

        s_freq_corr = compute_freq_corr(s_feat)
        t_freq_corr = compute_freq_corr(t_feat)

        return self.mse(s_freq_corr, t_freq_corr)

    def time_frequency_joint_loss(self, s_feat, t_feat):
        """时频联合分布损失：保持时频特征的整体分布"""

        # 计算时频特征的联合统计信息
        def compute_joint_stats(x):
            # 计算均值和标准差
            mean = x.mean(dim=[1, 2])  # [B, 768]
            std = x.std(dim=[1, 2])  # [B, 768]
            # 计算时频协方差
            x_centered = x - mean.unsqueeze(1).unsqueeze(2)
            cov = torch.einsum('btfd,btfg->bdg', x_centered, x_centered) / (self.time_bins * self.freq_bins)
            return mean, std, cov

        s_mean, s_std, s_cov = compute_joint_stats(s_feat)
        t_mean, t_std, t_cov = compute_joint_stats(t_feat)

        return (self.mse(s_mean, t_mean) +
                self.mse(s_std, t_std) +
                0.1 * self.mse(s_cov, t_cov))

    def dynamic_range_loss(self, s_feat, t_feat):
        """特征动态范围损失：确保特征值分布的动态范围一致"""

        def compute_dynamic_range(x):
            percentiles = torch.quantile(x, torch.tensor([0.05, 0.95], device=x.device), dim=1)
            percentiles = percentiles[1] - percentiles[0]
            percentiles = torch.quantile(percentiles, torch.tensor([0.05, 0.95], device=x.device), dim=1)
            percentiles = percentiles[1] - percentiles[0]
            return percentiles

        s_range = compute_dynamic_range(s_feat)
        t_range = compute_dynamic_range(t_feat)

        return self.mse(s_range, t_range)

    def local_structure_loss(self, s_feat, t_feat):
        """局部结构保持损失：确保局部时频模式的一致性"""

        # 使用2D卷积来提取局部特征
        def extract_local_patterns(x):
            x = x.permute(0, 3, 1, 2)
            unfold = nn.Unfold(kernel_size=(3, 3), stride=1, padding=1)
            patterns = unfold(x)
            return patterns

        s_patterns = extract_local_patterns(s_feat)
        t_patterns = extract_local_patterns(t_feat)

        return self.mse(s_patterns, t_patterns)


if __name__ == '__main__':
    x = torch.randn(3, 320, 256).cuda()
    target = torch.randn(3, 208, 768).cuda()
    model = FeatureProjector()
    # model = TransformerProjector(d_t=320, num_patches=208)
    model = model.cuda()
    # output = model(x, teacher_forcing=False, target=target)
    output = model(x)
    print(output.shape)  # 输出形状 (16, 208, 768)
    #
    # fn = AdapterLoss()
    # loss, log_dict = fn(output, target)
    # print(loss)
    # print(log_dict)