필사 모드: 自监督学习（Self-Supervised Learning）完全指南：SimCLR、MAE、DINO、CLIP

import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image

class RotationSSL(nn.Module):
    def __init__(self, backbone, num_classes=4):
        super().__init__()
        self.backbone = backbone
        self.classifier = nn.Linear(backbone.output_dim, num_classes)

    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)

def create_rotation_dataset(images):
    """将图像沿 4 个方向旋转，生成 (image, label) 对"""
    rotated_images = []
    labels = []

    for img in images:
        for k in range(4):  # 0, 90, 180, 270 度
            # 将 PIL Image 旋转 k*90 度
            rotated = T.functional.rotate(img, k * 90)
            rotated_images.append(rotated)
            labels.append(k)

    return rotated_images, labels

# 训练循环
def train_rotation_ssl(model, dataloader, optimizer, epochs=10):
    criterion = nn.CrossEntropyLoss()
    model.train()

    for epoch in range(epochs):
        total_loss = 0
        for images, labels in dataloader:
            images, labels = images.cuda(), labels.cuda()

            optimizer.zero_grad()
            logits = model(images)
            loss = criterion(logits, labels)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

import itertools
import numpy as np

class JigsawSSL(nn.Module):
    def __init__(self, backbone, num_permutations=100):
        super().__init__()
        self.backbone = backbone
        # 独立处理每个 patch
        self.patch_encoder = nn.Sequential(
            backbone,
            nn.Linear(backbone.output_dim, 512)
        )
        # 处理完全部 9 个 patch 后进行排列分类
        self.classifier = nn.Sequential(
            nn.Linear(512 * 9, 4096),
            nn.ReLU(),
            nn.Linear(4096, num_permutations)
        )

    def forward(self, patches):
        # patches: (batch, 9, C, H, W)
        batch_size = patches.size(0)
        patch_features = []

        for i in range(9):
            feat = self.patch_encoder(patches[:, i])  # (batch, 512)
            patch_features.append(feat)

        # 拼接所有 patch 特征
        combined = torch.cat(patch_features, dim=1)  # (batch, 512*9)
        return self.classifier(combined)

def create_jigsaw_dataset(image, grid_size=3):
    """将图像分割为 grid_size x grid_size 的 patch 并打乱顺序"""
    h, w = image.shape[-2:]
    patch_h, patch_w = h // grid_size, w // grid_size

    patches = []
    for i in range(grid_size):
        for j in range(grid_size):
            patch = image[...,
                         i*patch_h:(i+1)*patch_h,
                         j*patch_w:(j+1)*patch_w]
            patches.append(patch)

    # 从预定义的排列集合中选择
    permutation_idx = np.random.randint(0, len(PREDEFINED_PERMUTATIONS))
    perm = PREDEFINED_PERMUTATIONS[permutation_idx]
    shuffled = [patches[p] for p in perm]

    return torch.stack(shuffled), permutation_idx

class ColorizationSSL(nn.Module):
    def __init__(self):
        super().__init__()
        # 输入 L 通道（亮度）
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.ReLU(),
        )
        # 输出 ab 通道（颜色）
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 2, 3, padding=1),  # ab 通道
            nn.Tanh()
        )

    def forward(self, l_channel):
        features = self.encoder(l_channel)
        ab_channels = self.decoder(features)
        return ab_channels

# 在 LAB 色彩空间中处理
from skimage.color import rgb2lab, lab2rgb

def prepare_colorization_data(rgb_image):
    lab = rgb2lab(rgb_image)
    l_channel = lab[:, :, 0:1]   # 亮度通道
    ab_channels = lab[:, :, 1:]  # 颜色通道
    return l_channel, ab_channels

class VideoSSL(nn.Module):
    def __init__(self, frame_encoder, temporal_model):
        super().__init__()
        self.frame_encoder = frame_encoder
        self.temporal_model = temporal_model  # LSTM 或 Transformer
        self.decoder = nn.ConvTranspose2d(...)

    def forward(self, frames):
        # frames: (batch, T, C, H, W)
        T = frames.size(1)

        # 对每一帧编码
        frame_features = []
        for t in range(T - 1):
            feat = self.frame_encoder(frames[:, t])
            frame_features.append(feat)

        # 时间建模
        features = torch.stack(frame_features, dim=1)
        next_feat = self.temporal_model(features)

        # 解码下一帧
        next_frame_pred = self.decoder(next_feat[:, -1])
        return next_frame_pred

图像 x
├── 增强视图 v1 → z1（嵌入）
└── 增强视图 v2 → z2（嵌入）

目标：z1 与 z2 应接近，z1 与其他图像的 z3 应远离

import torch
import torch.nn.functional as F

def info_nce_loss(z1, z2, temperature=0.5):
    """
    InfoNCE / NT-Xent Loss 实现
    z1, z2: (batch_size, embedding_dim) - 同一图像的两个视图
    """
    batch_size = z1.size(0)

    # L2 归一化
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)

    # 拼接两个视图，生成 2N x D 矩阵
    # [z1_1, z1_2, ..., z1_N, z2_1, z2_2, ..., z2_N]
    z = torch.cat([z1, z2], dim=0)  # (2N, D)

    # 计算所有配对的相似度
    similarity = torch.mm(z, z.t()) / temperature  # (2N, 2N)

    # 排除与自身的相似度（对角线设为 -inf）
    mask = torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
    similarity.masked_fill_(mask, float('-inf'))

    # 正样本对：第 i 个与第 i+N 个，第 i+N 个与第 i 个
    labels = torch.cat([
        torch.arange(batch_size, 2 * batch_size),
        torch.arange(batch_size)
    ]).to(z.device)

    loss = F.cross_entropy(similarity, labels)
    return loss

