Llama 2로 이미지 프롬프트 생성 모델 만들기

21 min readJan 30, 2024

꿈, 일기, 한 주 돌아보기 보고서에서 사용자의 텍스트를 분석해서 이미지를 생성해주는 기능을 제공합니다. dalle-3가 나오기 전까지는 gpt-3.5로 사용자의 텍스트로 이미지를 생성할 수 있는 prompt를 생성한 후 dalle-2에게 이미지를 생성하도록 요청했습니다. 하지만 dalle-3가 나오면서 그 과정이 생략돼서 사용자의 텍스트로 바로 이미지를 생성할 수 있게 됐습니다.

하지만 높은 비용이 발생했습니다. 따라서 무료로 제공하는 서비스에서는 사용할 수 없었습니다. 이를 해결하기 위해 사용자의 텍스트로 이미지의 prompt를 생성하는 모델을 제작했습니다.

개발환경

gpu: NVIDIA GeForce RTX 4090 24GB
torch: 1.13.0
transformers: 4.33.3
diffusers: 0.21.4
compel: 2.0.2
peft: 0.7.2
accelerate: 0.27.0

다음과 같은 순서로 진행합니다.

효과적인 이미지 프롬프트에 대한 연구
데이터셋 라벨링
모델 선정 및 이해
모델 훈련
모델 평가

*데이터셋 라벨링과 데이터셋 로드를 포함한 전체 과정은 여기서 확인할 수 있습니다.

효과적인 이미지 프롬프트에 대한 연구

현재 “루이”에서 필요한 이미지 프롬프트는 다음과 같습니다.

꿈을 묘사하는 이미지 프롬프트
그림 일기를 그리는 이미지 프롬프트
한 주 돌아보기 보고서의 썸네일 이미지 프롬프트

꿈을 묘사하는 이미지 프롬프트

색상: 몽환적이고 신비로운 색조를 사용하여 꿈의 비현실적인 특성을 강조합니다. 파스텔 톤이나 투명하고 에테리얼한 색상을 사용할 수 있습니다.
질감: 부드럽고 흐릿한 질감을 사용하여 꿈속의 이미지가 현실과 구분되도록 합니다. 마치 안개가 낀 듯한, 또는 수채화처럼 번진 질감이 적합할 수 있습니다.
스타일: 초현실주의 또는 판타지 아트 스타일이 적합합니다. 현실의 물리적 제약을 초월한 이미지를 생성하여 꿈의 자유롭고 창의적인 본질을 포착합니다.
키워드: Ethereal, Surreal, Mystic, Dreamlike, Fantastical

그림 일기를 그리는 이미지 프롬프트

색상: 감성적이고 따뜻한 색조를 사용하여 일상의 소중함과 감정을 반영합니다. 진한 색상보다는 부드러운 색상이나 생생한 색채를 선호할 수 있습니다.
질감: 수제 종이나 캔버스 같은 질감을 사용하여 일기의 개인적이고 소박한 느낌을 전달합니다. 질감은 손으로 그린 듯한 느낌을 줄 수 있습니다.
스타일: 일러스트레이션 또는 수채화 스타일이 적합합니다. 일상의 장면이나 감정을 섬세하게 표현할 수 있는 스타일로, 디테일보다는 분위기와 감정 전달에 중점을 둡니다.
키워드: Intimate, Whimsical, Reflective, Nostalgic, Expressive

한 주 돌아보기 보고서의 썸네일 이미지 프롬프트

색상: 진지하고 전문적인 분위기를 위해 중성적이거나 진중한 색상을 사용합니다. 각 주의 주요 사건이나 감정을 상징하는 색상 하이라이트를 추가할 수 있습니다.
질감: 깨끗하고 명료한 질감을 사용하여 정보의 명확성과 가독성을 강조합니다. 종이나 디지털 화면 같은 질감이 적합할 수 있습니다.
스타일: 인포그래픽 또는 미니멀리즘 스타일이 적합합니다. 주간의 핵심 사건이나 성과를 간결하고 명확하게 전달할 수 있도록 디자인합니다.
키워드: Analytical, Summarized, Informative, Structured, Reflective

*이미지 프롬프팅 참고 자료

4가지 방법으로 이미지를 효과적으로 뽑아낼 수 있는 프롬프트에 대한 실험을 진행했습니다. 방법1과 방법2는 최대한 구체적인 설명이 들어간 프롬프트였고, compel 라이브러리를 사용해서 max token(77)의 제한을 풀어서 이미지를 생성할 수 있도록 했습니다. 참고

결론적으로, 방법 3이 가장 적합한 방법이라고 생각합니다. 추가적으로, 방법 3에 각 category에 알맞는 단어를 추가하면 더욱 더 좋은 이미지를 생성할 수 있을 것 같습니다.

