Rustのproc-macroで自動微分してGPT-2を学習してみた
はじめに
深層学習フレームワークにおいては、自動微分によってロスの勾配を求め、勾配降下法によってモデルのパラメータを最適化することで学習が進みます。
例えばPyTorchの場合、torch.Tensorを使うことで計算グラフが記録されていき、この情報から誤差逆伝播法によって勾配が計算されます[1]。
ところで、自動微分はよりコード生成的なアプローチ、即ち、関数の定義を受け取り、その導関数の定義を自動生成するという方法で実装することもできます。
このアプローチが機能するかどうかは言語のメタプログラミングのやり易さに依存すると言って良いでしょう。
もしCやFortranでこのアプローチを取る場合、Tapenadeのようにソースコード変換を行う外部ツールとして実装する必要があります。
他方、メタプログラミング機能が充実している言語においては、コード生成アプローチによってシームレスな体験を実現できます。
例えばJulia言語ではZygote.jlというsource-to-sourceの自動微分ライブラリが広く使われており、これを用いると言語標準の型を使って普通に書いたプログラムをそのまま微分することができます。
julia> f(x) = 5x^3 + 2x + 1 # 関数を定義
f (generic function with 1 method)
julia> f(10)
5021
julia> f'(10) # 'をつけることで、f(x)の微分を計算する関数が自動でコンパイルされ実行される
1502.0
julia> @code_llvm f'(10) # 生成された関数のllvm irを見ると5.0*(3x)*x+2.0に相当するコードが生成されているのがわかる
; Function Signature: (::Zygote.var"#97#98"{typeof(Main.f)})(Int64)
; @ /Users/xxxxx/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:159 within `#97`
define double @"julia_#97_6689"(i64 signext %"x::Int64") #0 {
top:
%0 = mul i64 %"x::Int64", 3
%1 = mul i64 %0, %"x::Int64"
%2 = sitofp i64 %1 to double
%3 = fmul double %2, 5.000000e+00
%4 = fadd double %3, 2.000000e+00
ret double %4
}
さて、最近の自分は勉強がてらちょっとしたツールをRustで作ることが多いです。
そして、Rustもproc-macroによる強力な静的メタプログラミングをサポートしています。
となると、Rustのproc-macroでも同様のsource-to-sourceの自動微分が実現できるのかが気になってきます。
ということで 実際に自動微分っぽいことをするproc-macroを作成しました。
そして自動微分ができるようになったからには、何かしら深層学習モデルを学習してみたいところです。
今回は、GPT-2（アーキテクチャの小さなdecoder-only transformer）をndarrayを使って実装し、実際に学習を行ってみました。
自動微分を実装する
Zygote.jlの動作
実装方針を考えるためにも、まずはZygote.jlの動作原理について見て行きます。
これについては論文や寺崎敏志氏によるわかりやすい日本語解説記事が大変参考になりますが、ここでは非常にざっくりと説明します。
主役となるのはrruleという関数で、これは「関数f」と「引数x」を受け取り、「f(x)の値」と「yからy * f'(x)を計算する関数（これをpullbackとよぶ[2]）」のペアを返します[3]。
ちなみにxがベクトルの場合はy * f'(x)はvector–Jacobian productになります。
さて、ここで
function f(x)
y = g(x)
z = h(y)
return z
end
のような関数があり、gとhに対してはrruleが定義されているとしましょう。するとfに対するrruleは
function rrule(f, x)
y, pullback_1 = rrule(g, x) # pullback_1 = y_bar -> y_bar * g'(x) と同義
z, pullback_2 = rrule(h, y) # pullback_2 = z_bar -> z_bar * h'(y) と同義
return z, function(z_bar)
y_bar = pullback_2(z_bar) # = z_bar * h'(y)
x_bar = pullback_1(y_bar) # = z_bar * h'(y) * g'(x) = z_bar * h'(g(x)) * g'(x) = z_bar * f'(x)
return x_bar
end
end
のように、内部で呼び出される関数gとhのrruleの組み合わせとして表現できます。
このようにして微分を求めたい関数のrruleを求めていき、最後に_, pullback = rrule(f, x); pullback(1.0)を計算すれば、知りたかった勾配f'(x)が求まるというわけです。
Juliaでは関数はそれぞれが固有の型を持っており、なおかつ多重定義された関数において使用される定義は全ての引数の型を見て決まります（多重ディスパッチ）。そのため、単にrruleという関数を多重定義するだけでこのような挙動が実現できます。
適切にrruleが定義できるのであれば、多重定義によってどんな型、どんな関数でも微分することができるというのはこのアプローチの興味深いポイントです。
ところで、先にあげた例は単純化のために関数を呼び出す度にその結果を変数に格納していました。
実際のZygote.jlはg(h(g(x, x)), x) + h(x)のような複雑でネストした関数呼び出しでも動作します。
また、ifやforなどの制御構造を含むコードも正しく変換することができます。
これが実現できるのは、ソースコードを直接変換しているのではなく、コンパイラをフックして静的単一代入(SSA)形式の中間表現を獲得し、これに対して変換を行なっているからです。
実装方針
さて、ここまでを見るとproc-macroで自動微分するにあたっての問題点が見えてきます。
まず、Rustのproc-macroはあくまでもトークン列からトークン列への変換です。コンパイラをフックして中間表現を引っ張ってきてそれを変換するようなことはできません。
また、型システムやディスパッチの仕組みの違いから、関数呼び出しf(x)を対応するrruleの呼び出しrrule(f, x)に置き換える部分も実装が難しそうです。
色々考えた結果、ユーザーが何も考えずに書いたコードがそのまま微分できる、という体験を実現することは諦め、以下のように妥協することにしました。
- 自動微分したい関数の内部では、値とpullbackを返す関数を明示的に呼び出す
-
y = f(x)と書くのではなく、最初からy, _ = rrule(f, x)に相当するコードをユーザーが書く
-
- pullbackは必ず決まった名前の引数で受け取り（今回は
__とした）、トークン列を見るだけでそれがrruleの呼び出しであることがわかるようにする-
let (y, __) = rrule_of_f(x);という感じ
-
この他、微分するのは第一引数だけ（複数引数を扱いたいときはタプルにする）、参照は扱わない、など問題を単純化しマクロで扱いやすくしました。
実際の動作
具体的に、今回作成したマクロad!がどのような動作をするのかを解説します。動作イメージは以下のような形です。
ad!マクロは微分したいクロージャーを受け取り、
ad!(|(x, y)| {
let (x, __) = foo(x, 1); // 自動微分の対象になるのは第一引数のみ
let (y, __) = bar((x, y)); // 複数の引数はタプルにまとめる
let (z, __) = baz(y);
z
})
以下のようなコードに変換します。
|(x, y)| {
let (x, __pb1) = foo(x, 1);
let (y, __pb2) = bar((x, y));
let (z, __pb3) = baz(y);
(
z,
move |mut z| {
let mut y = __pb3(z);
let mut (x, y) = __pb2(y);
let mut x = __pb1(x);
(x, y)
}
)
}
構文としてのifはサポートしていませんが、
let (y, __) = (if y == 0 {foo} else {bar})(x)
のように書くことはできます。
今回はtransformerを実装したかったので、レイヤーを順番に作用させるためのループ処理は必要になります。
ループについての構文は少し不恰好ですが以下のようにしました。
ループに入る前に、ループ内のpullbackを受け取るためのmut変数__を定義します。ループ内での関数呼び出しは一つまでとし、mut変数と__のペアを更新していくように書きます。
ad!(|x| {
let (mut y, __) = foo(x);
let mut __;
for i in 0..10 {
(y, __) = bar(y, i);
}
y
})
これは以下のように展開されます。
|x| {
let (mut y, __pb1) = foo(x);
let mut __pb2_stack = Vec::new();
for i in 0..10 {
let __pb2;
(y, __pb2) = bar(y, i);
__pb2_stack.push(__pb2);
}
(
y,
move |mut y| {
while let Some(__pb2) = __pb2_stack.pop() {
y = __pb2(y);
}
let x = __pb1(y);
x
}
)
}
具体的な実装はここにあります。仕様を簡略化したおかげで500行くらいでいけました。
この仕様を見て、マクロのユーザーがここまで気を遣ってコードを書く必要があるなら、もういっそDSLにしてしまったほうが良いのでは？ という感想を抱く人も多いと思います。
まあ、自分もそう思います...
とはいえ、今回はproc-macroでsource-to-source自動微分ができるかどうかを模索する、というのが目的だったので、ad!マクロには合法なRustのコードを渡すように仕様を決めました。
だって、そうじゃないとsource-to-source感が薄くなっちゃうかなって...
GPT-2を実装する
何はともあれ自動微分ができるようになったので、果たしてこれが本当に使い物になるか確かめるためにも、GPT-2を実装して学習してみましょう。
幸いなことに、GPT-2を自分で実装するチュートリアルはインターネットにたくさんあり、資料には苦労しません。
特に参照したのはllm.cです。ここにはC言語のみで書かれたGPT-2の学習コードがあります。
アテンション、レイヤーノルム、ソフトマックス、MLPといった基本的な構成モジュールについての勾配計算が手書きで定義されており、最後にそれを組み合わせてロスの勾配を求めています。
今回はなるべく自動微分を活用してコードを書きたかったので、多次元配列に対する基本的な演算にのみ微分を定義し、これらのモジュールについての勾配はad!で自動的に導出することにします。
Rustでの多次元配列は素直にndarrayクレートを使用します。ここも自作しているといつまでもゴールできないので...
演算としては、四則演算（＋ブロードキャスト）、（バッチ）行列積、log、exp、gelu、pow2、sqrt、sum、split、permute、reshape、advanced indexing（x[[3,1,4],:]みたいなやつ）があればOKです。なので、まずはこのあたりのpullbackを頑張って実装します。
あとはこれらを組み合わせて必要なモジュールを作って行きます。cargo expandを使って一番複雑なアテンションの微分計算がどのように展開されたかを見てみるとこんな感じです。
これを見ると、確かに自動で微分を導出してくれているな...と少し良い気分になれます。
attention全体
let f = |(x, p)| {
let (((w_qkv, w_out, b_out), p), __0) = AttentionList::get(p, idx);
let (b_out, __1) = reshape(b_out, [1, emb]);
let (x, __2) = reshape(x, [bat * seq, emb]);
let (qkv, __3) = mm((x, w_qkv));
let ((q, kv), __4) = split_at(qkv, Axis(1), emb);
let ((k, v), __5) = split_at(kv, Axis(1), emb);
let (q, __6) = reshape(q, [bat, seq, num_heads, emb / num_heads]);
let (q, __7) = swap_axes(q, Axis(1), Axis(2));
let (q, __8) = reshape(q, [bat * num_heads, seq, emb / num_heads]);
let (k, __9) = reshape(k, [bat, seq, num_heads, emb / num_heads]);
let (k, __10) = swap_axes(k, Axis(1), Axis(2));
let (k, __11) = reshape(k, [bat * num_heads, seq, emb / num_heads]);
let (v, __12) = reshape(v, [bat, seq, num_heads, emb / num_heads]);
let (v, __13) = swap_axes(v, Axis(1), Axis(2));
let (v, __14) = reshape(v, [bat * num_heads, seq, emb / num_heads]);
let (kt, __15) = swap_axes(k, Axis(1), Axis(2));
let (a, __16) = bmm((q, kt));
let (a, __17) = div_num(a, scaling);
let (a, __18) = reshape(a, [bat * num_heads, seq, seq]);
let (m, __19) = attach(mask_causal)(());
let (a, __20) = add((m, a));
let (a, __21) = softmax(a, Axis(2));
let (m, __22) = attach(mask_dropout_attn)(());
let (a, __23) = mul((a, m));
let (a, __24) = bmm((a, v));
let (a, __25) = reshape(a, [bat, num_heads, seq, emb / num_heads]);
let (a, __26) = swap_axes(a, Axis(1), Axis(2));
let (a, __27) = reshape(a, [bat * seq, emb]);
let (a, __28) = mm((a, w_out));
let (a, __29) = add((a, b_out));
let (a, __30) = reshape(a, [bat, seq, emb]);
let (m, __31) = attach(mask_dropout_resid)(());
let (a, __32) = mul((a, m));
(
(a, p),
move |(mut a, mut p)| {
let (mut a, mut m) = __32(a);
let () = __31(m);
let mut a = __30(a);
let (mut a, mut b_out) = __29(a);
let (mut a, mut w_out) = __28(a);
let mut a = __27(a);
let mut a = __26(a);
let mut a = __25(a);
let (mut a, mut v) = __24(a);
let (mut a, mut m) = __23(a);
let () = __22(m);
let mut a = __21(a);
let (mut m, mut a) = __20(a);
let () = __19(m);
let mut a = __18(a);
let mut a = __17(a);
let (mut q, mut kt) = __16(a);
let mut k = __15(kt);
let mut v = __14(v);
let mut v = __13(v);
let mut v = __12(v);
let mut k = __11(k);
let mut k = __10(k);
let mut k = __9(k);
let mut q = __8(q);
let mut q = __7(q);
let mut q = __6(q);
let mut kv = __5((k, v));
let mut qkv = __4((q, kv));
let (mut x, mut w_qkv) = __3(qkv);
let mut x = __2(x);
let mut b_out = __1(b_out);
let mut p = __0(((w_qkv, w_out, b_out), p));
(x, p)
},
)
}
最後にモジュールを組み合わせてtransformer全体を組み立てます。こちらもcargo expandの結果の一部をお見せするとこのような感じです。
transformer全体
|(embedding, pos_emb, mut attn, mut ln, mut ffn, linear)| {
let (e, __0) = advanced_indexing_along(embedding, Axis(0), &data_in);
let (e, __1) = reshape(e, [batch_size, seq, emb]);
let ((p, rest), __2) = split_at(pos_emb, Axis(0), seq);
let ((), __3) = detach(rest);
let (p, __4) = reshape(p, [1, seq, emb]);
let (x, __5) = add((e, p));
let (emb_dropout, __6) = attach(emb_dropout)(());
let (mut x, __7) = mul((x, emb_dropout));
let mut __8 = Vec::new();
for l in 0..num_layers {
let mut f_inner = |(x, attn, ln, ffn)| {
let ((x, res), __0) = fork(x);
let ((x, ln), __1) = LayerNormList::call((x, ln), 2 * l);
let ((x, attn), __2) = AttentionList::call((x, attn), l, dropout);
let (x, __3) = add((x, res));
let ((x, res), __4) = fork(x);
let ((x, ln), __5) = LayerNormList::call((x, ln), 2 * l + 1);
let ((x, ffn), __6) = FeedForwardList::call((x, ffn), l, dropout);
let (x, __7) = add((x, res));
(
(x, attn, ln, ffn),
move |(mut x, mut attn, mut ln, mut ffn)| {
let (mut x, mut res) = __7(x);
let (mut x, mut ffn) = __6((x, ffn));
let (mut x, mut ln) = __5((x, ln));
let mut x = __4((x, res));
let (mut x, mut res) = __3(x);
let (mut x, mut attn) = __2((x, attn));
let (mut x, mut ln) = __1((x, ln));
let mut x = __0((x, res));
(x, attn, ln, ffn)
},
)
};
__8.push({
let __;
((x, attn, ln, ffn), __) = f_inner((x, attn, ln, ffn));
__
});
}
let ((x, ln), __9) = LayerNormList::call((x, ln), 2 * num_layers);
let (x, __10) = reshape(x, [batch_size * seq, emb]);
let (logit, __11) = mm((x, linear));
let (prob, __12) = softmax(logit, Axis(1));
let (prob, __13) = reshape(prob, [batch_size, seq, vocab]);
let ((), __14) = LayerNormList::detach(ln);
let ((), __15) = AttentionList::detach(attn);
let ((), __16) = FeedForwardList::detach(ffn);
(
prob,
move |mut prob| {
let mut ffn = __16(());
let mut attn = __15(());
let mut ln = __14(());
let mut prob = __13(prob);
let mut logit = __12(prob);
let (mut x, mut linear) = __11(logit);
let mut x = __10(x);
let (mut x, mut ln) = __9((x, ln));
while let Some(__8) = __8.pop() {
(x, attn, ln, ffn) = __8((x, attn, ln, ffn));
}
let (mut x, mut emb_dropout) = __7(x);
let () = __6(emb_dropout);
let (mut e, mut p) = __5(x);
let mut p = __4(p);
let mut rest = __3(());
let mut pos_emb = __2((p, rest));
let mut e = __1(e);
let mut embedding = __0(e);
(embedding, pos_emb, attn, ln, ffn, linear)
},
)
}
このようにして、最終的にはdecoder-only transformerのクロスエントロピー損失の勾配が計算できるようになりました。
GPT-2を学習する
ここまでくれば、あとは実際にモデルを学習するだけです。
参考にしたllm.cでは本物のGPT-2と同じくWebTextコーパスで学習を回すコードが書かれていますが、このコーパスは巨大すぎてちょっとした検証には辛いです。特に今回の実装はCPU上でしか動作しないので尚更...
実際、llm.cのサンプルは学習済みモデルを読み込んで継続学習するという形になっています。
動作確認としても、できればゼロから学習を回してみたいところです。しかし、巨大なコーパスは扱えない。
ということで良いものがないか色々探した（ChatGPTでDeepResearchした）ところ、TinyStoriesというデータセットが良さそうと分かりました。
TinyStoriesはGPT-3.5/4で生成された、3-4歳児でも分かる単語からなる短編小説の合成データセットで、小規模なモデルの学習に向いています。
GitHubで検索するとこのデータセットを使って小さな言語モデルを学習してみた具体例を色々とみることができます。
特にtanaydesai/plutoがコンパクトに良く書かれており、出現頻度の低いトークンを除外して語彙数を絞って学習する部分など、参考になりました。
ちなみに、学習データをトークナイズする部分については普通にpythonで書きました。
結果
モデルを学習しつつ、学習50ステップごとに文章を生成させてみました。
"One day,"までを初期入力として与え、そこから貪欲法で50トークン生成しています。
最初の50ステップではほんの少し文章を生成したあとでひたすら改行を連打していましたが、300ステップくらいまでくるとぱっと見流暢な英語を生成するようになります。
2000ステップくらいまで進むと、文章として成立しているだけでなく、文の間の意味の繋がりもしっかりとしてきています。
50: One day, there was a was a little girl.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
100: One day, a little girl named Lily named Lily named Lily. She was very very happy to play. She was a big park. She was very happy. She was very happy to the park.\n\n\n\n\n\n\n\n\n\n\n\n
150: One day, a little girl named Lily was very excited. She was very happy to play with her mom. She was very excited. She was very excited to the park. She was very excited. She was very excited. She was very happy and said,
200: One day, a little girl named Lily was very excited to play with her mom. She was very excited. She saw a big mommy and the park. She saw a big, "I'm it.\n\n\n\n\n\n\n\n\n\n
250: One day, a little girl named Lily was very excited to the park. She was so happy and the park. She was so happy and he was so happy. She was so happy to the park. She was so happy. She was so happy to the
300: One day, a little girl named Lily was very excited to play with her friends. She was very happy and she was very happy. She was so excited to the park. She was so excited to the park. She was so excited to the park. She
350: One day, a little girl named Lily was very happy. She was very happy and she was very happy. She was very happy and she was very happy. She was very happy to the park.\n\nOne day, she was so happy and she was
400: One day, a little girl named Lily was very happy. She was very happy and she was so excited to the park. She was so happy to the park. She was so happy to the park.\n\nOne day, she was so happy, she
450: One day, a little girl named Lily went to play with her mommy. She was very happy and she had a big tree. She wanted to play with her mommy.\n\nOne day, Lily saw a big tree. She saw a big tree
500: One day, a little girl named Lily went to the park. She was very excited to the park. She was very excited to the park. She was very excited to the park and she was so excited.\n\nOne day, she saw a big,
550: One day, a little girl named Lily went to the park. She was very excited to the park. She was very excited to the park. She was so excited to the park.\n\nOne day, she saw a little girl was so excited to the
600: One day, a little girl named Lily went to the park. She was very excited to the park. She was very excited and saw a big tree. She was very excited and wanted to play with her mom.\n\nLily was very excited to the
650: One day, a little girl named Lily went to the park. She was very excited and she had a big and she had a big smile. She was very happy and she wanted to play with her mom.\n\nOne day, she saw a big,
700: One day, a little girl named Lily was very excited to the park. She was so excited to the park with her mom. She was so excited and she saw a big, she was so she wanted to go on the park.\n\nLily was
750: One day, a little girl named Lily went to the park. She was very excited to the park. She saw a big tree. She wanted to play with her mom.\n\n"I want to play with me?" Lily said.\n\n"I
800: One day, a little girl named Lucy was walking in the park. She was playing in the park and she saw a big tree. She was very excited and she was very excited. She was so excited and she was so excited to the little girl.\n
850: One day, a little girl named Lily went to the park. She was playing with her mom and she had a big, but she was very excited. She was very excited to her mom said, "Mommy, I have a big, but I have
900: One day, a little girl named Lily went to the park. She was playing with her mommy and she had a big red dress. She wanted to play with her mommy. She was so excited to go to the store and she could see what she
950: One day, a little girl named Lucy was walking in the park. She was walking in the park with her mommy and saw a big, and a big, a big, a big tree. She was a big, and she was a big, and
1000: One day, a little girl named Lily went to the park. She saw a big tree. She was very excited and wanted to go to the park. She saw a big, but she was scared. She was so excited to the little girl was so she
1050: One day, a little girl named Lily went to the park. She saw a big tree. She wanted to play with her mom.\n\nLily asked her mom if she could go to the park. She said, "I want to play with me
1100: One day, a little girl named Lily went to the park. She saw a big tree with a big tree. She was very excited and wanted to go to the park.\n\n"Mommy, Lily, I'm going to the park. I'm
1150: One day, a little girl named Lily went to the park. She saw a big, red ball. She wanted to go on the park. She asked her mom, "What is a big, but I can I have a big box."\n\nL
1200: One day, a little girl named Lily went to the park with her mommy. She was very excited to play with her mommy. She was very excited and she wanted to play with her mommy.\n\nLily's mommy said, "
1250: One day, a little girl named Lily went to the park. She saw a big tree. She was very excited. She wanted to go to the park. She asked her mom if she could go to the store.\n\nLily said, "I
1300: One day, a little girl named Lily went to the park. She saw a big tree. She was very happy. She wanted to go to the park. She asked her mom if she could go to the park.\n\nLily was very excited.
1350: One day, a little girl named Lily went to the park with her mom. She saw a big tree. She was very excited. She wanted to play with her mom.\n\n"Mom, Lily, Lily, Lily. I want to play with me
1400: One day, a little girl named Lily went to the park. She saw a big tree and a big, red dog. She was very excited. She wanted to play with her mom.\n\n"Mom, Lily, Lily, Lily. I can I
1450: One day, a little girl named Lily went to the park. She saw a big tree with a big, red ball. Lily was very happy and wanted to play with it.\n\nLily's mom said, "Mom, I have a big,
1500: One day, a little girl named Lily went to the park with her mommy. She saw a big, red ball with a big, red ball. Lily was very happy and wanted to play with it.\n\nLily's mommy said, "
1550: One day, a little girl named Lily went to the park with her mommy. She saw a big, red ball that was very pretty. She wanted to play with her mommy, but she didn't know what was.\n\nLily asked her
1600: One day, a little girl named Lily went to the park. She saw a big tree with a big, red ball. She was very excited to see what it was.\n\nLily was so excited to see the tree. She ran to the tree
1650: One day, a little boy named Tim went to the park with his mom. He saw a big slide. He wanted to play with it. He asked his mom if he could not find it.\n\nTim said, "I'm sorry for you to
1700: One day, a little girl named Lily went to the park with her mommy. She saw a big, red ball. It was a big and shiny and shiny. Lily was very happy.\n\n"Look, Lily, Lily, I want to play
1750: One day, a little girl named Lily went to the park with her mom. She saw a big tree with a big tree. She was very excited and wanted to see what was inside.\n\n"What is a big tree?" asked her mom.\n
1800: One day, a little girl named Lily went to the park. She saw a big tree with a big tree. She was very curious and wanted to play with it.\n\n"Mom, Lily, I can't find it?" Lily asked.\n\n
1850: One day, a little girl named Lily went to the park. She saw a big tree. She was very excited and wanted to go on the tree.\n\n"What is that?" she asked.\n\n"I'm going to go to the tree
1900: One day, a little girl named Lily went to the park with her mommy. She saw a big, red ball with a big, red ball. Lily was very happy. She wanted to play with it, but she didn't know what it was.
1950: One day, a little girl named Lily went to the park with her mom. She saw a big, red ball. She wanted to play with it. She asked her mom if she could help her mom.\n\nLily said, "Lily,
2000: One day, a little girl named Lily went to the park with her mom. She saw a big tree with a big tree. She wanted to climb the tree, but it was too high.\n\nLily's mom said, "Lily, I
ということで、ちゃんと生成モデルが学習できていると言って良さそうです。
ちなみに、手元のPCだと学習1ステップ（バッチサイズは64）に50秒くらいかかるので、ここまでに1日くらいかかってます。
ちゃんとしたライブラリを使った場合と比べると、速度は雲泥の差と言うほかないでしょう。
ちなみにこれでも初期の実装よりはかなり高速化しています。
最初はndarrayのメソッドを素直にラップしていたのですが、こうすると計算がシングルスレッドになってしまい、恐ろしく遅いです。
ということでrayonを使い色々な演算をマルチスレッド化していったのですが、こうなると基本的な演算も全てpar_iterを使って自分で書く必要があり、コードがすさまじくごちゃごちゃしてしまい大変つらかった...
その他色々反省点は多いですが、とりあえず文章が生成できたのでヨシ！という感じです。
まとめ
PyTorchって本当によくできてますね〜
-
これは微分幾何学に由来した用語ですが、自動微分と微分幾何学の関係についてはちゃんと理解できていません。おそらく https://arxiv.org/abs/1812.11592 あたりを真面目に読むと良いのだと思います。 ↩︎
-
基本的な演算についての
rruleはChainRules.jlというプロジェクトで管理されており、ここで具体例を見ることができます。 ↩︎
Discussion