# 使用示例
z1 = torch.randn(32, 128)  # 批大小 32，嵌入维度 128
z2 = torch.randn(32, 128)
loss = info_nce_loss(z1, z2, temperature=0.5)
print(f"InfoNCE Loss: {loss.item():.4f}")

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as T

class SimCLR(nn.Module):
    def __init__(self, backbone='resnet50', projection_dim=128, temperature=0.5):
        super().__init__()
        self.temperature = temperature

        # 主干编码器
        if backbone == 'resnet50':
            resnet = models.resnet50(pretrained=False)
            self.encoder = nn.Sequential(*list(resnet.children())[:-1])  # 去掉 FC 层
            self.feature_dim = 2048
        elif backbone == 'resnet18':
            resnet = models.resnet18(pretrained=False)
            self.encoder = nn.Sequential(*list(resnet.children())[:-1])
            self.feature_dim = 512

        # Projection Head（重要：微调时需去掉）
        self.projection_head = nn.Sequential(
            nn.Linear(self.feature_dim, self.feature_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.feature_dim, projection_dim)
        )

    def encode(self, x):
        h = self.encoder(x)
        h = h.squeeze(-1).squeeze(-1)  # (batch, feature_dim)
        return h

    def forward(self, x):
        h = self.encode(x)
        z = self.projection_head(h)
        return z

    def contrastive_loss(self, z1, z2):
        return info_nce_loss(z1, z2, self.temperature)


class SimCLRDataAugmentation:
    """SimCLR 的强数据增强"""
    def __init__(self, image_size=224, s=1.0):
        color_jitter = T.ColorJitter(
            brightness=0.8*s,
            contrast=0.8*s,
            saturation=0.8*s,
            hue=0.2*s
        )
        self.transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply([color_jitter], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.GaussianBlur(kernel_size=int(0.1 * image_size)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)


# SimCLR 训练
def train_simclr(model, dataloader, optimizer, epochs=200):
    model.train()
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=epochs
    )

    for epoch in range(epochs):
        total_loss = 0
        for (x1, x2), _ in dataloader:
            x1, x2 = x1.cuda(), x2.cuda()

            z1 = model(x1)
            z2 = model(x2)

            loss = model.contrastive_loss(z1, z2)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step()
        avg_loss = total_loss / len(dataloader)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")


# 完整流水线
model = SimCLR(backbone='resnet50', projection_dim=128).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-6)
augmentation = SimCLRDataAugmentation(image_size=224)

# 微调时去掉 projection head
def get_finetuning_model(pretrained_simclr, num_classes):
    # 仅使用编码器
    encoder = pretrained_simclr.encoder
    classifier = nn.Linear(pretrained_simclr.feature_dim, num_classes)
    return nn.Sequential(encoder, nn.Flatten(), classifier)

class MoCo(nn.Module):
    def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07):
        super().__init__()
        self.K = K  # 队列大小
        self.m = m  # 动量系数
        self.T = T  # 温度

        # query 编码器与 key 编码器
        self.encoder_q = base_encoder(num_classes=dim)
        self.encoder_k = base_encoder(num_classes=dim)

        # key 编码器不计算梯度，仅做动量更新
        for param_q, param_k in zip(
            self.encoder_q.parameters(),
            self.encoder_k.parameters()
        ):
            param_k.data.copy_(param_q.data)
            param_k.requires_grad = False

        # 负样本队列
        self.register_buffer("queue", torch.randn(dim, K))
        self.queue = F.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """以动量方式更新 key 编码器"""
        for param_q, param_k in zip(
            self.encoder_q.parameters(),
            self.encoder_k.parameters()
        ):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        """向队列中加入新 key，移除旧 key"""
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)

        self.queue[:, ptr:ptr + batch_size] = keys.T
        ptr = (ptr + batch_size) % self.K
        self.queue_ptr[0] = ptr

    def forward(self, im_q, im_k):
        # query 嵌入
        q = self.encoder_q(im_q)
        q = F.normalize(q, dim=1)

        # key 嵌入（不计算梯度）
        with torch.no_grad():
            self._momentum_update_key_encoder()
            k = self.encoder_k(im_k)
            k = F.normalize(k, dim=1)

        # 正样本 logit：(batch, 1)
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        # 负样本 logit：(batch, K)
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])

        logits = torch.cat([l_pos, l_neg], dim=1) / self.T
        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()

        loss = F.cross_entropy(logits, labels)

        self._dequeue_and_enqueue(k)
        return loss

