見出し画像

RoBERTaで記事分類のAIモデルを構築

記事分類のAIモデル構築をChatGPTと相談しながら進めてみました。

環境

  • Google Colab Pro (GPUのA100を使用)

  • TensorFlow (Google Colabに元から入っているまま)

  • Transformers (4.29.2)

  • xlm-roberta-base

  • トレーニングデータを独自で事前に用意

Google ColabでインストールしたパッケージはTransformersだけです。

!pip install transformers

事前に作成していた学習データをGoogle Driveに格納しておきます。
今回作成したデータとしては全部で5,219件あり、そこで記事1に分類するものが3,760件、それ以外が1,459件となっています。
また、データの形式はJSON Lineで作成していました。

ある程度チューニングした後のコード

引き続き試しているので完成版ではないのですが、ある程度チューニングした後のコードです。

current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = LOGS_DIR + "/run-" + current_time

# データのロードと前処理
def load_data(filename):
    with open(filename, "r") as f:
        lines = f.readlines()

    texts = []
    labels = []
    for line in lines:
        data = json.loads(line)
        texts.append(data["title"] + " " + data["text"])  # タイトルとテキストデータを連結
        labels.append(1 if data["label"] == "記事1" else 0)  # ラベルデータ(記事1の記事:1、それ以外の記事:0)

    return texts, labels

initial_learning_rate = 0.0001
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate, decay_steps=10000, decay_rate=0.95, staircase=True
)

# TensorBoardコールバックの設定
tensorboard_callback = TensorBoard(log_dir=log_dir)

# データの読み込み
texts, labels = load_data(TRAIN_JSON)
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)

# class weightsを計算
total = len(train_labels)
music_count = sum(train_labels) * 0.8 # "記事1" label is 1
not_music_count = total - music_count
class_weight = {0: 1 / not_music_count, 1: 1 / music_count}

# トークナイザーの初期化
tokenizer = RobertaTokenizerFast.from_pretrained(PRETRAINED_MODEL)

# トークン化
train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=512)
val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=512)

# TensorFlowのデータセットを作成
train_dataset = tf.data.Dataset.from_tensor_slices(
    (dict(train_encodings), train_labels)
).shuffle(len(train_texts)).batch(32)
val_dataset = tf.data.Dataset.from_tensor_slices((dict(val_encodings), val_labels)).batch(32)

# ドロップアウト率を設定
dropout_rate = 0.5

# モデルの設定を初期化
config = RobertaConfig.from_pretrained(
    PRETRAINED_MODEL,
    num_labels=2,
    hidden_dropout_prob=dropout_rate,  # ここでドロップアウト率を指定
    attention_probs_dropout_prob=dropout_rate  # アテンションのドロップアウトも設定
)

# モデルの初期化
model = TFRobertaForSequenceClassification.from_pretrained(
    PRETRAINED_MODEL,
    config=config
)

# モデルのコンパイル
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"]
)

# EarlyStoppingコールバックの設定
early_stopping_callback = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',  # 監視対象の指標(ここでは検証データに対する損失)
    patience=2,  # 指標が改善されないエポックを許容する数
    restore_best_weights=True  # 最適な重みを復元するかどうか
)

# 学習開始
model.fit(
    train_dataset, epochs=100,
    validation_data=val_dataset,
    callbacks=[tensorboard_callback, early_stopping_callback],
    class_weight=class_weight
)


# モデルの保存
model.save_pretrained(MODEL_DIR)

TensorBoardで確認もしていました。

%load_ext tensorboard
%tensorboard --logdir /gdrive/MyDrive/data/article/classification/logs

ChatGPTとの相談

基本的には以下のような形式で聞いていました。
最初は丁寧に聞いていたのですが、途中から面倒になりこのような聞き方になってます。

コード

(上記コードを貼り付け)

結果

(GoogleColabの実行結果を貼り付け、以下のようなもの)

所感と提案を下さい

そうすると次のような回答をくれるので、コードを変更してみたり、分からないことはさらに聞いたりします。

考慮したこと

結果的に考慮したことは以下のようなものがありました。また、一つ変更する毎に試して変化を確認しながら進めました。

TensorBoardの利用

