📎
日本語対応のJina CLIP v2をローカルGPU環境で動かしてみた
はじめに
CLIP（Contrastive Language-Image Pre-Training）について調べているときに「Jina CLIP v2」という比較的新しい（2024年11月発表）モデルを見つけたのでローカルのGPU環境で動かしてみました。
Jina CLIP v2とは？
詳しくは以下の公式ブログ記事を参照して頂くとして、私が気になったのは以下の点です。
- 日本語を含む89言語に対応している（OpenAI CLIPは英語のみ）
- 8,192トークンと長めの入力に対応している（OpenAI CLIPは77トークン）
- 512x512ピクセルという比較的高い解像度の画像の入力に対応している（OpenAI CLIPは224x224ピクセル）
- マトリョーシカ表現（Matryoshka Representations）により64〜1024次元の範囲で埋め込み表現を用途により選んで使うことができる（OpenAI CLIPは512次元固定）
詳しくは公式ブログ記事（日本語）を参照してください。
また、本記事ではローカルのGPU環境でJina CLIP v2を動かしていますが、APIとして使用する場合には以下のスクラップが参考になります。
ローカルで動かす
環境構築
今回は以下の環境で試しました。
- GPU: NVIDIA GeForce RTX 4080 16GB
- OS: Ubuntu 22.04.5 LTS
- NVIDIAドライバ: 550.127.05
- Docker: 24.0.7
- Docker Compose: 2.18.0
- NVIDIA Container Toolkit: 1.13.5
- Python: 3.12.8
色々ハマりつつも依存関係を整理してcompose.yaml、Dockerfile、requirements.txtを作りました。
以下のコマンドでPythonのREPL環境を起動できます。なお、ビルド後のイメージサイズは約14.6GBです。
docker-compose run --build --rm jina-clip-v2
compose.yaml
services:
jina-clip-v2:
build: "."
deploy:
resources:
reservations:
devices:
- driver: "nvidia"
capabilities: ["gpu"]
count: 1
Dockerfile
FROM python:3.12-slim
RUN \
rm -f /etc/apt/apt.conf.d/docker-clean \
&& apt-get update \
&& DEBIAN_FRONTEND=noninteractive apt-get install --yes --no-install-recommends \
git-core \
wget
RUN \
cd /root/ \
&& wget --quiet https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb \
&& dpkg -i cuda-keyring_1.1-1_all.deb \
&& rm cuda-keyring_1.1-1_all.deb \
&& apt-get update \
&& DEBIAN_FRONTEND=noninteractive apt-get install --yes --no-install-recommends \
cuda-toolkit-12-6
WORKDIR /app
COPY ./requirements.txt ./
RUN \
python -m pip install --requirement requirements.txt
RUN \
python -m pip install --no-build-isolation flash-attn==2.7.2.post1
requirements.txt
anykeystore==0.2
apex==0.9.10.dev0
certifi==2024.8.30
charset-normalizer==3.4.0
cryptacular==1.6.2
defusedxml==0.7.1
einops==0.8.0
filelock==3.16.1
fsspec==2024.10.0
greenlet==3.1.1
huggingface-hub==0.26.5
hupper==1.12.1
idna==3.10
Jinja2==3.1.4
MarkupSafe==3.0.2
mpmath==1.3.0
networkx==3.4.2
numpy==2.2.0
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127
oauthlib==3.2.2
packaging==24.2
PasteDeploy==3.1.0
pbkdf2==1.3
pillow==11.0.0
plaster==1.1.2
plaster-pastedeploy==1.0.1
psutil==6.1.0
pyramid==2.0.2
pyramid-mailer==0.15.1
python3-openid==3.2.0
PyYAML==6.0.2
regex==2024.11.6
repoze.sendmail==4.4.1
requests==2.32.3
requests-oauthlib==2.0.0
safetensors==0.4.5
setuptools==75.6.0
SQLAlchemy==2.0.36
sympy==1.13.1
timm==1.0.12
tokenizers==0.21.0
torch==2.5.1
torchvision==0.20.1
tqdm==4.67.1
transaction==5.0
transformers==4.47.0
translationstring==1.4
triton==3.1.0
typing_extensions==4.12.2
urllib3==2.2.3
velruse==1.1.1
venusian==3.1.1
WebOb==1.8.9
WTForms==3.2.1
wtforms-recaptcha==0.3.2
xformers==0.0.28.post3
zope.deprecation==5.0
zope.interface==7.2
zope.sqlalchemy==3.1
なお、手動で環境を構築する際には、以下のPythonパッケージをインストールする必要があります。
pip install torch timm transformers einops xformers flash-attn apex
依存ライブラリについては、以下に記載があります。
実行
OpenAI CLIPの実行例をJina CLIP v2向けに改変して実行してみます。
from transformers import AutoModel, AutoProcessor
from PIL import Image
import requests
model = AutoModel.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True).to("cuda")
processor = AutoProcessor.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True)
outputs = model(
input_ids=inputs["input_ids"].to("cuda"),
attention_mask=inputs["attention_mask"].to("cuda"),
pixel_values=inputs["pixel_values"].to("cuda"),
)
probs = outputs.logits_per_image.softmax(dim=1)
Dockerコンテナ内での実行結果は以下の通りです。
なお、初回は約2GBのモデルのダウンロードが行われますのでしばし待ちます。
Python 3.12.8 (main, Dec 4 2024, 20:39:59) [GCC 12.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from transformers import AutoModel, AutoProcessor
>>> from PIL import Image
>>> import requests
>>> model = AutoModel.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True).to("cuda")
...
model.safetensors: 100%|████████████████████████████████████████████████████████████████████████████████| 1.73G/1.73G [05:12<00:00, 5.53MB/s]
...
>>> processor = AutoProcessor.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True)
...
>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)
>>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True)
>>> inputs
{'input_ids': tensor([[ 0, 10, 16186, 111, 10, 7515, 2],
[ 0, 10, 16186, 111, 10, 10269, 2]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1]]), 'pixel_values': tensor([[[[ 0.5143, 0.5727, 0.6603, ..., -0.0259, -0.1280, -0.0696],
[ 0.5435, 0.6165, 0.6603, ..., 0.0325, 0.0179, -0.0113],
[ 0.5435, 0.5873, 0.6019, ..., 0.0179, 0.0179, 0.1347],
...,
[ 1.8135, 1.8573, 1.9157, ..., 1.4778, 1.4486, 1.5654],
[ 1.9157, 1.8573, 1.8865, ..., 1.2588, 1.1712, 1.6238],
[ 1.8719, 1.8573, 1.9011, ..., 1.1712, 1.4486, 1.5654]],
[[-1.4069, -1.3469, -1.2568, ..., -1.4970, -1.5870, -1.4970],
[-1.3769, -1.2718, -1.2268, ..., -1.4519, -1.4369, -1.4519],
[-1.3469, -1.2718, -1.2418, ..., -1.5120, -1.4669, -1.3769],
...,
[ 0.0789, 0.0939, 0.1389, ..., -0.6565, -0.6565, -0.5665],
[ 0.1689, 0.1089, 0.0789, ..., -0.8816, -0.9117, -0.4164],
[ 0.0939, 0.0789, 0.1239, ..., -0.9117, -0.5965, -0.4464]],
[[-0.6839, -0.5275, -0.3426, ..., -0.8545, -0.8830, -0.8972],
[-0.4279, -0.4137, -0.4137, ..., -0.8972, -0.8688, -0.7977],
[-0.4279, -0.4422, -0.4564, ..., -0.8972, -0.8403, -0.6981],
...,
[ 1.6340, 1.5771, 1.6482, ..., 0.8803, 0.6528, 0.7523],
[ 1.6340, 1.5913, 1.7193, ..., 0.9656, 0.5248, 0.7808],
[ 1.7193, 1.6340, 1.6055, ..., 0.8661, 0.8661, 0.9941]]]])}
>>> outputs = model(
... input_ids=inputs["input_ids"].to("cuda"),
... attention_mask=inputs["attention_mask"].to("cuda"),
... pixel_values=inputs["pixel_values"].to("cuda"),
... )
>>> outputs
CLIPOutput(loss=None, logits_per_image=tensor([[16.1250, 9.7500]], device='cuda:0', dtype=torch.bfloat16,
grad_fn=<TBackward0>), logits_per_text=tensor([[16.1250],
[ 9.7500]], device='cuda:0', dtype=torch.bfloat16,
grad_fn=<MulBackward0>), text_embeds=tensor([[-0.1514, 0.0217, 0.0742, ..., 0.0015, -0.0176, 0.0089],
[-0.1069, -0.0010, -0.0154, ..., 0.0131, -0.0114, 0.0130]],
device='cuda:0', dtype=torch.bfloat16, grad_fn=<DivBackward0>), image_embeds=tensor([[ 0.0054, 0.0986, -0.0322, ..., 0.0054, 0.0027, 0.0065]],
device='cuda:0', dtype=torch.bfloat16, grad_fn=<DivBackward0>), text_model_output=None, vision_model_output=None)
>>> probs = outputs.logits_per_image.softmax(dim=1)
>>> probs
tensor([[1.0000, 0.0017]], device='cuda:0', dtype=torch.bfloat16,
grad_fn=<SoftmaxBackward0>)
「a photo of a cat」と「a photo of a dog」であれば前者の方が強く関連しているという結果が出力されており、上手く動作しているようです。
せっかくなので、日本語でも確認してみましょう。
inputs = processor(text=["風景の写真", "猫の写真", "犬の写真"], images=image, return_tensors="pt", padding=True)
outputs = model(
input_ids=inputs["input_ids"].to("cuda"),
attention_mask=inputs["attention_mask"].to("cuda"),
pixel_values=inputs["pixel_values"].to("cuda"),
)
probs = outputs.logits_per_image.softmax(dim=1)
>>> inputs = processor(text=["風景の写真", "猫の写真", "犬の写真"], images=image, return_tensors="pt", padding=True)
>>> inputs["input_ids"]
tensor([[ 0, 6, 80791, 137146, 2],
[ 0, 6, 35076, 137146, 2],
[ 0, 6, 37084, 137146, 2]])
>>> outputs = model(
... input_ids=inputs["input_ids"].to("cuda"),
... attention_mask=inputs["attention_mask"].to("cuda"),
... pixel_values=inputs["pixel_values"].to("cuda"),
... )
>>> probs = outputs.logits_per_image.softmax(dim=1)
>>> probs
tensor([[0.0020, 0.9766, 0.0203]], device='cuda:0', dtype=torch.bfloat16,
grad_fn=<SoftmaxBackward0>)
「風景の写真」、「猫の写真」、「犬の写真」のうち2番目の値が高く、こちらも正常に動作しているようです。
また、inputs["input_ids"]を見てみると日本語が適切にトークン化されていることが確認できます。
ハマったポイント
Jina CLIP v2のローカルGPU環境での実行にあたり、ハマったポイントを記録しておきます。
ローカル環境で動かす方法が載っていない
Jina CLIP v2のモデルファイル自体はHugging Faceにホスティングされていますが、動かし方に関する記載はありませんでした。
設定ファイルを参照したところ、AutoModel、AutoProcessorに対応していることが分かったので、無事に動かすことができました。
flash-attnのインストールに失敗する
依存ライブラリのインストールに際して、flash-attnのインストール（ビルド）に失敗してしまうという問題に遭遇しました。
ModuleNotFoundError: No module named 'torch'と出力されていますが、PyTorchは問題なくインストールされており、動作する状況です。
# pip install flash-attn
Collecting flash-attn
Downloading flash_attn-2.7.2.post1.tar.gz (3.1 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.1/3.1 MB 6.9 MB/s eta 0:00:00
Installing build dependencies ... done
Getting requirements to build wheel ... error
error: subprocess-exited-with-error
× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> [20 lines of output]
Traceback (most recent call last):
File "/usr/local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
main()
File "/usr/local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
json_out['return_val'] = hook(**hook_input['kwargs'])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 118, in get_requires_for_build_wheel
return hook(config_settings)
^^^^^^^^^^^^^^^^^^^^^
File "/tmp/pip-build-env-bimk_e2h/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 334, in get_requires_for_build_wheel
return self._get_build_requires(config_settings, requirements=[])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/tmp/pip-build-env-bimk_e2h/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 304, in _get_build_requires
self.run_setup()
File "/tmp/pip-build-env-bimk_e2h/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 522, in run_setup
super().run_setup(setup_script=setup_script)
File "/tmp/pip-build-env-bimk_e2h/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 320, in run_setup
exec(code, locals())
File "<string>", line 21, in <module>
ModuleNotFoundError: No module named 'torch'
[end of output]
note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error
× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.
note: This error originates from a subprocess, and is likely not a problem with pip.
最終的に以下を行うことで解決しました。
- CUDA Toolkitをインストールする
-
pip install --no-build-isolation flash-attnのように--no-build-isolationオプション付きでインストールする
おわりに
Jina CLIP v2の基本的な動作を確認することができました。
日本語に対応し、比較的解像度の高い画像も扱えるので色々と応用が広がりそうです。
なお、Jina CLIP v2はAPIを介しての利用は商用利用が可能ですが、モデルファイルをダウンロードしての利用は商用利用に制限があるようです。
詳しくはライセンスをご確認ください。
本記事が何らかの参考になれば幸いです。
Discussion