class BYOL(nn.Module):
    def __init__(self, backbone, projection_dim=256, prediction_dim=128, tau=0.996):
        super().__init__()
        self.tau = tau  # 动量系数

        # 在线网络
        self.online_encoder = backbone
        self.online_projector = nn.Sequential(
            nn.Linear(2048, 4096), nn.BatchNorm1d(4096), nn.ReLU(),
            nn.Linear(4096, projection_dim)
        )
        self.online_predictor = nn.Sequential(
            nn.Linear(projection_dim, 4096), nn.BatchNorm1d(4096), nn.ReLU(),
            nn.Linear(4096, prediction_dim)
        )

        # 目标网络（不计算梯度）
        self.target_encoder = copy.deepcopy(backbone)
        self.target_projector = copy.deepcopy(self.online_projector)

        for param in self.target_encoder.parameters():
            param.requires_grad = False
        for param in self.target_projector.parameters():
            param.requires_grad = False

    @torch.no_grad()
    def update_target(self):
        """EMA 更新"""
        for online, target in zip(
            self.online_encoder.parameters(),
            self.target_encoder.parameters()
        ):
            target.data = self.tau * target.data + (1 - self.tau) * online.data

        for online, target in zip(
            self.online_projector.parameters(),
            self.target_projector.parameters()
        ):
            target.data = self.tau * target.data + (1 - self.tau) * online.data

    def forward(self, x1, x2):
        # 在线网络前向
        online_feat1 = self.online_encoder(x1).squeeze()
        online_proj1 = self.online_projector(online_feat1)
        online_pred1 = self.online_predictor(online_proj1)

        online_feat2 = self.online_encoder(x2).squeeze()
        online_proj2 = self.online_projector(online_feat2)
        online_pred2 = self.online_predictor(online_proj2)

        # 目标网络（stop gradient）
        with torch.no_grad():
            target_feat1 = self.target_encoder(x1).squeeze()
            target_proj1 = self.target_projector(target_feat1)

            target_feat2 = self.target_encoder(x2).squeeze()
            target_proj2 = self.target_projector(target_feat2)

        # BYOL Loss：最大化余弦相似度
        loss1 = byol_loss(online_pred1, target_proj2.detach())
        loss2 = byol_loss(online_pred2, target_proj1.detach())

        return (loss1 + loss2).mean()

def byol_loss(p, z):
    """Negative cosine similarity"""
    p = F.normalize(p, dim=-1)
    z = F.normalize(z, dim=-1)
    return -(p * z).sum(dim=-1)

class SimSiam(nn.Module):
    def __init__(self, backbone, dim=2048, pred_dim=512):
        super().__init__()
        self.encoder = backbone

        # Projector
        self.projector = nn.Sequential(
            nn.Linear(dim, dim, bias=False),
            nn.BatchNorm1d(dim),
            nn.ReLU(inplace=True),
            nn.Linear(dim, dim, bias=False),
            nn.BatchNorm1d(dim),
            nn.ReLU(inplace=True),
            nn.Linear(dim, dim, bias=False),
            nn.BatchNorm1d(dim, affine=False)  # 最后一层 BN：affine=False
        )

        # Predictor
        self.predictor = nn.Sequential(
            nn.Linear(dim, pred_dim, bias=False),
            nn.BatchNorm1d(pred_dim),
            nn.ReLU(inplace=True),
            nn.Linear(pred_dim, dim)
        )

    def forward(self, x1, x2):
        z1 = self.projector(self.encoder(x1).squeeze())
        z2 = self.projector(self.encoder(x2).squeeze())

        p1 = self.predictor(z1)
        p2 = self.predictor(z2)

        # Stop gradient —— 非对称结构的核心
        loss = simsiam_loss(p1, z2.detach()) / 2 + \
               simsiam_loss(p2, z1.detach()) / 2
        return loss

def simsiam_loss(p, z):
    """Negative cosine similarity"""
    z = z.detach()  # Stop gradient
    p = F.normalize(p, dim=1)
    z = F.normalize(z, dim=1)
    return -(p * z).sum(dim=1).mean()

输入图像（196 个 patch）
    ↓ 75% 掩码
49 个可见 patch → 编码器（ViT-Large）→ 编码
                 ↓
196 个 patch（含掩码 token）→ 解码器（浅层 ViT）→ 像素复原

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat

class PatchEmbedding(nn.Module):
    """将图像转换为 patch 嵌入"""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2

        self.projection = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # x: (batch, C, H, W)
        x = self.projection(x)  # (batch, embed_dim, H/p, W/p)
        x = x.flatten(2)        # (batch, embed_dim, num_patches)
        x = x.transpose(1, 2)  # (batch, num_patches, embed_dim)
        return x


class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., dropout=0.):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        mlp_dim = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # Self-attention
        norm_x = self.norm1(x)
        attn_out, _ = self.attn(norm_x, norm_x, norm_x)
        x = x + attn_out
        # MLP
        x = x + self.mlp(self.norm2(x))
        return x


