ここでは、フィードフォワードネットワークでファジータイリングアクティベーションを使用するトランスフォーマーをトレーニングします。これを言語モデルとして使用し、Tiny Shakespeareデータセットでトレーニングしてデモンストレーションを行います
。ただし、これはおそらくFTAにとって理想的なタスクではなく、連続変数を含むデータのモデル化にはFTAの方が適していると考えています。
21import copy
22
23import torch
24import torch.nn as nn
25
26from labml import experiment
27from labml.configs import option
28from labml_helpers.module import Module
29from labml_nn.activations.fta import FTA
30from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
31from labml_nn.transformers import MultiHeadAttention, TransformerLayer
32from labml_nn.transformers.utils import subsequent_mask
d_model
トークン埋め込みに含まれる機能の数d_ff
は FFN の隠れレイヤーにあるフィーチャの数ですactivation
FTA アクティベーションモジュールですかdropout
は隠れ層のドロップアウト確率です40 def __init__(self, d_model: int, d_ff: int,
41 activation: FTA,
42 dropout: float = 0.1):
49 super().__init__()
重みとバイアスでパラメータ化されたレイヤー 1
51 self.layer1 = nn.Linear(d_model, d_ff)
重みとバイアスでパラメータ化されたレイヤー 2
53 self.layer2 = nn.Linear(d_ff * activation.expansion_factor, d_model)
隠しレイヤーのドロップアウト
55 self.dropout = nn.Dropout(dropout)
アクティベーション機能
57 self.activation = activation
59 def forward(self, x: torch.Tensor):
61 x = self.activation(self.layer1(x))
ドロップアウトを適用
63 x = self.dropout(x)
65 return self.layer2(x)
68class AutoregressiveTransformer(Module):
n_tokens
ボキャブラリ内のトークンの数ですd_model
は埋め込みサイズn_layers
変圧器層の数ですlayer
レイヤーです。n_layers
変圧器にはこれのコピーを使います76 def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: TransformerLayer):
83 super().__init__()
n_layers
層付き変圧器
85 self.transformer_layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])
トークン埋め込みレイヤー
88 self.emb = nn.Embedding(n_tokens, d_model)
読み出し層
90 self.readout = nn.Linear(d_model, n_tokens)
マスクは最初の呼び出しで初期化されます
93 self.mask = None
x
形状の入力トークンです [seq_len, batch_size]
95 def forward(self, x: torch.Tensor):
自動回帰マスクの作成
100 if self.mask is None or self.mask.size(0) != len(x):
次にマスクすると、トークンがマスクされ、将来のトークンが見えなくなります
102 self.mask = subsequent_mask(len(x)).to(x.device)
トークンの埋め込みを入手
105 x = self.emb(x)
トランスエンコーダー
107 for layer in self.transformer_layers:
108 x = layer(x=x, mask=self.mask)
ロジットを取得
110 x = self.readout(x)
結果を返す
113 return x, None
116class Configs(NLPAutoRegressionConfigs):
モデル
125 model: AutoregressiveTransformer
レイヤー数
128 n_layers: int = 4
そしてディープノーム用
131 deep_norm_alpha: float
132 deep_norm_beta: float
注目されているヘッドの数
135 n_heads: int = 4
埋め込みサイズ
137 d_model: int = 256
各アテンションヘッドのサイズ
139 d_k: int = 16
フィードフォワードレイヤーサイズ
141 d_ff: int = 256
自由貿易協定
144 fta_lower_limit: float = -1.
145 fta_upper_limit: float = +1.
146 fta_delta: float = 0.2
147 fta_eta: float = 0.05
150@option(Configs.model)
151def _model(c: Configs):
FTA アクティベーションモジュールを作成
157 fta = FTA(c.fta_lower_limit, c.fta_upper_limit, c.fta_delta, c.fta_eta)
トランスフォーマーを作成します。TransformerLayer
MultiHeadAttention
再利用して実装します
161 m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
162 TransformerLayer(d_model=c.d_model,
163 feed_forward=FeedForwardFTA(d_model=c.d_model,
164 d_ff=c.d_ff,
165 activation=fta,
166 dropout=0.1),
167 self_attn=MultiHeadAttention(c.n_heads, c.d_model,
168 dropout_prob=0.0),
169 dropout_prob=0.0))
デバイスに移動
172 return m.to(c.device)
175def main():
実験を作成
180 experiment.create(name="fta", writers={'screen', 'labml'})
コンフィグの作成
182 conf = Configs()
オーバーライド設定
184 experiment.configs(conf, {
キャラクターレベルのトークナイザーを使う
186 'tokenizer': 'character',
プロンプトセパレータが空白
188 'prompt_separator': '',
サンプリングの開始プロンプト
190 'prompt': 'It is ',
タイニー・シェイクスピア・データセットを使う
192 'text': 'tiny_shakespeare',
コンテキストサイズを次の値にしてください
195 'seq_len': 256,
32 エポックのトレーニング
197 'epochs': 32,
バッチサイズ
199 'batch_size': 16,
エポックごとにトレーニングと検証を切り替える
201 'inner_iterations': 10,
ウォームアップなしの Adam オプティマイザー
204 'optimizer.optimizer': 'Adam',
205 'optimizer.learning_rate': 3e-4,
206 })
保存および読み込み用のモデルを設定します
209 experiment.add_pytorch_models({'model': conf.model})
実験を始める
212 with experiment.start():
トレーニングを実行
214 conf.run()
218if __name__ == '__main__':
219 main()