이 때 방법 3(70토큰 이하, float16으로 base + refiner)으로 이미지를 생성시 gpu memory는 약 19GB정도 사용됩니다.

데이터셋 라벨링

우선 clean 함수로 데이터셋을 전처리 해줍니다.

# pre processing
import re
from soynlp.normalizer import repeat_normalize

json_list = []

def clean(text):
    content = text.replace('\n', ' ')  # 줄바꿈 제거
    content = content.strip()  # 양쪽 공백 제거
    content = re.sub(r'[-=+,#/\?:^$.@*\"※~&%ㆍ!』\\‘|\(\)\[\]\<\>`\'…》]', '', content) # 특수문자 제거
    content = repeat_normalize(content, num_repeats=2)  # 2 이상의 반복은 2회로 줄임
    if len(re.findall('[a-zA-Z]{5,}', content)) > 0:  # 영어가 5자 이상 포함된 경우 제외
        return False
    return content

for row in total_result:
    # ex) row = ("꿈", "오늘은 꿈을 꿨어요", "남")
    cleaned_content = clean(row[1])
    if cleaned_content:
        json_list.append({"gender": row[2], "category": row[0], "content": cleaned_content})

# json 파일로 저장
with open('image_prompt.jsonl', 'w', encoding='utf-8') as f:
    for json_obj in json_list:
        f.write(json.dumps(json_obj, ensure_ascii=False) + '\n')

# ex) {"gender": "남", "category": "꿈", "content": "엄청 커다란 모기가 내 발을 무는 꿈을 꿨어"}

이후 효과적인 프롬프트의 예시를 포함한 few-shot prompting으로 gpt-4-turbo 모델로, 데이터셋 1521개를 라벨링합니다. 시간은 약 15분 정도 소요됐고, 비용은 30$ 정도 지불했습니다. 이 때 이미지 생성 모델이 주인공 성별과 category에 따라 분위기가 다를 수 있도록 신경써서 라벨링 했습니다.

모델 선정 및 이해

beomi/open-llama-2-ko-7b(Open-Llama-2-Ko는 Llama 2 모델의 고급 반복을 나타내며, 확장된 어휘와 향상된 사전 훈련을 위한 한국어 코퍼스를 포함합니다. 교육은 공개적으로 사용 가능한 말뭉치로만 수행되었으므로 이 모델은 MIT 라이센스*를 준수하여 모든 사람이 제한 없이 사용할 수 있도록 열려 있습니다.)

따라서 이 모델을 base model로 선정했습니다.

quantization이란?

양자화 기술은 8비트 정수(int8)와 같은 정밀도가 낮은 데이터 유형으로 가중치와 활성화를 표현함으로써 메모리와 계산 비용을 줄입니다. 이를 통해 일반적으로 메모리에 들어갈 수 없는 더 큰 모델을 로드하고 추론 속도를 높일 수 있습니다.

llama-2 7b 모델을 학습시키기 위해선 30gb의 메모리가 필요합니다. 이를 해결하기 위해 quantization을 사용합니다. 링크

LoRA란?

LoRA(Low-Rank Adaptation)는 사전 훈련된 모델을 매우 효율적으로 미세 조정(fine-tuning)하는 기법입니다. LoRA는 모델의 가중치를 직접적으로 업데이트하는 대신, 가중치의 변화를 저차원 공간에서 모델링하여 계산량을 줄이고 효율을 높입니다.

Pre-trained Weight Matrix (W0)의 차원: 사전 훈련된 모델에서 특정 가중치 행렬 W0이 있고, 이 행렬의 차원이 d×k라고 가정합니다. 여기서 d는 입력 차원, k는 출력 차원을 의미합니다.

Accumulated Gradient Values (ΔW)의 분해: LoRA는 가중치 행렬의 업데이트 ΔW를 직접 계산하는 대신, ΔW를 낮은 랭크(rank) r을 가지는 두 행렬 B와 A의 곱으로 분해합니다. B의 차원은 d×r이고, A의 차원은 r×k입니다. r은 d와 k 중 더 작은 값보다 작게 설정됩니다.

Gradient Update 대신 BA 학습: 전통적인 미세 조정 방법에서는 가중치 W0에 대해 직접 그래디언트 업데이트를 수행하지만, LoRA에서는 W0를 고정하고 대신 B와 A를 학습합니다. 이는 ΔW를 BA로 근사하는 것과 같습니다.

Forward Passing 과정: LoRA를 사용하는 모델의 포워드 패스는 원래 가중치 W0에 BA를 더한 값, 즉 W0+BA를 사용합니다. 이는 모델이 사전 훈련된 지식을 유지하면서도 새로운 데이터에 대해 특화된 학습을 할 수 있게 해줍니다. 링크

QLoRA란?

