x-clip

一个简洁但完整的<a href="https://openai.com/blog/clip/">CLIP</a>实现，包含了来自最近论文的各种实验性改进

安装

$ pip install x-clip

使用方法

import torch
from x_clip import CLIP

clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 10000,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8,
    visual_patch_dropout = 0.5,             # 图像块dropout概率，用于Kaiming He的FLIP中以节省计算并改善最终结果 - 0.5是一个好值，0.75是可接受的上限
    use_all_token_embeds = False,           # 是否使用细粒度对比学习（FILIP）
    decoupled_contrastive_learning = True,  # 使用解耦对比学习（DCL）目标函数，从InfoNCE损失的分母中移除正样本对（CLOOB + DCL）
    extra_latent_projection = True,         # 是否为文本到图像和图像到文本的比较使用单独的投影（CLOOB）
    use_visual_ssl = True,                  # 是否对图像进行自监督学习
    use_mlm = False,                        # 对文本使用掩码语言学习（MLM）（DeCLIP）
    text_ssl_loss_weight = 0.05,            # 文本MLM损失的权重
    image_ssl_loss_weight = 0.05            # 图像自监督学习损失的权重
)

# 模拟数据

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

# 训练

loss = clip(
    text,
    images,
    freeze_image_encoder = False,   # 如果使用预训练的图像网络，是否冻结图像编码器，由LiT论文提出
    return_loss = True              # 需要设置为True以返回对比损失
)

loss.backward()

你也可以传入外部的视觉transformer或残差网络。你只需确保你的图像编码器返回形状为batch x seq x dim的一组嵌入，并确保正确指定dim_image为返回嵌入的维度。以下是使用vit_pytorch中的视觉transformer的示例

$ pip install vit_pytorch>=0.25.6

import torch
from x_clip import CLIP

from vit_pytorch import ViT
from vit_pytorch.extractor import Extractor

base_vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 512,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

vit = Extractor(
    base_vit,
    return_embeddings_only = True
)

clip = CLIP(
    image_encoder = vit,
    dim_image = 512,           # 必须设置为与上面的视觉transformer相同的维度
    dim_text = 512,
    dim_latent = 512,
    num_text_tokens = 10000,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

loss = clip(text, images, return_loss = True)
loss.backward()

最后，也可以外部定义文本transformer。目前，它需要返回包括CLS令牌在内的嵌入。

import torch
from x_clip import CLIP, TextTransformer

from vit_pytorch import ViT
from vit_pytorch.extractor import Extractor

base_vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 512,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

image_encoder = Extractor(
    base_vit,
    return_embeddings_only = True
)

text_encoder = TextTransformer(
    dim = 512,
    num_tokens = 10000,
    max_seq_len = 256,
    depth = 6,
    heads = 8
)

clip = CLIP(
    image_encoder = image_encoder,
    text_encoder = text_encoder,
    dim_image = 512,
    dim_text = 512,
    dim_latent = 512
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

loss = clip(text, images, return_loss = True)
loss.backward()

多视图对比学习损失

本仓库还支持多视图对比学习损失，如<a href="https://arxiv.org/abs/2110.05208">DeCLIP</a>中提出的。只需传入增强的文本和/或增强的图像，它就会自动计算，并按初始化时设置的multiview_loss_weight进行加权。例如：

import torch
from x_clip import CLIP, TextTransformer

from vit_pytorch import ViT
from vit_pytorch.extractor import Extractor

base_vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 512,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

image_encoder = Extractor(
    base_vit,
    return_embeddings_only = True
)

text_encoder = TextTransformer(
    dim = 512,
    num_tokens = 10000,
    max_seq_len = 256 + 1,
    depth = 6,
    heads = 8
)

clip = CLIP(
    image_encoder = image_encoder,
    text_encoder = text_encoder,
    dim_image = 512,
    dim_text = 512,
    dim_latent = 512,
    extra_latent_projection = True,
    multiview_loss_weight = 0.1         # 将多视图对比损失的权重设为0.1
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

aug_text = torch.randint(0, 10000, (4, 256))  # 增强文本（回译或EDA），与text维度相同
aug_images = torch.randn(4, 3, 256, 256)      # 增强图像，与上面的images维度相同
loss = clip(
    text,
    images,
    aug_text = aug_text,           # 传入增强文本
    aug_image = aug_images,        # 传入增强图像
    return_loss = True,
    freeze_image_encoder = True
)

loss.backward()

你甚至可以传入多个增强文本或图像

# ...

aug_texts = (
    torch.randint(0, 10000, (4, 256)),
    torch.randint(0, 10000, (4, 256)),
)

aug_images = (
    torch.randn(4, 3, 256, 256),
    torch.randn(4, 3, 256, 256),
)

loss = clip(
    text,
    images,
    aug_text = aug_texts,
    aug_image = aug_images,
    return_loss = True,
    freeze_image_encoder = True
)

loss.backward()

自定义视觉自监督学习模块

你可以通过visual_ssl关键字传入自己的视觉自监督学习模块，如下所示：

import torch
from x_clip import CLIP
from x_clip.visual_ssl import SimSiam

from vit_pytorch import ViT
from vit_pytorch.extractor import Extractor

base_vit = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 512,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)

image_encoder = Extractor(
    base_vit,
    return_embeddings_only = True
)

visual_ssl = SimSiam(                 # 外部定义的SimSiam - 需要是一个接受与CLIP相同维度图像并返回标量损失的模块
    image_encoder,
    image_size = 256,
    hidden_layer = -1
)

clip = CLIP(
    image_encoder = image_encoder,
    dim_image = 512,
    dim_text = 512,
    dim_latent = 512,
    use_mlm = True,
    visual_ssl = visual_ssl,           # SSL模块传入CLIP
    use_all_token_embeds = False,
    extra_latent_projection = False,
    mlm_random_token_prob = 0.1
)

text = torch.randint(0, 10000, (4, 256))
images = torch.randn(4, 3, 256, 256)

loss = clip(text, images, return_loss = True)
loss.backward()

引用

@misc{radford2021learning,
    title   = {Learning Transferable Visual Models From Natural Language Supervision}, 
    author  = {Alec Radford and Jong Wook Kim and Chris Hallacy and Aditya Ramesh and Gabriel Goh and Sandhini Agarwal and Girish Sastry and Amanda Askell and Pamela Mishkin and Jack Clark and Gretchen Krueger and Ilya Sutskever},
    year    = {2021},
    eprint  = {2103.00020},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}

@misc{yao2021filip,
    title   = {FILIP: Fine-grained Interactive Language-Image Pre-Training}, 
    author  = {Lewei Yao and Runhui Huang and Lu Hou and Guansong Lu and Minzhe Niu and Hang Xu and Xiaodan Liang and Zhenguo Li and Xin Jiang and Chunjing Xu},
    year    = {2021},
    eprint  = {2111.07783},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}

@misc{fürst2021cloob,
    title   = {CLOOB: Modern Hopfield Networks with InfoLOOB Outperform CLIP},
    author  = {Andreas Fürst and Elisabeth Rumetshofer and Viet Tran and Hubert Ramsauer and Fei Tang and Johannes Lehner and David Kreil and Michael Kopp and Günter Klambauer and Angela Bitto-Nemling and Sepp Hochreiter},
    year    = {2021},
    eprint  = {2110.11316},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG}
}

@misc{yeh2021decoupled,
    title   = {解耦对比学习},
    author  = {叶骏晓 and 洪承耀 and 许彦齐 and 刘庭伦 and 陈宇北 and Yann LeCun},
    year    = {2021},
    eprint  = {2110.06848},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG}
}

