CLIP：打通图文壁垒的多模态神器，原理与实战全解析 原创

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import clip
import os
from tqdm import tqdm


# 1. 数据集定义
class CLIPDataset(Dataset):
    def __init__(self, image_dir, text_file, transform=None):
        """
        image_dir: 图像文件夹路径
        text_file: 文本文件路径，每行格式为"图像文件名,文本描述"
        """
        self.image_dir = image_dir
        self.transform = transform
        self.data = []
        with open(text_file, 'r', encoding='utf-8') as f:
            for line in f:
                img_name, text = line.strip().split(',', 1)
                self.data.append((img_name, text))


    def __len__(self):
        return len(self.data)


    def __getitem__(self, idx):
        img_name, text = self.data[idx]
        img_path = os.path.join(self.image_dir, img_name)
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, text


# 2. 数据增强与加载
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, cnotallow=0.2, saturatinotallow=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# 假设训练数据存放在./train_images，文本文件为./train_texts.txt
dataset = CLIPDataset(
    image_dir='./train_images',
    text_file='./train_texts.txt',
    transform=transform
)
dataloader = DataLoader(
    dataset,
    batch_size=256,  # 根据GPU显存调整
    shuffle=True,
    num_workers=8,
    pin_memory=True
)


# 3. 加载预训练编码器（或初始化）
device = "cuda" if torch.cuda.is_available() else "cpu"
image_encoder, text_encoder = clip.load("ViT-B/32", device=device)


# 若从头训练，可初始化模型
# from clip.model import VisionTransformer, TextTransformer
# image_encoder = VisionTransformer(...).to(device)
# text_encoder = TextTransformer(...).to(device)


# 4. 定义损失函数与优化器
temperature = 0.07  # 温度参数
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(
    list(image_encoder.parameters()) + list(text_encoder.parameters()),
    lr=1e-4,
    weight_decay=1e-6  # L2正则化
)


# 学习率衰减
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)


# 5. 训练循环
num_epochs = 30
for epoch in range(num_epochs):
    image_encoder.train()
    text_encoder.train()
    total_loss = 0.0


    for images, texts in tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}"):
        images = images.to(device)
        texts = clip.tokenize(texts).to(device)  # 文本tokenize


        # 编码图像和文本
        with torch.no_grad():  # 若冻结部分层，在此处设置
            image_features = image_encoder(images)
            text_features = text_encoder(texts)


        # L2归一化
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)


        # 计算相似度矩阵
        logits = (image_features @ text_features.T) / temperature


        # 计算对比损失（双向）
        labels = torch.arange(logits.shape[0], device=device)
        loss = (criterion(logits, labels) + criterion(logits.T, labels)) / 2


        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


        total_loss += loss.item() * images.size(0)


    # 计算平均损失
    avg_loss = total_loss / len(dataset)
    print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
    scheduler.step()


    # 保存模型
    if (epoch + 1) % 5 == 0:
        torch.save({
            'image_encoder': image_encoder.state_dict(),
            'text_encoder': text_encoder.state_dict(),
            'optimizer': optimizer.state_dict()
        }, f"clip_epoch_{epoch+1}.pth")


print("训练完成！")

pip install torch ftfy regex tqdm  # 基础依赖
pip install git+https://github.com/openai/CLIP.git  # CLIP库
pip install gradio faiss-cpu  # 可视化与向量检索

import clip
import torch
from PIL import Image


# 加载模型和预处理函数
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)


# 图片预处理：缩放、归一化等
def process_image(img_path):
    image = Image.open(img_path).convert("RGB")
    return preprocess(image).unsqueeze(0).to(device)


# 文本预处理：tokenize（转换成模型可理解的格式）
def process_text(text):
    return clip.tokenize([text]).to(device)

import faiss
import os
import numpy as np


# 图片文件夹路径
img_dir = "path/to/your/images"
img_paths = [os.path.join(img_dir, f) for f in os.listdir(img_dir) if f.endswith(('jpg', 'png'))]


# 提取所有图片的向量
img_features = []
for path in img_paths:
    img = process_image(path)
    with torch.no_grad():
        feat = model.encode_image(img)  # 图像编码
        feat /= feat.norm(dim=-1, keepdim=True)  # 归一化
        img_features.append(feat.cpu().numpy())


# 构建FAISS索引
d = 512  # CLIP向量维度通常为512
index = faiss.IndexFlatL2(d)  # 用L2距离衡量相似度
index.add(np.vstack(img_features))  # 加入所有图片向量

import gradio as gr


def search_images(query):
    # 文本编码
    text = process_text(query)
    with torch.no_grad():
        text_feat = model.encode_text(text)
        text_feat /= text_feat.norm(dim=-1, keepdim=True)


    # 检索最相似的5张图
    _, indices = index.search(text_feat.cpu().numpy(), k=5)
    return [img_paths[i] for i in indices[0]]


# 搭建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("## 🔍 CLIP以文搜图演示")
    query = gr.Textbox(label="输入文本描述（如'a red dress'）")
    btn = gr.Button("搜索")
    gallery = gr.Gallery(label="检索结果")
    btn.click(fn=search_images, inputs=query, outputs=gallery)


demo.launch()