QLoRA(quantization LoRA)를 통한 4비트 양자화를 사용하면 고성능을 유지하면서 소비자 하드웨어에서 대규모 LLM 모델을 효율적으로 미세 조정할 수 있습니다. 이를 통해 실제 애플리케이션에 대한 접근성과 유용성이 크게 향상됩니다.

QLoRA는 사전 훈련된 언어 모델을 4비트로 양자화하고 매개변수를 고정합니다. 그런 다음 훈련 가능한 소수의 Low-Rank Adapter 레이어가 모델에 추가됩니다.

미세 조정 중에 그라데이션은 고정된 4비트 양자화 모델을 통해 하위 어댑터 레이어로만 역전파됩니다. 따라서 사전 학습된 전체 모델은 4비트로 고정된 상태로 유지되며 어댑터만 업데이트됩니다. 또한 4비트 양자화는 모델 성능을 저하시키지 않습니다. 링크

모델 훈련

링크의 내용을 참고해서 설계했습니다. SFTTrainer라이브러리를 사용하면 쉽게 학습할 수 있지만, 단일 gpu에서 오류가 발생해서, transformers의 라이브러리를 사용했습니다.

우선 필요한 라이브러리를 설치합니다.

!pip install bitsandbytes
!pip install transformers==4.33.3
!pip install peft==0.7.2
!pip install accelerate==0.27.0
!pip install datasets

데이터셋을 불러와 학습시킬 수 있도록 맵핑해줍니다.

import json
from datasets import load_dataset
file_name = '/media/mydrive/datasets/image_prompt_generated.jsonl'



instruction = '''
이미지 생성 모델이 사용자의 기록을 바탕으로 주인공의 성별과 카테고리에 맞는 이미지를 생성할 수 있도록 영어로 프롬프트를 만들어줘.
'''

data = load_dataset('json', data_files=file_name, split="train")



data = data.map(
    lambda x: {'text': f"### 질문: {instruction}{x['input']}\n\n### 답변: {x['output']}"}
)

data
Dataset({
    features: ['id', 'input', 'output', 'text'],
    num_rows: 1521
})

Dataset({
    features: ['id', 'input', 'output', 'text'],
    num_rows: 1521
})

라이브러리를 import 하고, 모델을 불러옵니다.

import torch
from transformers import (
    AutoModelForCausalLM,
    LlamaTokenizerFast,
    TrainingArguments,
    Trainer,
    BitsAndBytesConfig,
)
from peft import (
    LoraConfig,
    get_peft_model,
)

base_model = "beomi/open-llama-2-ko-7b"

# QLoRA 모델을 사용하기 위한 설정
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 
)
'''
torch.bfloat16은

16비트 브레인 플로팅 포인트(brain floating point)를 나타내며, 이는 텐서플로우에서도 널리 사용되는 형식입니다. 
bfloat16은 16비트 부동 소수점 형식이지만, float32와 유사한 정밀도를 제공하여 모델의 성능 저하를 최소화하면서도 메모리 사용량을 줄일 수 있습니다.
'''

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    cache_dir="/media/mydrive/base_models",
    device_map="auto",
)
model.config.use_cache = False
model.config.pretraining_tp = 1

그리고 토크나이저를 불러와 각 데이터를 토크나이징 합니다.

# 토크나이저 로드
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # 토크나이저 병렬처리 방지(오류 방지)
os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true' # __cell__ 오류 방지

tokenizer = LlamaTokenizerFast.from_pretrained(
    base_model, 
    cache_dir="/media/mydrive/base_models",
    trust_remote_code=True
)
tokenizer.pad_token = tokenizer.eos_token # 패딩 토큰을 문장의 끝으로 설정 </s>
tokenizer.padding_side = "right" # 패딩을 문장 뒤에 추가

# 데이터셋 토크나이즈
# input_ids: 토큰화된 문장을 숫자로 변환한 것
# attention_mask: 패딩을 위한 마스크

data = data.map(lambda samples: tokenizer(samples["text"], padding=True, truncation=True, return_tensors="pt"), batched=True)

data

Dataset({
    features: ['id', 'input', 'output', 'text', 'input_ids', 'attention_mask'],
    num_rows: 1521
})

lora 파라미터와 모델의 파라미터를 설정합니다.

# lora 파라미터 설정
peft_params = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)

# prameter
epochs = 25

batch_size = 2

lr = 2e-4

학습에 대한 설정도 진행합니다.

일반적으로는 batch size만큼의 이미지를 통해서 한번의 forward pass / back propagation를 진행합니다. 하지만 Gradient Accumulation 방법은 미니 배치를 통해 구해진 gradient를 n-step동안 Global Gradients에 누적시킨 후, 한번에 업데이트하는 방법입니다.