class MAE(nn.Module):
    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_channels=3,
        # 编码器配置
        encoder_dim=768,
        encoder_depth=12,
        encoder_heads=12,
        # 解码器配置（比编码器小得多）
        decoder_dim=512,
        decoder_depth=8,
        decoder_heads=16,
        # 掩码比例
        mask_ratio=0.75,
    ):
        super().__init__()
        self.mask_ratio = mask_ratio
        self.patch_size = patch_size
        num_patches = (img_size // patch_size) ** 2

        # patch 嵌入
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, encoder_dim)

        # [CLS] token 与位置嵌入（编码器）
        self.cls_token = nn.Parameter(torch.zeros(1, 1, encoder_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, encoder_dim),
            requires_grad=False
        )

        # 编码器（ViT）
        self.encoder_blocks = nn.ModuleList([
            TransformerBlock(encoder_dim, encoder_heads)
            for _ in range(encoder_depth)
        ])
        self.encoder_norm = nn.LayerNorm(encoder_dim)

        # 编码器-解码器连接
        self.decoder_embed = nn.Linear(encoder_dim, decoder_dim, bias=True)

        # 掩码 token
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))

        # 解码器位置嵌入
        self.decoder_pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, decoder_dim),
            requires_grad=False
        )

        # 解码器（浅层 ViT）
        self.decoder_blocks = nn.ModuleList([
            TransformerBlock(decoder_dim, decoder_heads)
            for _ in range(decoder_depth)
        ])
        self.decoder_norm = nn.LayerNorm(decoder_dim)

        # 像素预测头
        self.decoder_pred = nn.Linear(decoder_dim, patch_size**2 * in_channels)

        self._init_weights()

    def _init_weights(self):
        """初始化 sinusoidal 位置嵌入"""
        # 初始化位置嵌入（sinusoidal）
        pos_embed = self._get_2d_sincos_pos_embed(
            self.pos_embed.shape[-1],
            int(self.patch_embed.num_patches**0.5)
        )
        self.pos_embed.data.copy_(
            torch.from_numpy(pos_embed).float().unsqueeze(0)
        )
        # 初始化其余部分
        nn.init.normal_(self.cls_token, std=.02)
        nn.init.normal_(self.mask_token, std=.02)

    def _get_2d_sincos_pos_embed(self, embed_dim, grid_size):
        """2D sinusoidal 位置嵌入"""
        import numpy as np
        grid_h = np.arange(grid_size, dtype=np.float32)
        grid_w = np.arange(grid_size, dtype=np.float32)
        grid = np.meshgrid(grid_w, grid_h)
        grid = np.stack(grid, axis=0)
        grid = grid.reshape([2, 1, grid_size, grid_size])

        emb_h = self._get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
        emb_w = self._get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
        emb = np.concatenate([emb_h, emb_w], axis=1)
        # 为 CLS token 添加零嵌入
        emb = np.concatenate([np.zeros([1, embed_dim]), emb], axis=0)
        return emb

    def _get_1d_sincos_pos_embed_from_grid(self, embed_dim, pos):
        import numpy as np
        omega = np.arange(embed_dim // 2, dtype=np.float32)
        omega /= embed_dim / 2.
        omega = 1. / 10000**omega
        pos = pos.reshape(-1)
        out = np.einsum('m,d->md', pos, omega)
        emb_sin = np.sin(out)
        emb_cos = np.cos(out)
        emb = np.concatenate([emb_sin, emb_cos], axis=1)
        return emb

    def random_masking(self, x, mask_ratio):
        """随机掩码：仅保留部分 patch"""
        N, L, D = x.shape  # batch, length, dim
        len_keep = int(L * (1 - mask_ratio))

        # 用随机噪声生成打乱索引
        noise = torch.rand(N, L, device=x.device)
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        # 仅选择可见 patch
        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(
            x, dim=1,
            index=ids_keep.unsqueeze(-1).repeat(1, 1, D)
        )

        # 生成掩码（1：被掩盖，0：可见）
        mask = torch.ones(N, L, device=x.device)
        mask[:, :len_keep] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    def forward_encoder(self, x, mask_ratio):
        """编码器：仅处理可见 patch"""
        # patch 嵌入
        x = self.patch_embed(x)

        # 添加位置嵌入（不含 CLS）
        x = x + self.pos_embed[:, 1:, :]

        # 随机掩码
        x, mask, ids_restore = self.random_masking(x, mask_ratio)

        # 添加 CLS token
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # 编码器 Transformer
        for block in self.encoder_blocks:
            x = block(x)
        x = self.encoder_norm(x)

        return x, mask, ids_restore

    def forward_decoder(self, x, ids_restore):
        """解码器：包含掩码 token，复原整体"""
        # 编码器维度 → 解码器维度
        x = self.decoder_embed(x)

        # 添加掩码 token
        mask_tokens = self.mask_token.repeat(
            x.shape[0],
            ids_restore.shape[1] + 1 - x.shape[1],
            1
        )
        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # 不含 CLS
        x_ = torch.gather(
            x_, dim=1,
            index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])
        )
        x = torch.cat([x[:, :1, :], x_], dim=1)  # 重新加入 CLS

        # 解码器位置嵌入
        x = x + self.decoder_pos_embed

        # 解码器 Transformer
        for block in self.decoder_blocks:
            x = block(x)
        x = self.decoder_norm(x)

        # 像素预测
        x = self.decoder_pred(x)
        x = x[:, 1:, :]  # 去掉 CLS token
        return x

    def forward(self, imgs, mask_ratio=0.75):
        # 编码（含掩码）
        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
        # 解码（复原像素）
        pred = self.forward_decoder(latent, ids_restore)
        # 计算损失
        loss = self.forward_loss(imgs, pred, mask)
        return loss, pred, mask

    def forward_loss(self, imgs, pred, mask):
        """针对被掩码 patch 的 MSE Loss"""
        # 目标：patch 化后的图像
        target = self.patchify(imgs)

        # 归一化（patch normalization）
        mean = target.mean(dim=-1, keepdim=True)
        var = target.var(dim=-1, keepdim=True)
        target = (target - mean) / (var + 1.e-6).sqrt()

        # 仅对被掩码的 patch 计算损失
        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # 每个 patch 的平均值
        loss = (loss * mask).sum() / mask.sum()  # 被掩码 patch 的平均值
        return loss

    def patchify(self, imgs):
        """将图像转换为 patch 序列"""
        p = self.patch_size
        h = w = imgs.shape[2] // p
        x = imgs.reshape(imgs.shape[0], 3, h, p, w, p)
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(imgs.shape[0], h * w, p**2 * 3)
        return x


# 创建并训练 MAE 模型
mae_model = MAE(
    img_size=224,
    patch_size=16,
    encoder_dim=768,
    encoder_depth=12,
    encoder_heads=12,
    decoder_dim=512,
    decoder_depth=8,
    decoder_heads=16,
    mask_ratio=0.75
).cuda()

optimizer = torch.optim.AdamW(
    mae_model.parameters(),
    lr=1.5e-4,
    betas=(0.9, 0.95),
    weight_decay=0.05
)