学習状況を視覚的に確認したかったので、TensorBoardで確認できるようにしました。

current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
log_dir = LOGS_DIR + "/run-" + current_time

tensorboard_callback = TensorBoard(log_dir=log_dir)

早期停止の組み込み

どれだけ学習を回せばいいか最初から設定するのは難しいので、一定回数学習が進まない場合は停止するように組み込みました。

early_stopping_callback = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',  # 監視対象の指標(ここでは検証データに対する損失)
    patience=2,  # 指標が改善されないエポックを許容する数
    restore_best_weights=True  # 最適な重みを復元するかどうか
)

トレーニングデータの改善

データの正則化などは実施していたのですが、現状でもまだ件数のバランスが悪いですが、最初はもっとバランスが悪かったのでデータを追加しました。

トレーニングデータのシャッフル

ラベルでソートされてしまっているような状態だったのでデータをシャッフルするようにしました。

train_dataset = tf.data.Dataset.from_tensor_slices(
    (dict(train_encodings), train_labels)
).shuffle(len(train_texts)).batch(32)

検証データの利用

フレームワークで勝手に検証データも分割して使ってくれるのかと勘違いしていたので、検証データを利用するようにしました。

texts, labels = load_data(TRAIN_JSON)
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)

val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=512)

val_dataset = tf.data.Dataset.from_tensor_slices((dict(val_encodings), val_labels)).batch(4)

ラベルの重み付け

トレーニングデータのラベル毎の件数のバランスが悪いため重みをつけました。

# class weightsを計算
total = len(train_labels)
music_count = sum(train_labels) * 0.8 # "記事1" label is 1
not_music_count = total - music_count
class_weight = {0: 1 / not_music_count, 1: 1 / music_count}

学習率の調整

他のパラメータとの関係もあるので、常に調整し続けています。また、スケジューリングを持ちいて学習が進む毎に学習率を変化するようにしました。

initial_learning_rate = 0.0001
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate, decay_steps=10000, decay_rate=0.95, staircase=True
)

バッチ数の調整

GPUのメモリ不足のエラーを気にしながら調整しました。

train_dataset = tf.data.Dataset.from_tensor_slices(
    (dict(train_encodings), train_labels)
).shuffle(len(train_texts)).batch(32)
val_dataset = tf.data.Dataset.from_tensor_slices((dict(val_encodings), val_labels)).batch(32)

ドロップアウトの追加

過学習と思われる状態が続いたので、ドロップアウトを追加しました。

dropout_rate = 0.5

config = RobertaConfig.from_pretrained(
    PRETRAINED_MODEL,
    num_labels=2,
    hidden_dropout_prob=dropout_rate,  # ここでドロップアウト率を指定
    attention_probs_dropout_prob=dropout_rate  # アテンションのドロップアウトも設定
)

損失関数の選択

今回の分類は2値分類になっているため、"binary_crossentropy"も勧められました(ChatGPTより)が、学習済みモデルの出力を確認して"SparseCategoricalCrossentropy"を利用するようにしました。

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"]
)

ちなみに、学習済みモデルの出力の確認方法や、それによってどの損失関数が望ましいかなどをChatGPTに聞きながら進めました。

さいごに

調整の繰り返しのため完成したとはまだ考えていないのですが、それっぽくなってきたのでひとまずここに残しておこうと思います。

事前学習モデルにBERTを使うことも考えましたが、今回はRoBERTaを利用してみました。ただ、途中で一度BERTに切り替えてみましたが学習評価上の数値はそれほど変わってなかったのでRoBERTaに戻しました。

そしてGoogle ColabのA100はいいですね。1 epochが60秒くらいで進んでくれて快適です。その代わりコンピューティングユニット(13.08/h)がガンガン減っていくので、後半はT4(1.96/h)にして試してました。
でも、(起動マシンに寄ると思いますが)効率は悪くなってたような・・・

ランタイム、コンピューティングユニット、処理時間
A100, 13.08/h, 約60s/epoch
T4, 1.96/h, 約600s/epoch

消費は約1/6なのに、時間は10倍・・・

好きに使える環境があるといいですよね。


この記事が気に入ったらサポートをしてみませんか?