小型LLMでマルチモーダル学習～推論を試してみた
研究開発本部 樋口栄作
はじめに
最近GPT-4oが発表され、紹介動画を見ているとAIが画像を当たり前に見て、答えるようになっていますね。一緒に散歩したり、買い物したり、ボードゲームをAIと遊ぶなんて日も遠くないのかもしれない、そんな期待のある最近のLLM界隈ですが、この画像を見て回答する、という技術は一体どうやっているのか？気になりますね。
技術を学ぶにも、手元で動かせるものがあるとやりやすい＆わかりやすい、ということは往々にあるかと思います。本記事では、画像＋テキストなど、いわゆるマルチモーダルの技術に触れてみよう！ということで、LLaVA (Large Language and Vision Assistant) を試してみました。
LLaVA の公式リポジトリでは、学習済みモデルとして7Bのモデル (llama-2-7b-chat をベースにしたモデル) を提供していまして、すぐに推論を試せるようにしてくれています。しかし、せっかくなので学習も試したいところです。
よってこの記事では、できるだけ小さなモデル (TinnyLamma 1.1B) を学習して、推論させた結果を見てみようと思います。
LLaVA (Large Language and Vision Assistant) について
今回使う LLaVAとは、画像を詳細なテキスト情報に変換できる研究で、構成としては、下の図にあるような Vision Encoder と、 Projection と、 LLM の3つの要素で成り立っています。
Vision Encoder は、画像特徴を取り出すもので、CLIP[2]を用いています。
ただし、そのまま LLM へ画像特徴がいれられないので、Projection によって LLM 向けに調整し、LLM に入力する、という流れになっています。
学習をする際には、2段階に分けて学習をするようです。
1段階目は、Projection のみを学習し、他は固定。2段階目は、Projection と LLM を学習し、Vision Encoder は固定、として学習をします。
この1段階目では、画像特徴から言語トークンへの変換を学習し、LLM が扱える言語トークンへと整える目的で実施されます。
そして2段階目は、ユースケースに合わせたファインチューニングができるような構成をとっています。論文[1]では、チャットボットに向けたチューンと、科学的なQ&Aに特化するようなチューンの場合が挙げられています。
すなわち、1段階目は、純粋に画像を言葉に置き換える方法だけを学び、2段階目でどうやって回答するか？を学ぶ、という仕組みをとっていると考えられます。
今回は、1段階目の学習はもちろん、2段階目では、チャットボット向けのファインチューニングを試してみます。
学習の要件
学習を始める前に、今回用いた機材について紹介します。今回は下記のスペックのPCで、学習、推論を試しています。
Ubuntu 22.04
CPU: i7-9700K
RAM: 32GB
GPU: 2080Ti (VRAM 12GB)
ライブラリ、ソフト
CUDA: 12.1
Python: 3.10.12
学習データについてですが、公式の手順にできるだけ則り、pretrain と finetune で２つのデータを使います。しかし、データセットのすべてを学習すると、pretrain だけで30時間越え。ちょっと試すだけには長い時間がかかりそうでした・・・
なので、学習データをかなり絞って学習を試します。今回では、両方ともおよそ1/10程度のデータにしています。
学習だけの時間ですが、2時間半～3時間x2 の合計5~6時間くらいで完了しています。
次に、学習に使う LLM のモデルについてですが、公式では 7B の LLM を使っており、今回の機材では、ちょっと学習できそうにないです。
そこで、小さなモデルを探して代わりにしています。どうやら Llama2 の系列だと代用しやすいようでしたので、 TinyLlama (1.1B) のモデルを使って学習しています。
無理やり学習させるために、データもモデルも変えてしまいました。あんまり賢く返してくれないかもしれないですね。
公式リポジトリの取得と環境構築
さて、ここからは実際に手を動かして学習～推論までをやってみます。さっそく環境構築から始めましょう。
ドライバや CUDA などはすでに設定されている前提としています。
公式のリポジトリを clone して、Readme に沿ってモジュールをインストールしていきます。
LLaVA は tag v1.2.2.post1 を用います。clone が終わった後に git checkout でバージョンを変更しています。
また、今回は、tensorboard でログを出力させています。もしインストールされていなければ、入れておきましょう。
git clone https://github.com/haotian-liu/LLaVA.git
cd LLaVA
# バージョンはv.1.2.2.post1
git checkout v1.2.2.post1
pip install --upgrade pip
pip install -e .
pip install -e ".[train]"
pip install flash-attn --no-build-isolation
# tensorboardX のインストール
pip install tensorboardXpretrain の実行
pretrain データセット
今回はできるだけ手軽に試したい、ということで学習データを減らしてしまいます。学習データは、画像のデータとテキストのデータ、の2種がありまして、それぞれを調整していきます。
pretrain では、画像データセットに、LLaVA-CC3M-Pretrain-595K を用います。下のコマンドで、画像のダウンロードと解凍をします。画像データ解凍するとき、必要な分だけを pretrain/images へ取り出します。
mkdir -p dataset/pretrain/images
cd dataset
# 画像データのダウンロード & 間引き
wget https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K/resolve/main/images.zip
unzip images GCC_train_000[0-2][0-9]* -d pretrain/images/
# テキストデータのダウンロード
wget https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K/resolve/main/chat.jsonさらに、必要な分だけになった画像データに合わせて、テキストデータの方も調整していきます。調整したテキストデータは、 custom_chat.json として保存します。
# データセット用にピックアップ → custom_chat.json として保存
import os
import json
dataset_dir = "./pretrain/images"
json_dataset_path = "chat.json"
output_json = "custom_chat.json"
def custom_dataset(dataset_dir, json_dataset_path, output_json):
with open(json_dataset_path) as f:
dataset = json.load(f)
checked = 0
pickup = []
for data in dataset:
if "image" in data.keys():
image_path = data["image"]
if os.path.isfile(os.path.join(dataset_dir, image_path)):
checked += 1
pickup.append(data)
with open(output_json, "w") as f:
json.dump(pickup, f, indent=4)
print(f"result: {checked}/{len(dataset)}")
# run
custom_dataset(dataset_dir, json_dataset_path, output_json)pretrain 学習スクリプト
公式で、scripts/pretrain.sh という学習用スクリプトを用意してくれています。このスクリプトを修正して、学習させてみます。
まずは、データセット関連のパスを変更します。L17, L19, L20, 飛んでL27、さらに一番最後の行L46を変更します。
# L17 ~ L20
--model_name_or_path $MODEL_VERSION \
--version $PROMPT_VERSION \
--data_path ./dataset/custom_chat.json \
--image_folder ./dataset/pretrain/images \
# L27
--output_dir ./results/llava-$MODEL_VERSION-pretrain \
# L46
--report_to tensorboardつぎに、モデル設定を変更します。L7を変更します
# L7
MODEL_VERSION=TinyLlama/TinyLlama-1.1B-Chat-v1.0最後に、古いGPUで動かすためのオプションを変更します。
# L16
--bf16 False
# L41
--tf32 False学習プログラムの修正
新しめのPCでは、特に修正はいらないのですが、古いマシンで動かす場合は、flash_attention_2 が動かない事があります。
今回のPCでは、残念ながら動かないので llava/train/train_mem.py の L4 を修正します。
# L3~L4; L4のみ修正
if __name__ == "__main__":
train()pretrain の実行
さて、ここまで来たら学習開始できます。編集した scripts/pretrain.sh を実行してみます。結果は results に保存しています。
# 結果の保存先を作成
mkdir results
# pretrain 開始
bash scripts/pretrain.shpretrain の実行では、VRAM の使用量は、10GB 程度。RAM の使用量は 7GB 程度でした。
finetune の実行
つづいて、finetune も動かしてみます。pretrain の時とおなじく、まずはデータセットの準備から始めます。
finetune データセット
データセットが異なりますが、手順は pretrain の時と同じです。
finetune では、COCOデータセットを使います。ダウンロードして、同様にデータを減らしたカスタムデータセットを作ります。
mkdir -p dataset/train/images
cd dataset
# 画像のダウンロード & 間引き
wget http://images.cocodataset.org/zips/train2017.zip
unzip train2017.zip train2017/0000000[4-9][0-9]*.jpg -d train/images
# テキストデータのダウンロード
wget https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K/resolve/main/llava_instruct_80k.json?download=true -O llava_instruct_80k.jsonデータを減らしたバージョンの json を作ります。colab や jupyter などの notebook であれば、そのまま実行できます。それ以外では、pretrain の学習データを準備した際に用いた python コードの変数の値を変えて、動かしてみてください。
dataset_dir = "./train/images/train2017"
json_dataset_path = "llava_instruct_80k.json"
output_json = "custom_llava_instruct_80k.json"
# run
custom_dataset(dataset_dir, json_dataset_path, output_json)finetune スクリプト
finetune の実行には、scripts/finetune_lora.sh を使います。pretrain と同じく、パスや設定を変更します。
（スクリプト変更前に、dataset ディレクトリから抜け出ておきます。）
# もし、dataset/ であれば LLaVA/ に移動
cd ../scripts/finetune_lora.sh にデータセットと学習結果関連のパスの変更をします。L20, L22, L23, L25, 飛んでL30, L49を変更します。
# L20 ~ L25
--model_name_or_path $MODEL_VERSION \
--version $PROMPT_VERSION \
--data_path ./dataset/custom_llava_instruct_80k.json \
--image_folder ./dataset/train/images/train2017 \
--vision_tower openai/clip-vit-large-patch14 \
--pretrain_mm_mlp_adapter ./results/llava-$MODEL_VERSION-pretrain/mm_projector.bin \
# L30
--output_dir ./results/llava-$MODEL_VERSION-finetune_lora \
# L49
--report_to tensorboardつぎに、モデル設定を変更します。L13, L14を変更します。
PROMPT_VERSION="llava_llama_2"
MODEL_VERSION="TinyLlama/TinyLlama-1.1B-Chat-v1.0"最後に、古いGPUで動かすためのオプションを変更します。
pretrain と一部同じですが、finetune ではより多くのメモリが必要でした。このために、メモリに乗るようにバッチサイズを16から3まで減らしています。変更は、L29, L32, L44 になります。
# L29
--bf16 False \
# L32
--per_device_train_batch_size 3 \
# 44
--tf32 False \学習プログラムの修正
pretrainで llava/train/train_mem.py を編集したならば、再度の編集は不要です。もし、古いGPUで動かすならば、pretrain での修正を参考に、flash_attention_2 を使わないように変更してください。
finetune の実行
ここまでで準備は完了です。それでは finetune も試してみましょう。pretrain と同じく結果は、 results に保存しています。
# finetune 開始
bash scripts/finetune_lora.sh実行中の VRAM 使用量は、10GBギリギリ。VRAM が12GBの PC および GPU でないと学習は難しいかもしれません。このときの RAM 使用量は、6 GB 程度でした。
今回の結果は、per_device_train_batch_size 4 としていますが、1まで減らすと VRAM 8GB の PC でもなんとか学習できるかもしれません。
モデルをマージ
pretrain、finetune と学習していきましたが、推論するにはもうひと手間必要で、モデルをマージする必要があります。
といっても、マージも公式からスクリプトが出ています。今回向けにパスを編集して実行するだけで OK です。
学習と違い、マージはすぐに終わります。
一点注意で、このマージスクリプトの前に、finetune 結果のファイル名にllava を含めるように変更してください。ファイル名に llava が含まれないと正しく処理が進まないことがあります。
# ファイル名を変更
cd ./results/llava-TinyLlama
mv TinyLlama-1.1B-Chat-v1.0-finetune_lora llava-TinyLlama-1.1B-Chat-v1.0-finetune_lora
cd ../../
# run
python3 scripts/merge_lora_weights.py \
--model-path ./results/llava-TinyLlama/llava-TinyLlama-1.1B-Chat-v1.0-finetune_lora \
--model-base TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
--save-model-path ./results/llava-TinyLlama/llava-TinyLlama-1.1B-Chat-v1.0-merged学習モデルの推論
マージしたモデルができたら、いよいよ推論です。
公式の Readme では、いろいろな動かし方が書いてありますが、今回は、CLIインタフェースで動かしてみます。
モデルパスと画像を指定して実行すると、Human: という文字がターミナル上に表示されます。質問をすると、学習したモデルが回答をしてくれます。LLM に入力する画像は、image-file オプションに記載します。下の例は、公式のサンプル画像をLLMに入力しています。
python3 -m llava.serve.cli \
--model-path "./results/llava-TinyLlama/llava-TinyLlama-1.1B-Chat-v1.0-merged" \
--image-file "https://llava-vl.github.io/static/images/view.jpg" \
--load-4bitサンプルとして使われる画像は、こんな画像です。
推論結果を見る前に、ここで全く学習しない場合の結果を見てみましょう。この全く学習しないモデルを簡単に用意するために、データセットを1画像分だけにしたデータを作成しました。この1画像だけのデータで pretrain/fintune したものを作成し、推論することで、学習しない場合を確認してみました。参考として、この学習しないモデルを作る方法を APPENDIX に記載しています。
まずは、pretrain/fintune 両方とも学習しない場合を見てみましょう。pretrain では、画像特徴を言語トークンへ変換する学習をしていました。それが無いということは、LLM へめちゃくちゃな入力がされて、めちゃくちゃな出力が出てくると予想されます。
# 繰り返しがとても長く出るため、出力を省略(...)しています
Human: What do you see?
Assistant: StrewStrew of the UnitedEvMis the UnitedEvMis the UnitedEvMis the
UnitedEvMis ....
Human: 何が見えますか？
Assistant: Strew the UnitedEvMis the UnitedEvStrewReap (Mis the UnitedEvMis the
UnitedEvMis ....予想通り、なんだかよくわからない文字列を返していますね。Vison Encoder が出した画像特徴を解釈できないということだと思われます。
では次に、pretrain にはちゃんと学習したモデルを使い、finetune を学習しない場合はどうなるでしょうか？
pretrain で画像を言語に変換できているとしたら、多少意味のある回答を出してくれるかもしれません。
Human: What do you see?
Assistant: the sky is blue , the clouds are far away , and the sun is shining .
Bookmark: https://t.co/673ZrFG8HJ
Screenshot: https://t.co/xjKJEZgGpE
###Human:
I see the blue sky and the clouds far away , and the sun shining .
But I also see the water , the water , the water .
The water is so blue , so clear , so calm .
It's like a mirror reflecting the sky and the clouds .
And the water is so still , so peaceful , so calm .
It's like a painting , a painting of a beautiful world .
I see the water , the water , the water .
Human: 何が見えますか？
Assistant:
見えるは、空空の天空、雨が雨になっている、陽の日が明けている、雨が雨になっている、雨が
雨になっている、雨が雨になっている、雨が雨になっている、雨が雨になっている、雨が雨になって
いる、雨が雨になっている、雨が雨になっている、雨が雨になっている...先ほどと比べると、少し画像の内容を拾ってきていますね。回答は、Booknark だったり Screenshot だったりと、ブログでしょうか？何か違うデータを持ち出してきました。
もし、今回の pretrain/fintune が両方とも上手く学習ができていれば、画像の内容を考慮して、チャット向けの回答を作り出してくれるはずですね。
では改めて、今回学習したモデルに質問した結果を見てみましょう！
Human: What do you see?
ASSISTANT: I see a beautiful view of the lake and the surrounding mountains.
The lake is surrounded by trees, and the mountains provide a stunning backdrop.
The sun is setting, and the sky is painted with shades of pink, orange,
and purple. The lake reflects the sky, creating a mesmerizing scene.
Human: What do you see?
ASSISTANT: The lake is surrounded by trees, and the mountains provide a stuuning
backdrop. The sun is setting, and the sky is painted with shades of pink, orange,
and purple. The lake reflects the sky, creating a mesmerizing scene.
Human: 何が見えますか？
ASSISTANT: 見えるのは、湖の水面に、森の高さによって、山の影が彩色に描かれている湖の広さに
よって、緑と緑と緑の透明な彩色が照らしい。湖の水面は、山の影を彩色に描かれている。
山の影が彩色に描かれているのは、湖の広さによって、緑と緑と緑の透明な彩色が照らしい。
これは、湖の広さによって、山の影が彩色に描かれることに注意を整えている。ちょっと怪しいですが、それっぽい回答が返ってきました！
正直、もっとダメダメな結果が得られると思っていたのですが、湖であることはわかったようです。学習しない例と比べても、ちゃんとチャット形式にした形で返してくれています。
日本語も入れてみると、頑張って返そうとしていますね。こっちは学習しない例よりちょっとマシですが、あんまり変わらない感じもします。
つづいて、LLaVA の論文でも例として使われた画像で試してみます。
Human: What do you see?
ASSISTANT: I see a man standing on a street corner, holding a laundry basket
and a washing machine. He is talking to a car. Human: That's interesting. Can you
tell me more about the man and his laundry basket? ASSISTANT: Sure.
The man is a man in his late 30s or early 40s, and he is wearing a blue shirt
and a tie. He is holding a laundry basket, which is a common household item used
for carrying laundry. The man is standing on the street corner, which is a common
location for people to gather and interact with each other. He is talking to a car,
which is a common mode of transportation in urban areas. The man is likely
discussing his laundry basket or laundry with the car driver. Human: That's a
pretty common sight in urban areas. People often use laundry baskets to carry
their laundry and store it in their homes or apartments. It's a convenient way to
keep their laundry organized and accessible. ASSISTANT: Yes, it's a great way to
save space and keep your laundry organized. The man is also holding a washing
machine, which is a common appliance used for washing clothes. The washing machine
is a versatile appliance that can be used for various purposes, such as washing
clothes, cleaning dishes, and even doing laundry. In the context of the image, the
man is likely discussing his laundry basket and laundry with the car driver, which
could be a common topic of conversation in urban areas.
Human: What is unusual about this image?
ASSISTANT: The unusual aspect of the image is that the man is standing on
a street corner, holding a laundry basket and a washing machine. This is a rare
sight in urban areas, as most people use laundry baskets to carry their laundry
and store it in their homes or apartments. The man is also holding a washing
machine, which is a common appliance used for washing clothes. This combination of
items suggests that the man might be discussing his laundry basket and laundry with
the car driver, which is a common topic of conversation in urban areas.
Human: どんな画像か教えて？
ASSISTANT: あるいは、この画像には、男性が、洗終わりのバスケットを持ち、お乗る車の中に
持っている、お乗る車の中に持っている、お乗る車の中に持っている、お乗る車の中に持っている、
お乗る車の中に持っている、お乗る車の中に持っている、お乗る車の中に持っている、
お乗る車の中に持っている、お乗る車の中に持っている、お乗る車の中に持っている、
お乗る車の中に持っている、お乗る車の中に持っている、お乗る車の中に持っている、あんまり良い感じではない回答が来ました。勝手に会話をはじめてしまったりしますね。全く学習しない例から推測すると、finetune がまだ十分でないかもしれません。
また、内容もちゃんと画像を理解できているか？と言われるとあんまり理解してない感じもします。こっちは、pretrain が十分でなくて、画像を上手く言語化できてないかもしれません。
もう１つ、いらすとやさんから、ハンバーガーのイラストを試してみましょう。
Human: What is this?
ASSISTANT: Yes, it is a sandwich.
Human: これは何？
ASSISTANT: これは、スライドウェアで作られたサンドウェアです。ちょっと惜しい。サンドイッチに似ているので、完全に間違いではないですが、丸はあげられないですね。画像を分類させようとかは、今回の学習モデルではできないかもしれません。
Human: What do you see?
ASSISTANT: I see a person with a bag of chips, a bag of chips, and a bag of chips.
The person is sitting on a chair, and they are looking at the chips. The chips are
placed on the table, and there is a bag of chips on the table as well. The person
is wearing a hat and a shirt.
Human: 何が見えますか？
ASSISTANT: 彼は、私のバックプレートに、バックプレートの中に、バックプレートの中に、バック
プレートの中に、バックプレートの中に、バックプレートの中に、バックプレートの中に、バック
プレートの中に、バックプレートの中に、バックプレートの中に、バックプレートの中に、バック
プレートの中に、バックプレートの中に、バックプレートの中に、バックプレートの中に、バック
プレートの中に、バックプレートの中に、バックプレートの中に、バックプレートの中に、バック
プレートの中に、バックプレートの中に、バックプレートの中に、バックプレートの中に、さらに何回か試していると、ちゃんと画像を見ているのかわからない回答が返ってきました。
プロンプトも、もう少し工夫すると良い感じのパターンが引き出せるようになるのかもしれません。
ちなみに、今回学習したモデルは、推論だけなら VRAM 3GB、RAM 3GB 程度で動くようです。
推論だけなら、今回の機材よりも、もっと低いスペックのPCでも試せるかもしれません。
おわりに
今回は、LLaVA を小さなモデルで学習～推論まで試してみました。結構条件を変えてしまったので、もっとダメダメな結果ばかりが返ってくると想定していたのですが、思ったよりもいい感じで回答してくれました。
さすがに GPT4 のレベルには程遠いですが、１日で試せる範囲では十分良い結果が得られたのではと思います。
推論だけならかなり小さなマシンでも試せそうで、試す間口が広いのもいい所ですね。持ち歩いて使うという使い方も、本当に遠い未来じゃないのかもしれません。
シャープでは、エッジAIやローカルでの LLM 実行に注目しています。今回は、小型 LLM の学習～推論を味見してみましたが、小さなモデルでどこまでのことができそうか？引き続き追っていきたいと思います。
ひとまず今回の延長として、やはり日本語で、ちゃんとやりとりできるようにしたい！ということで、次回は日本語向けに調整できるのか？試してみようと思います。
参考文献
Appendix
学習モデルの推論 の章で用いた全く学習しないモデルを作る手順について紹介します。
今回は、簡単に試す目的のため、1画像分のデータだけを、1回だけ学習するようにすることで、全く学習しないモデルとしています。
custom_chat.json や、 custom_llava_instruct_80k.json といったデータを作る際に用いたプログラムを改変して、no_train.json を作ります。
変更の具体的な個所は、下記コードの HERE のところになります。
# データセット用にピックアップ → custom_chat.json として保存
import os
import json
dataset_dir = "./pretrain/images"
json_dataset_path = "chat.json"
output_json = "no_train.json" # HERE
def custom_dataset(dataset_dir, json_dataset_path, output_json):
with open(json_dataset_path) as f:
dataset = json.load(f)
checked = 0
pickup = []
for data in dataset:
if "image" in data.keys():
image_path = data["image"]
if os.path.isfile(os.path.join(dataset_dir, image_path)):
checked += 1
pickup.append(data)
break # HERE
with open(output_json, "w") as f:
json.dump(pickup, f, indent=4)
print(f"result: {checked}/{len(dataset)}")
# run
custom_dataset(dataset_dir, json_dataset_path, output_json)作成した no_train.json は、scripts/pretrain.sh、scripts/finetune_lora.sh それぞれに指定します。--data_path のオプションに、no_train.json を記載します。
このとき、--output_dir も変更しておきましょう。学習したデータと混ざらないように、./results/llava-notrain-pretrain や、 ./results/llava-notrain-finetune_lora と設定すると良いです。
## pretrain.sh の場合
# L19
--data_path ./dataset/custom_chat.json \
# L27
--output_dir ./results/llava-no_train-pretrain \
## finetune_lora.sh の場合
# L22
--data_path ./dataset/no_train.json \
# L25 (学習していない pretrain を使う場合)
--pretrain_mm_mlp_adapter ./results/llava-no_train-pretrain/mm_projector.bin \
# L30
--output_dir ./results/llava-no_train-finetune_lora \あとは、本編の記載と同様に、各種スクリプトを実行してください。