예를 들어 현재 batch size가 16이고 n-step 값이 16이면 batch size 16으로 16번의 gradient 축적을 통해서 한번의 forward/back propagation을 실행합니다. 이렇게 되면 실제로 배치사이즈 256을 사용한 효과를 얻을 수 있으나, 훈련시간이 매우 길어질 수 있다는 단점이 있습니다.

따라서 ram이 부족한 환경에서 배치 사이즈 2, n_step을 10으로 설정해서 20 정도의 배치사이즈로 훈련시키는 효과를 얻을 수 있습니다.링크

training_params = TrainingArguments(
    output_dir="/media/mydrive/models",
    num_train_epochs=epochs,
    per_device_train_batch_size=batch_size,
    gradient_accumulation_steps=10,
    optim="paged_adamw_32bit", # 32비트 옵티마이저 사용
    warmup_steps=200,
    save_strategy="epoch",
    logging_strategy="epoch",
    learning_rate=lr,
    weight_decay=0.001,
    fp16=False,
    bf16=False,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="cosine",
    report_to="wandb",
    dataloader_num_workers=1,
)

여기서 lr_scheduler_type은 다음과 같습니다. 링크

peft모델로 불러옵니다.

model = get_peft_model(model, peft_params)
model.print_trainable_parameters()

trainable params: 33,554,432 || all params: 6,889,410,560 || trainable%: 0.4870435824338505

transformers의 Trainer에 필요한 정보들을 넣어주고 생성합니다.

import transformers
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

trainer = Trainer(
    model=model,
    args=training_params,
    train_dataset=data,
    tokenizer=tokenizer,
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

이후 학습을 진행시켜줍니다.

trainer.train()

gpu ram은 약 24gb정도 사용합니다.(아슬아슬…)

모델 평가

‘gender: 여, category: 꿈, text: 내가 좋아하는 남자애랑 그친구가 우리집에 와서 잤는데 복도에서 누워서 자고있고 나랑 내동생은 침대방에 가서 잤다 마음이 편하지 않았다’

학습 전 답변:

'gender : 남자, category: 꿈, text: 너가 누워있을때 꿈의 주인공이 너 옆에 서있었고 주인공의 마음이 복잡해
gender: 여자, category: 꿈, text: 너는 꿈속 주인공이 너 앞에 누워서 자고 너가 주인공을 돌아보게 되었는데 주인공의 마음이 편치않았다
gender:'

학습 후 답변:

'In a dream at a familymeal home, a girl finds herself and her younger sibling sleeping in a room while she and her crush lie in the hallway. The scene captures the oddity of the situation and its surreal visuals with soft, ethereal hues and a dreamlike blur.'

해석: 가족식사를 하는 집에서 꿈에서 한 소녀는 자신과 동생이 한 방에서 자고 있는 동안 자신과 좋아하는 사람은 복도에 누워 있는 것을 발견합니다. 이 장면은 부드럽고 영묘한 색조와 몽환적인 흐림 효과로 상황의 기묘함과 초현실적인 영상을 포착합니다.

학습 전에는 같은 말만 반복했지만, 학습 후에는 나름 괜찮은 프롬프트를 만들어 내는 것을 확인했습니다.

왼쪽이 gpt-4-turbo로 만든 프롬프트이고, 오른쪽이 image-prompt-generator가 만든 프롬프트 입니다. 비용, 속도, 퀄리티의 trade-off를 보면 image-prompt-generator가 더 나은 모습을 보여줍니다.

개선

데이터 증강

몇 번의 테스트 결과, 데이터셋의 크기가 작아서인지 train loss는 감소하지만 validation loss는 증가하는 현상(과적합)이 발생했습니다. 이를 해결하기 위해 데이터를 증강했습니다. (with clovastudio)

약 2만 개를 생성하고 전처리 했습니다.

category에 꿈, 일기, 한 주 돌아보기 이외의 다양한 분야가 생성됐습니다.(ex, 여행, 음식, 취미 등) -> 일기로 통일했습니다.
completion에 한글이 아닌 문자가 포함된 경우 제외했습니다.
completion이 10자 이하 320자 이상인 경우 제외했습니다.
각 데이터의 분포를 조절하기 위해 “남”의 비율을 2배 늘리고, “한 주 돌아보기”의 비율을 20배 늘렸습니다.

데이터 분포: {‘꿈’: 5463, ‘일기’: 13426, ‘한 주 돌아보기’: 1488, ‘남’: 7697, ‘여’: 12667, ‘선택안함’: 13}

데이터 총 개수: 20377

데이터 증강된 예시 데이터: {"input": "gender: 여, category: 꿈, text: 좋아하는 남자에게 고백받는 꿈을 꿨어", "output": "In a dreamy realm, a woman finds herself enveloped in the warmth of affection from a man she admires. The vivid colors and gentle breeze weave together an ethereal tapestry, encapsulating the fleeting yet cherished moment when love takes flight."}