@misc{zhai2021lit,
    title   = {LiT: 使用锁定图像文本微调进行零样本迁移},
    author  = {翟晓华 and 王笑 and Basil Mustafa and Andreas Steiner and Daniel Keysers and Alexander Kolesnikov and Lucas Beyer},
    year    = {2021},
    eprint  = {2111.07991},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}

@misc{li2021supervision,
    title   = {监督无处不在：一种数据高效的对比语言-图像预训练范式},
    author  = {李阳光 and 梁峰 and 赵立晨 and 崔宇峰 and 欧阳万里 and 邵静 and 于凤伟 and 颜俊杰},
    year    = {2021},
    eprint  = {2110.05208},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}

@Article{mu2021slip,
    author  = {Norman Mu and Alexander Kirillov and David Wagner and 谢赛宁},
    title   = {SLIP: 自监督遇上语言-图像预训练},
    journal = {arXiv预印本 arXiv:2112.12750},
    year    = {2021},
}

@misc{su2021roformer,
    title   = {RoFormer: 具有旋转位置嵌入的增强型Transformer},
    author  = {苏剑林 and 卢钰 and 潘胜峰 and 温博 and 刘云峰},
    year    = {2021},
    eprint  = {2104.09864},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL}
}

@inproceedings{anonymous2022normformer,
    title   = {NormFormer: 通过额外归一化改进Transformer预训练},
    author  = {匿名},
    booktitle = {提交至第十届国际学习表示会议},
    year    = {2022},
    url     = {https://openreview.net/forum?id=GMYWzWztDx5},
    note    = {审核中}
}

@inproceedings{Li2022ScalingLP,
    title   = {通过掩码扩展语言-图像预训练},
    author  = {李扬豪 and 范浩琦 and 胡荣航 and Christoph Feichtenhofer and 何恺明},
    year    = {2022}
}

@article{Liu2022PatchDropoutEV,
    title   = {PatchDropout: 使用补丁丢弃来节约视觉Transformer资源},
    author  = {刘悦 and Christos Matsoukas and Fredrik Strand and Hossein Azizpour and Kevin Smith},
    journal = {ArXiv},
    year    = {2022},
    volume  = {abs/2208.07220}
}

@misc{shi2023enhance,
    title   = {通过表示相似性正则化增强音频生成的可控性}, 
    author  = {石阳阳 and Gael Le Lan and Varun Nagaraja and 倪昭恒 and 梅鑫浩 and 张义 and Forrest Iandola and 刘洋 and Vikas Chandra},
    year    = {2023},
    eprint  = {2309.08773},
    archivePrefix = {arXiv},
    primaryClass = {cs.SD}
}