Hugging Faceを使った自作モデル作成の事始め
はじめに
Hugging FaceにはBertForClassificationなどの既存のモデルが定義されていますが，自作のアーキテクチャでモデルを定義・学習したいこともあります．本記事では，ある程度ちゃんとした形式で実装をするための事始めについて書きます．
本記事はpython3.11.0，および以下のモジュールで検証しました．
torch==2.1.0 transformers==4.34.0
概要
（HuggingFaceらしく）モデルを定義するためにはConfigとModelの2つのclassが必須です．
Configには，モデルを初期化および訓練するための情報を定義します．例えば，隠れ層のサイズや分類層の出力サイズ，もしくはdropoutの確率などです．
Modelは，Configさえ受け取れば初期化ができてforwardが通せるように設計します．例えば，configに入力次元数と隠れ次元数が書いてあれば，それを使ってLinearを初期化できます．
コードベースでは，ざっくり以下のようにモデルを初期化できればいいです．（ConfigとModelはそれぞれclassとして定義されていると思ってください）
config = Config() model = Model(config)
以下，色々な例を出しながら説明します．
最小限の自作モデルの実装例
まず手始めに，線形層を2層重ねただけのどうでもいいモデルを作ってみます． 入力は適当な5次元ベクトルで，隠れ次元数を3，出力は1次元のベクトルだとします（回帰タスクのつもりで）．
Configの定義
Configを自作するときには，PretrainedConfigを継承する形で定義します．これにより，既に定義された変数が使えたり，セーブ・ロードがお手軽にできたり，transformersの既存Configとインターフェイスが合うなどの利点があります．
今回は2層積むので，入力サイズ，隠れサイズ，出力サイズの3つの情報があれば初期化できます．数値は適当です．
__init__()では必ずデフォルト引数を設定してください．
from transformers import PretrainedConfig class Config(PretrainedConfig): def __init__( self, input_size: int = 5, hidden_size: int = 3, output_size: int = 1, **kwards ): super().__init__(**kwards) self.input_size = input_size self.hidden_size = hidden_size self.output_size = output_size
(**kwardsはなんやねん，という方がいるかもしれません．PretrainedConfigにはどのモデルにも共通で使える（であろう）パラメータが既に多数定義されています．そのようなものについては，わざわざ自分でself.XXX =と書く必要はなく，**kwardsを経由してオーバーライドするだけで済みます．**kwardsの構文的な意味は，余った引数をまとめて辞書型で受け取ることです．その後，super().__init__(**kwards)で親クラスに（辞書を展開して）渡している感じです．構文的な意味についてはHuggin Faceではなくpythonの話なので，詳しくは検索してください．)
これで，Configのインスタンスを作ることができるようになりました．
config = Config() print(config) ''' printの結果： Config { "hidden_size": 3, "input_size": 5, "output_size": 1, "transformers_version": "4.34.0" } '''
もちろん，初期化のときに引数を与えれば好きな値を設定できます．
config = Config(
input_size=999,
hidden_size=1000,
output_size=999
)
print(config)
'''
printの結果：
Config {
"hidden_size": 1000,
"input_size": 999,
"output_size": 999,
"transformers_version": "4.34.0"
}
'''
また，PretrainedConfigはセーブ・ロードのメソッドを持っています．よく見るsave_pretrained()とfrom_pretrained()です．自作したクラスはPretrainedConfigを継承しているので，当然これらのメソッドを呼び出せます．
from_pretrained()はクラスメソッドであることに注意してください．
config = Config() outdir = 'sample' config.save_pretrained(outdir) loaded_config = Config.from_pretrained(outdir) print(config) print(loaded_config)
もしsave_pretrained()を呼んだ時にXXX.__init__() missing XX required positional arguments:のようなエラーが出る場合は，Configクラスの__init__()でデフォルト引数が設定されているか確認してください．
sampleディレクトリは勝手に作成されます．その中にconfig.jsonが保存されているので，中身を確認してみてください．
Configの定義は以上です．
モデル定義
モデルはPreTrainedModelを継承する形で定義します．冒頭でも述べたように，モデルは初期化の時にConfigを受け取るので，Configに書かれた情報を使いながら初期化するように__init__()を書きます．
from transformers import PreTrainedModel import torch.nn as nn class Model(PreTrainedModel): config_class = Config def __init__(self, config): super().__init__(config) self.config = config self.linear1 = nn.Linear( self.config.input_size, self.config.hidden_size ) self.linear2 = nn.Linear( self.config.hidden_size, self.config.output_size )
気をつけるべき点は以下です．
config_class = Configのように，クラス変数としてconfig_classにConfigを定義したクラスを設定してください．from_pretrained()でモデルを読み込む際に必ず必要になりますdef __init__(self, config):のように，初期化時にConfigのインスタンスを渡せるようにしてください．．super().__init__(config)のように，親クラスの初期化のためにもconfigのインスタンスを渡してください．
こうすることで，以下のようにしてモデルを定義できます．次元数がConfigの通りになっていることを確認してください．
config = Config() model = Model(config) print(model) ''' printの結果 Model( (linear1): Linear(in_features=5, out_features=3, bias=True) (linear2): Linear(in_features=3, out_features=1, bias=True) ) '''
PreTrainedModelにセーブ・ロードのメソッドが定義されているので，これを継承した自作モデルはセーブ・ロードのメソッドを呼び出せます．Configの時と同じく，save_pretrained()とfrom_pretrained()です．
モデルをセーブする時にはConfigのsave_pretrained()も呼ばれるため，モデルをセーブすればConfigも勝手にセーブされます．
config = Config()
model = Model(config)
outdir = 'sample'
model.save_pretrained(outdir)
loaded_model = Model.from_pretrained(outdir)
もう一度sampleディレクトリを確認してください．モデルの重みを保存したファイルであるpytorch_model.binが保存されています．
あとはforward()を書けば完成です．forwardは学習ループの書き方にもよるのでかなり自由に書けますが，Hugging Faceらしい実装を目指すのであれば，
- forward()の引数にラベルを渡すことができて，ラベルを渡された時にはlossまで計算して返すようにする
- 返り値のために専用のOutputクラス（dataclass）を別に用意し，返り値はそれに入れて返す
をやっておくと良いと思います．
from transformers import PreTrainedModel import torch import torch.nn as nn from dataclasses import dataclass # 返り値用のクラス @dataclass class ModelOutput: loss: torch.Tensor = None logits: torch.Tensor = None class Model(PreTrainedModel): config_class = Config def __init__(self, config): super().__init__(config) self.config = config self.linear1 = nn.Linear( self.config.input_size, self.config.hidden_size ) self.linear2 = nn.Linear( self.config.hidden_size, self.config.output_size ) def forward( self, input, labels=None ): logits = self.linear1(input) logits = self.linear2(logits) loss = None if labels is not None: # labels=に何か与えられていたら損失計算 loss_fn = torch.nn.MSELoss() loss = loss_fn( logits, labels ) return ModelOutput( loss=loss, logits=logits )
こうすることで，次のように学習の1サイクルを回すことができます．
config = Config() model = Model(config) inputs = torch.rand((5)) # 何かの入力 labels = torch.tensor([0.5]) # 何かのラベル output = model(inputs, labels) print(output.loss) # tensor(0.9413, grad_fn=<MseLossBackward0>) output.loss.backward()
ここまで来たら，後はDataLoader周りを書けばいつも通り学習できます（この記事ではそこまでは踏み込みません）．
既存モデルを組み合わせる場合の実装例（マルチタスク文分類を例に）
次に，既存モデル（例えばBERTなど）を組み合わせる場合の実装をしてみます．ここではマルチタスクの文分類をするような気持ちで，BERTなどのエンコーダから得られた表現を二種類の分類層に独立に入力するようなモデルを作りましょう．
要件は，
- 1つ目のタスクは2値分類．ラベルはA0, A1とする
- 2つ目のタスクは3値分類，ラベルはB0, B1, B2とする
- ドロップアウト層を分類層の直前に入れる
こととします．
Configの定義
モデルがConfigに書かれた情報だけで初期化できるようにConfigを定義しないといけないので，始めにどのような情報があれば初期化できるかを考えます．今回は，
- エンコーダとして使うモデルのid
- それぞれの分類タスクにおけるクラスのidとラベルのマッピング情報
- dropoutの確率
があれば良さそうです．（実際には，モデルを先に組んでみて，その後で必要な情報を含むようにConfigを定義するのも良いと思います）
実装としてはPretrainedConfigを継承して定義しますが，ここでPretrainedConfigがどのようなパラメータを既に持っているのか見てみます．以下のconfiguration_utils.pyのPretrainedConfigのソースがありますので，ざっと確認してみてください．
この中に既に使えそうなものがあれば自分で定義する必要はありません．今回はid2labelsとlabel2idがidとラベルのマッピング情報を表すので，これを片方の分類層のために使いましょう．それで，もう片方の分類層の情報は自分で定義しましょう．dropoutも自分で定義しましょう．
from transformers import PretrainedConfig class Config(PretrainedConfig): def __init__( self, model_id=None, id2label_second=None, label2id_second=None, dropout=0.1, **kwards ): super().__init__(**kwards) self.model_id = model_id self.id2label_second = id2label_second self.label2id_second = label2id_second self.dropout = dropout
初期化するときは，id2labelとlabel2idも忘れずに，計6つの引数を与えます． セーブとロードもできることを確認します．
config = Config(
model_id='bert-base-uncased',
id2label={0:'A0', 1:'A1'},
label2id={'A0':0, 'A1': 1},
id2label_second={0:'B0', 1:'B1', 2:'B2'},
label2id_second={'B0':0, 'B1':1, 'B2':2},
dropout=0.1
)
config.save_pretrained('out_test')
loaded_config = Config.from_pretrained('out_test')
print(config)
'''
printの結果：
Config {
"dropout": 0.1,
"id2label": {
"0": "A0",
"1": "A1"
},
"id2label_second": {
"0": "B0",
"1": "B1",
"2": "B2"
},
"label2id": {
"A0": 0,
"A1": 1
},
"label2id_second": {
"B0": 0,
"B1": 1,
"B2": 2
},
"model_id": "bert-base-uncased",
"transformers_version": "4.34.0"
}
'''
コンフィグの定義は以上です．
モデルの定義
次にモデルを定義します．Configの内容をうまく使いながら初期化していきます．
from transformers import PreTrainedModel, AutoModel import torch.nn as nn class Model(PreTrainedModel): config_class = Config def __init__(self, config): super().__init__(config) self.config = config self.transformer = AutoModel.from_pretrained( config.model_id ) hidden_size = self.transformer.config.hidden_size self.classifier = nn.Linear( hidden_size, len(config.id2label) ) self.classifier2 = nn.Linear( hidden_size, len(config.id2label_second) ) self.dropout = nn.Dropout(config.dropout)
モデルの初期化ができるか確認します．ついでにセーブとロードができるか確認します．
config = Config(
model_id='bert-base-uncased',
id2label={0:'A0', 1:'A1'},
label2id={'A0':0, 'A1': 1},
id2label_second={0:'B0', 1:'B1', 2:'B2'},
label2id_second={'B0':0, 'B1':1, 'B2':2},
dropout=0.1
)
model = Model(config)
model.save_pretrained('out_test')
loaded_model = Config.from_pretrained('out_test')
print(model)
'''
printの結果：
Model(
(transformer): BertModel(
.....中略.....
)
(classifier): Linear(in_features=768, out_features=2, bias=True)
(classifier2): Linear(in_features=768, out_features=3, bias=True)
(dropout): Dropout(p=0.1, inplace=False)
'''
最後にforwardを定義します．分類層が2つあるためラベルも2種類与えられるようにする点に注意します．また，BERTの入力を与える必要があるので，tokenizerの返り値をそのまま入力できるようにします．今回はかなりサボって，input_idsとattention_maskだけ受け取るようにしましょう．
from transformers import PreTrainedModel import torch import torch.nn as nn from dataclasses import dataclass class Model(PreTrainedModel): config_class = Config def __init__(self, config): super().__init__(config) self.config = config self.transformer = AutoModel.from_pretrained( config.model_id ) hidden_size = self.transformer.config.hidden_size self.classifier = nn.Linear( hidden_size, len(config.id2label) ) self.classifier2 = nn.Linear( hidden_size, len(config.id2label_second) ) self.dropout = nn.Dropout(config.dropout) def forward( self, input_ids, attention_mask, labels_first=None, labels_second=None ): logits = self.transformer( input_ids=input_ids, attention_mask=attention_mask ).pooler_output logits = self.dropout(logits) logits_first = self.classifier(logits) logits_second = self.classifier2(logits) loss = None if labels_first is not None and labels_second is not None: loss_fn = nn.CrossEntropyLoss() num_labels = len(self.config.id2label) loss_first = loss_fn( logits_first.view(-1, num_labels), labels_first.view(-1) ) num_labels = len(self.config.id2label_second) loss_second = loss_fn( logits_second.view(-1, num_labels), labels_second.view(-1) ) loss = loss_first + loss_second return ModelOutput( loss=loss, logits_first=logits_first, logits_second=logits_second )
最後に，適当な入力でforwardが通ることを確かめます．（本当は小さいデータで過学習するか，なども検証した方がいいですが割愛）
from transformers import AutoTokenizer config = Config( model_id='bert-base-uncased', id2label={0:'A0', 1:'A1'}, label2id={'A0':0, 'A1': 1}, id2label_second={0:'B0', 1:'B1', 2:'B2'}, label2id_second={'B0':0, 'B1':1, 'B2':2}, dropout=0.1 ) model = Model(config) tokenizer = AutoTokenizer.from_pretrained(config.model_id) # 何かの文 sentences = ['This is a sample sentence.', 'This is also sentence.'] encode = tokenizer(sentences, return_tensors='pt', padding='max_length') del encode['token_type_ids'] # 何かのラベル encode['labels_first'] = torch.tensor([0, 1]) encode['labels_second'] = torch.tensor([2, 1]) loss = model(**encode).loss print(loss) # tensor(2.1259, grad_fn=<AddBackward0>) loss.backward()
おわり
Hugging Faceを用いて自作モデルを作る方法について書きました．単純な線形層を積むだけのモデルから始まり，実際にBERTを用いてマルチタスクで文分類するモデルを実装してみました．