# 训练步骤
def train_step(model, imgs, optimizer):
    model.train()
    optimizer.zero_grad()

    loss, pred, mask = model(imgs, mask_ratio=0.75)
    loss.backward()
    optimizer.step()

    return loss.item()

Student Network（学生）        Teacher Network（教师）
      ↑                              ↑
 Local Views（多个）          Global Views（2 个）

class DINOHead(nn.Module):
    """DINO Projection Head"""
    def __init__(self, in_dim, out_dim=65536, bottleneck_dim=256):
        super().__init__()
        layers = [nn.Linear(in_dim, 2048), nn.GELU()]
        layers += [nn.Linear(2048, 2048), nn.GELU()]
        layers += [nn.Linear(2048, bottleneck_dim)]
        self.mlp = nn.Sequential(*layers)

        # 经过 weight normalization 的最后一层
        self.last_layer = nn.utils.weight_norm(
            nn.Linear(bottleneck_dim, out_dim, bias=False)
        )
        self.last_layer.weight_g.data.fill_(1)
        self.last_layer.weight_g.requires_grad = False

    def forward(self, x):
        x = self.mlp(x)
        x = F.normalize(x, dim=-1)
        x = self.last_layer(x)
        return x


class DINO(nn.Module):
    def __init__(self, backbone, out_dim=65536, teacher_temp=0.04, student_temp=0.1):
        super().__init__()
        self.student_temp = student_temp
        self.teacher_temp = teacher_temp

        # Student 网络
        self.student = backbone
        self.student_head = DINOHead(self.student.embed_dim, out_dim)

        # Teacher 网络（EMA）
        self.teacher = copy.deepcopy(backbone)
        self.teacher_head = copy.deepcopy(self.student_head)

        for p in self.teacher.parameters():
            p.requires_grad = False
        for p in self.teacher_head.parameters():
            p.requires_grad = False

        # 用于 centering 的 center 向量
        self.register_buffer("center", torch.zeros(1, out_dim))

    @torch.no_grad()
    def update_teacher(self, momentum):
        """EMA 更新"""
        for student_ps, teacher_ps in zip(
            self.student.parameters(), self.teacher.parameters()
        ):
            teacher_ps.data.mul_(momentum).add_(
                (1 - momentum) * student_ps.data
            )
        for student_ps, teacher_ps in zip(
            self.student_head.parameters(), self.teacher_head.parameters()
        ):
            teacher_ps.data.mul_(momentum).add_(
                (1 - momentum) * student_ps.data
            )

    @torch.no_grad()
    def update_center(self, teacher_output):
        """Centering 更新（防止 collapse）"""
        batch_center = teacher_output.mean(dim=0, keepdim=True)
        self.center = self.center * 0.9 + batch_center * 0.1

    def dino_loss(self, student_output, teacher_output):
        """DINO Cross-entropy Loss"""
        student_out = student_output / self.student_temp
        student_out = student_out.chunk(2)  # global crops

        # Teacher: centering + sharpening
        teacher_out = F.softmax(
            (teacher_output - self.center) / self.teacher_temp, dim=-1
        )
        teacher_out = teacher_out.detach().chunk(2)

        total_loss = 0
        n_loss_terms = 0

        for iq, q in enumerate(teacher_out):
            for v in range(len(student_out)):
                if v == iq:
                    continue  # 排除相同视图
                loss = torch.sum(
                    -q * F.log_softmax(student_out[v], dim=-1), dim=-1
                )
                total_loss += loss.mean()
                n_loss_terms += 1

        return total_loss / n_loss_terms

    def forward(self, global_views, local_views, momentum=0.996):
        # Student：处理所有视图
        all_views = global_views + local_views
        student_output = torch.cat([
            self.student_head(self.student(view))
            for view in all_views
        ])

        # Teacher：仅处理全局视图
        with torch.no_grad():
            teacher_output = torch.cat([
                self.teacher_head(self.teacher(view))
                for view in global_views
            ])

        loss = self.dino_loss(student_output, teacher_output)

        # 更新
        self.update_center(teacher_output)
        self.update_teacher(momentum)

        return loss


# DINO 数据增强（多裁剪）
class DINOAugmentation:
    def __init__(self, global_crops_scale=(0.4, 1.), local_crops_scale=(0.05, 0.4),
                 n_local_crops=8, image_size=224, local_size=96):
        self.n_local_crops = n_local_crops

        # 全局裁剪（2 个，大视图）
        self.global_transform = T.Compose([
            T.RandomResizedCrop(image_size, scale=global_crops_scale, interpolation=T.InterpolationMode.BICUBIC),
            T.RandomHorizontalFlip(),
            T.RandomApply([T.ColorJitter(0.4, 0.4, 0.2, 0.1)], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

        # 局部裁剪（8 个，小视图）
        self.local_transform = T.Compose([
            T.RandomResizedCrop(local_size, scale=local_crops_scale, interpolation=T.InterpolationMode.BICUBIC),
            T.RandomHorizontalFlip(),
            T.RandomApply([T.ColorJitter(0.4, 0.4, 0.2, 0.1)], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def __call__(self, image):
        global_views = [self.global_transform(image) for _ in range(2)]
        local_views = [self.local_transform(image) for _ in range(self.n_local_crops)]
        return global_views, local_views

图像编码器（ViT / ResNet）
        ↓
  图像嵌入

文本编码器（Transformer）
        ↓
  文本嵌入

目标：匹配的 (图像, 文本) 对应相似
     不匹配的对应不同

def clip_loss(image_embeddings, text_embeddings, temperature):
    """
    CLIP 对称对比学习损失
    image_embeddings: (N, D)
    text_embeddings: (N, D)
    """
    # L2 归一化
    image_embeddings = F.normalize(image_embeddings, dim=1)
    text_embeddings = F.normalize(text_embeddings, dim=1)

    # 计算相似度矩阵：(N, N)
    logits = torch.matmul(image_embeddings, text_embeddings.T) / temperature

    # 标签：对角线元素（第 i 个图像 = 第 i 个文本）
    labels = torch.arange(len(logits)).to(logits.device)

    # 图像→文本、文本→图像双向损失
    loss_i2t = F.cross_entropy(logits, labels)
    loss_t2i = F.cross_entropy(logits.T, labels)

    return (loss_i2t + loss_t2i) / 2

from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import requests
import torch
import torch.nn.functional as F

# 加载预训练 CLIP 模型
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# ---- Zero-shot 图像分类 ----
def zero_shot_classify(image_path, candidate_labels):
    """
    在没有标签样本的情况下对图像分类
    """
    image = Image.open(image_path)

    # 生成文本提示
    texts = [f"a photo of a {label}" for label in candidate_labels]

    # 编码
    inputs = processor(
        text=texts,
        images=image,
        return_tensors="pt",
        padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)

    # 计算相似度
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)

    results = {}
    for label, prob in zip(candidate_labels, probs[0]):
        results[label] = prob.item()

    return results

# 使用示例
labels = ["cat", "dog", "car", "airplane", "bird"]
results = zero_shot_classify("example.jpg", labels)
for label, prob in sorted(results.items(), key=lambda x: -x[1]):
    print(f"{label}: {prob:.4f}")


# ---- 图文相似度检索 ----
class CLIPRetrieval:
    def __init__(self, model_name="openai/clip-vit-base-patch32"):
        self.model = CLIPModel.from_pretrained(model_name)
        self.processor = CLIPProcessor.from_pretrained(model_name)
        self.model.eval()

        self.image_embeddings = None
        self.image_paths = []

    def index_images(self, image_paths):
        """为图像数据库建立索引"""
        self.image_paths = image_paths
        embeddings = []

        for path in image_paths:
            image = Image.open(path)
            inputs = self.processor(images=image, return_tensors="pt")
            with torch.no_grad():
                embedding = self.model.get_image_features(**inputs)
            embeddings.append(F.normalize(embedding, dim=1))

        self.image_embeddings = torch.cat(embeddings, dim=0)
        print(f"已完成 {len(image_paths)} 张图像的索引")

    def search(self, query_text, top_k=5):
        """通过文本查询检索图像"""
        inputs = self.processor(text=[query_text], return_tensors="pt", padding=True)
        with torch.no_grad():
            text_embedding = self.model.get_text_features(**inputs)
        text_embedding = F.normalize(text_embedding, dim=1)

        # 余弦相似度
        similarities = torch.matmul(text_embedding, self.image_embeddings.T).squeeze()
        top_indices = similarities.argsort(descending=True)[:top_k]

        results = [
            (self.image_paths[i], similarities[i].item())
            for i in top_indices
        ]
        return results

# 多模态嵌入可视化
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np

def visualize_clip_embeddings(images, texts):
    """CLIP 嵌入的 t-SNE 可视化"""
    image_inputs = processor(images=images, return_tensors="pt")
    text_inputs = processor(text=texts, return_tensors="pt", padding=True)

    with torch.no_grad():
        image_embs = F.normalize(model.get_image_features(**image_inputs), dim=1)
        text_embs = F.normalize(model.get_text_features(**text_inputs), dim=1)

    # 拼接图像与文本嵌入
    all_embs = torch.cat([image_embs, text_embs], dim=0).numpy()

    # t-SNE 降维
    tsne = TSNE(n_components=2, random_state=42)
    reduced = tsne.fit_transform(all_embs)

    n = len(images)
    plt.figure(figsize=(12, 8))
    plt.scatter(reduced[:n, 0], reduced[:n, 1], c='blue', label='Images', alpha=0.7)
    plt.scatter(reduced[n:, 0], reduced[n:, 1], c='red', label='Texts', alpha=0.7)

    for i, text in enumerate(texts):
        plt.annotate(text, reduced[n+i], fontsize=8)

    plt.legend()
    plt.title('CLIP Embeddings (t-SNE)')
    plt.savefig('clip_tsne.png', dpi=150, bbox_inches='tight')
    plt.show()

第 1 步：用 dVAE 将图像 → 转换为离散 token（词表 8192 个）
第 2 步：ViT 预测被掩码 patch 的 token

from transformers import BeitForMaskedImageModeling, BeitImageProcessor
import torch

# BEiT 掩码图像建模
class BEiTPretraining:
    def __init__(self, model_name="microsoft/beit-base-patch16-224"):
        self.model = BeitForMaskedImageModeling.from_pretrained(model_name)
        self.processor = BeitImageProcessor.from_pretrained(model_name)

    def create_bool_masked_pos(self, num_patches, mask_ratio=0.4):
        """生成掩码位置"""
        num_masked = int(num_patches * mask_ratio)
        mask = torch.zeros(num_patches, dtype=torch.bool)
        # 块状掩码（BEiT v2）
        masked_indices = torch.randperm(num_patches)[:num_masked]
        mask[masked_indices] = True
        return mask

    def forward(self, images):
        inputs = self.processor(images, return_tensors="pt")
        pixel_values = inputs.pixel_values

        num_patches = (224 // 16) ** 2  # 196
        bool_masked_pos = self.create_bool_masked_pos(num_patches)
        bool_masked_pos = bool_masked_pos.unsqueeze(0).expand(
            pixel_values.size(0), -1
        )

        outputs = self.model(
            pixel_values=pixel_values,
            bool_masked_pos=bool_masked_pos
        )

        return outputs.loss

import torch
import torch.nn as nn
from torch.nn import functional as F

class GPTLanguageModel(nn.Module):
    def __init__(self, vocab_size, block_size, n_embd, n_head, n_layer, dropout=0.1):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(*[
            TransformerBlock(n_embd, n_head, dropout=dropout)
            for _ in range(n_layer)
        ])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.block_size = block_size

    def forward(self, idx, targets=None):
        B, T = idx.shape

        tok_emb = self.token_embedding(idx)  # (B, T, C)
        pos_emb = self.position_embedding(
            torch.arange(T, device=idx.device)
        )  # (T, C)
        x = tok_emb + pos_emb

        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        if targets is None:
            return logits, None

        # 下一个 token 预测损失
        B, T, C = logits.shape
        logits_flat = logits.view(B * T, C)
        targets_flat = targets.view(B * T)
        loss = F.cross_entropy(logits_flat, targets_flat)

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature

            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

from transformers import BertForMaskedLM, BertTokenizer
import torch

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForMaskedLM.from_pretrained("bert-base-uncased")

def create_mlm_input(text, mask_prob=0.15):
    """生成掩码语言模型输入"""
    tokens = tokenizer.encode(text, return_tensors="pt")
    labels = tokens.clone()

    # 15% 随机掩码
    rand = torch.rand(tokens.shape)
    mask_arr = (rand < mask_prob) & (tokens != tokenizer.cls_token_id) & \
               (tokens != tokenizer.sep_token_id)

    # 掩码策略：
    # - 80%：替换为 [MASK] token
    # - 10%：替换为随机 token
    # - 10%：保留原始 token
    for i, masked in enumerate(mask_arr[0]):
        if masked:
            prob = torch.rand(1).item()
            if prob < 0.8:
                tokens[0][i] = tokenizer.mask_token_id
            elif prob < 0.9:
                tokens[0][i] = torch.randint(len(tokenizer), (1,))
            # 剩余 10%：保留原始 token

    # 未被掩码的位置不计入损失
    labels[~mask_arr] = -100

    return tokens, labels

# 计算 MLM 损失
text = "The quick brown fox jumps over the lazy dog"
input_ids, labels = create_mlm_input(text)

with torch.no_grad():
    outputs = model(input_ids=input_ids, labels=labels)
    print(f"MLM Loss: {outputs.loss.item():.4f}")

# 预测被掩码的 token
input_text = "The [MASK] brown fox"
inputs = tokenizer(input_text, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits

mask_idx = (inputs.input_ids == tokenizer.mask_token_id).nonzero()[0][1]
top_5 = logits[0, mask_idx].topk(5)
for token_id, score in zip(top_5.indices, top_5.values):
    print(f"{tokenizer.decode([token_id])}: {score:.3f}")

from transformers import T5ForConditionalGeneration, T5Tokenizer

tokenizer = T5Tokenizer.from_pretrained("t5-base")
model = T5ForConditionalGeneration.from_pretrained("t5-base")

# T5 将所有任务统一为文本到文本
tasks = {
    "翻译": "translate English to Korean: The weather is nice today",
    "摘要": "summarize: Scientists have discovered a new species of deep-sea fish...",
    "情感分析": "sentiment analysis: This movie was absolutely fantastic!",
    "问答": "question: What is the capital of France? context: Paris is the capital..."
}

for task_name, input_text in tasks.items():
    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
    with torch.no_grad():
        output_ids = model.generate(
            inputs.input_ids,
            max_length=128,
            num_beams=4,
            early_stopping=True
        )
    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    print(f"\n[{task_name}]")
    print(f"输入: {input_text[:60]}...")
    print(f"输出: {output}")

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import numpy as np

class LinearProbe(nn.Module):
    """在 SSL 特征之上的线性分类器"""
    def __init__(self, feature_dim, num_classes):
        super().__init__()
        self.linear = nn.Linear(feature_dim, num_classes)

    def forward(self, x):
        return self.linear(x)


def extract_features(model, dataloader, device='cuda'):
    """用预训练模型提取特征"""
    model.eval()
    features_list = []
    labels_list = []

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            features = model.encode(images)
            features_list.append(features.cpu())
            labels_list.append(labels)

    return torch.cat(features_list), torch.cat(labels_list)


def few_shot_evaluation(ssl_model, train_loader, val_loader,
                        n_shots=[1, 5, 10, 100], device='cuda'):
    """Few-shot 评估：用 n 个标签训练线性分类器"""
    ssl_model.eval()

    # 提取全部特征
    all_features, all_labels = extract_features(ssl_model, train_loader, device)
    val_features, val_labels = extract_features(ssl_model, val_loader, device)

    results = {}
    num_classes = len(all_labels.unique())

    for n_shot in n_shots:
        # 仅使用 n_shot 个标签
        selected_indices = []
        for cls in range(num_classes):
            cls_indices = (all_labels == cls).nonzero(as_tuple=True)[0]
            selected = cls_indices[torch.randperm(len(cls_indices))[:n_shot]]
            selected_indices.append(selected)

        selected_indices = torch.cat(selected_indices)
        train_feats = all_features[selected_indices]
        train_labs = all_labels[selected_indices]

        # 训练线性分类器
        probe = LinearProbe(all_features.shape[1], num_classes).to(device)
        optimizer = torch.optim.Adam(probe.parameters(), lr=1e-3)
        criterion = nn.CrossEntropyLoss()

        train_feats = train_feats.to(device)
        train_labs = train_labs.to(device)

        for epoch in range(100):
            probe.train()
            logits = probe(train_feats)
            loss = criterion(logits, train_labs)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # 验证
        probe.eval()
        with torch.no_grad():
            val_logits = probe(val_features.to(device))
            preds = val_logits.argmax(dim=1).cpu()
            acc = (preds == val_labels).float().mean().item()

        results[n_shot] = acc
        print(f"{n_shot}-shot 准确率: {acc * 100:.2f}%")

    return results


# k-NN 评估（无参数评估）
def knn_evaluation(ssl_model, train_loader, val_loader, k=20, device='cuda'):
    """用 k-NN 分类器评估 SSL 表示质量"""
    ssl_model.eval()

    train_features, train_labels = extract_features(ssl_model, train_loader, device)
    val_features, val_labels = extract_features(ssl_model, val_loader, device)

    # 归一化
    train_features = nn.functional.normalize(train_features, dim=1)
    val_features = nn.functional.normalize(val_features, dim=1)

    # 按批次做 k-NN 分类
    correct = 0
    total = 0
    batch_size = 256

    for i in range(0, len(val_features), batch_size):
        batch_feats = val_features[i:i+batch_size].to(device)
        batch_labels = val_labels[i:i+batch_size]

        # 计算余弦相似度
        sims = torch.mm(batch_feats, train_features.to(device).T)
        topk_sims, topk_indices = sims.topk(k, dim=1)

        # 多数投票
        topk_labels = train_labels[topk_indices.cpu()]
        preds = topk_labels.mode(dim=1).values

        correct += (preds == batch_labels).sum().item()
        total += len(batch_labels)

    accuracy = correct / total
    print(f"k-NN ({k}) 准确率: {accuracy * 100:.2f}%")
    return accuracy

# 用于医学图像的 SSL（以 CheXpert 胸部 X 光为例）
class MedicalSSL(nn.Module):
    """针对医学图像的自监督学习"""
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone

        # 针对医学图像的增强（弱增强）
        self.augmentation = T.Compose([
            T.RandomResizedCrop(224, scale=(0.8, 1.0)),  # 较小的裁剪变化
            T.RandomHorizontalFlip(p=0.5),
            # 医学图像应最小化颜色变化
            T.RandomApply([T.ColorJitter(0.1, 0.1, 0.0, 0.0)], p=0.3),
            T.ToTensor(),
            T.Normalize([0.485], [0.229])  # 黑白 X 光
        ])

    def forward(self, x):
        v1 = self.augmentation(x)
        v2 = self.augmentation(x)
        return v1, v2


# 用于卫星图像的 SSL
class SatelliteSSL:
    """针对卫星图像的增强（保留地理特性）"""
    def __init__(self):
        self.transform = T.Compose([
            T.RandomResizedCrop(256, scale=(0.4, 1.0)),
            # 卫星图像：旋转不变性很重要
            T.RandomRotation(90),
            T.RandomHorizontalFlip(),
            T.RandomVerticalFlip(),
            # 模拟时间段/季节变化
            T.ColorJitter(brightness=0.4, contrast=0.4),
            T.ToTensor(),
        ])


# 比较训练性能
def compare_ssl_methods(dataset, num_epochs=100):
    methods = {
        'SimCLR': SimCLR(backbone='resnet50'),
        'BYOL': BYOL(backbone=resnet50()),
        'MAE': MAE(img_size=224),
    }

    results = {}
    for name, model in methods.items():
        print(f"\n正在训练 {name}...")
        model = model.cuda()
        optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

        for epoch in range(num_epochs):
            train_one_epoch(model, dataset, optimizer)

        # 100-shot 评估
        acc = few_shot_evaluation(model, dataset, n_shots=[100])
        results[name] = acc[100]
        print(f"{name} 100-shot 准确率: {acc[100]*100:.2f}%")

    return results

# 使用 Hugging Face 上的预训练 SSL 模型
from transformers import (
    ViTModel, ViTFeatureExtractor,
    DeiTModel,
    CLIPModel, CLIPProcessor
)
from PIL import Image
import torch

# ViT（经 MAE 预训练的版本）
def load_mae_pretrained():
    model = ViTModel.from_pretrained("facebook/vit-mae-base")
    feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/vit-mae-base")
    return model, feature_extractor

# 提取 DINO ViT 特征
def extract_dino_features(image_path):
    from transformers import ViTModel
    model = ViTModel.from_pretrained("facebook/dino-vits16")
    processor = ViTFeatureExtractor.from_pretrained("facebook/dino-vits16")

    image = Image.open(image_path)
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)

    # [CLS] token 特征（整幅图像的表示）
    cls_features = outputs.last_hidden_state[:, 0, :]
    # 各 patch 特征（含空间信息）
    patch_features = outputs.last_hidden_state[:, 1:, :]

    return cls_features, patch_features

# 用 CLIP 提取图像特征后接下游任务
def clip_feature_extraction(images, texts=None):
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    if texts:
        inputs = processor(text=texts, images=images, return_tensors="pt", padding=True)
        with torch.no_grad():
            outputs = model(**inputs)
        return outputs.image_embeds, outputs.text_embeds
    else:
        inputs = processor(images=images, return_tensors="pt")
        with torch.no_grad():
            image_features = model.get_image_features(**inputs)
        return image_features