TypeScriptでテンソル演算
バックエンド
テンソルまわり
いかにnumpyの再実装を避けるかという観点で
- TensorFlow.js を使えばいいのでは?
- TensorFlowのAPI自体は個人的に受け付けないけれど、今回の目的にはtf.Tensor+αだけあればいいので問題ないはず
- TSでの利用に関して、ドキュメントはあまり整備されてない雰囲気
- https://js.tensorflow.org/api/latest/
- コード眺めると型定義は色々書かれてはいるっぽい
- https://github.com/tensorflow/tfjs/blob/master/tfjs-core/src/types.ts
- 実際どの程度ありものの型定義でやれるのかは触ってみた方が速そう
実行エンジン
「PythonよりV8の方が速いのでは?」という素朴な期待を持っている。ここは早めに検証しておきたい
CUDA
- CUDAは現状OSへのインストールしか手段が提供されてないっぽい?
- TensorFlow.js はCUDAインストールの面倒は見てくれない模様
- TFJSは例によって特定バージョンのCUDAでしか動かないので、脳死最新版インストールはNG
- npmでCUDAを入れる・ローカルインストールするものはなさげ
Python同様に必要なバージョンのCUDA/cudnnを突っ込んだDockerイメージを用意するという手もあるけれど、これはこれで手間はかかる
TypeScript上での構文
TSで構文上の自由度が思ったよりない
- 演算子オーバーロードがない。つまり t1: Tensor と t2: Tensor に関して t1 + t2 といった自然な表記ができず、Add(t1,t2) のような表記を強制されてしまう
- Pythonでは演算子オーバーロード相当の処理が入っていて、ここは劣っている点になる
- プログラム側からASTを触る手段がない
- 総じて素のTSのみだと内部DSLは作りづらそうという印象
解決案1 JSX路線
- Babelのcustom jsx pluginを書く
- 基本的な演算でこれで書きやすくなるかはちょっと謎
const x1 = Tensor()
const x2 = Tensor()
// y = exp (x1+x2)^2
const y = <Tensor>
<exp><pow p=2><sum ops=[x1,x2] /></pow></exp>
</Tensor>
やっぱりダメなんじゃないかな感が強い
ネットワーク部分の記述にはアリかも
<TensorOp>
<ModuleBlock1>
<Conv2d stride=2 kernel=3 />
<Conv2d stride=2 kernel=3 />
<Conv2d stride=2 kernel=3 />
</ModuleBlock1>
<ModuleBlock2>
<ConvTranspose2d />
...
</ModuleBlock2>
</TensorOp>
define by runとは…となるけれど、部分的に利用する限りは選択肢の一つとしてあってもよさそう
解決案2 Literal type でDSLっぽいものにする
expression(t1, '+', t2) // '+' など文字列を入れる部分をLiteral typeで制約する
expression(t1, function1, function2, ...) // 部分適用を繰り返すイメージ
こういう方向性の書き味でやっていく
- 頻出関数は記号とか使って省略記法を用意してもいいかも
- テンソル演算を都度意識させられるけれど、これ自体は良し悪しなのであってよさそう
- どの道テンソル値はVariable相当の何かで包む必要があり、型に乗せる以上ユーザーはどうせ意識せざるを得ないはず
- プロパティを使えばさらに自然な書き方ができるかもしれない(できないかもしれない)
解決案3 外部DSL化
- Babel + webpackでゴリゴリやる前提なら・・・
- できればやりたくない
代替案
やはり演算子オーバーロードがないのは相当厳しい、のでPureScriptに行ってしまうという手はあるかもしれない…
-
ただこれをやると、TypeScriptという大文明圏から外れる獣道になる。「TypeScript(のようなプログラマーフレンドリーな言語)でやる」というコンセプトからは脱線することにPSでやることはそれはそれで意義はあると思うのだけれど
マイナー別言語でやるのは問題設定から外れるので一旦スコープ外にする。自分の場合、だいたいこういうことを考えているのは現実逃避を始めた時でもある-
TypeScriptから外れるなら、最初からRustあたりで組めばいいという話になりそう
- Rustには演算子オーバーロードもまともな型システムもある
- GPUメモリ管理や実行速度が求められる点でも深層学習向き
- Rust / PyTorchバインディングは既にあり、この場合は既存ライブラリ活用路線になりそう
妥協できるかどうか
- 問題になるのは正味演算子をまともに使えなくなる点
- その他の部分は良好、だと思う
- JSXでのネットワーク記述などは普通によさそう
演算子オーバーロードなしだとそこまで辛いか
// e^((a+b)^2 / 2)
tf.exp(a.plus(b).pow(2).div(2))
Pythonだと
tf.exp((a+b)^2 / 2)
- メソッドチェーンだと、より「手続きの記述」っぽくなる
- 「数式で書かれても分からない、コードで書かれたら分かる」という一部の人にはもしかしたらアリなのかもしれない
- 私の場合はナシ
- 宣言的に書きたい部分が大きいので、ここだけ手続き成分強いとなんだかなという気持ちになる
- そこまでゴリゴリやりたいならHaskell/PSに行けという話は出てくる
結論としては、やっぱりつらそう
解決策4 JSON as a DSL
- S式モドキをJSONモドキで書く
- 「JSはLisp系でしょ」「JSONとS式って似てるよね」というのは割と言われてきたこと[要出典]
- JSONそのものにする必要性はないので、TSのオブジェクト扱いでもいい(はず)
- ちゃんとやれば型チェック効かせられる、はず
/// MSE: (x: Tensor, y: Tensor) => sum((x-y)^2)/(x.length)
{"div": [{"sum": [{"pow": [{"dec": [x, y]} , 2]}]}, x.length]}
書いてから気づいたけど、オペレーター/関数はそのまま生のTSオブジェクト渡す方が楽かも
[div, [sum, [pow, [dec, [x, y]], 2]], x.length]
カッコが多くなるけどみんなpareditを使えば問題ない
TypeScriptのレイヤーでは数式を扱わず、パッケージ化されたなにかだけ使うようにするイメージでやるのがいいのかもしれない
突き詰めると S 式になるやつ、好きです。
PureScriptに関するメモ
- PS <- JS の利用はいいのだけど、PS -> JSのエクスポートはかなり罠が多そうな気がする
- 生成されるコードなどがあまり素直でなくなりそうな気がした
- 識別子に使える記号が多い方の言語から少ない方の言語に変換する時特有のキツさもありそう
- 最終的な利用者をPS側のみで想定するなら気にならないんだろうけど、"inter" operability を考え始めた瞬間に終わりそう
type Op<C extends number> = {f: Function, argc: C}
type Data = number // Tensor in real
type DType = 'float16' | 'float32'
type Var<T extends DType> = { dtype: T, data: Data }
type Arg2<T extends DType> = [Var<T> | Sexp<T>, Var<T> | Sexp<T>]
type Arg1<T extends DType> = [Var<T> | Sexp<T>]
type Arg3<T extends DType> = [Var<T> | Sexp<T>, Var<T> | Sexp<T>, Var<T> | Sexp<T>]
type Var2<T extends DType> = [Var<T>, Var<T>]
type VArgs<T extends DType> = Arg2<T> | Arg1<T> | Arg3<T> // | Arg4<T>|Arg0<T>|Arg6<T> ...
type Sexp<T extends DType> = [Op<1>, Arg1<T>] | [Op<2>, Arg2<T>] | [Op<3>, Arg3<T>]
type VarConvertable = number // |Tensor ...
function isVar<T extends DType>(s: any): s is Var<T> {
return s.dtype !== undefined
}
function createVar<T extends DType>(x: VarConvertable, dtype: T): Var<T> {
return { dtype: dtype, data: x }
}
function _div<T extends DType>([a, b]: Var2<T>): Var<T> {
return { dtype: a.dtype, data: a.data / b.data }
}
function _plus<T extends DType>([a, b]: Var2<T>): Var<T> {
return { dtype: a.dtype, data: a.data + b.data }
}
const div: Op<2> = { f: _div, argc: 2 }
const plus: Op<2> = { f: _plus, argc: 2 }
function evalSexp<T extends DType>(s: Sexp<T>): Var<T> {
const [op, args] = s
const evaledArgs: Var<T>[] = args.map(arg => {
if (isVar(arg)) {
return arg
} else {
return evalSexp(arg)
}
})
return op.f(evaledArgs)
}
const x1 = createVar<'float16'>(1, 'float16')
const x2 = createVar<'float16'>(2, 'float16')
console.log(x1, x2)
const sample: Sexp<'float16'> = [div, [[plus, [x1, x2]], x2]]
console.log(evalSexp(sample))
TypeScript型芸
やりたかったこと
- Tensor型でnumber[]のshapeを受け取り、型レベルでshapeの情報を持ちつつバックエンド(ここではTensorFlow)のコンストラクタにshapeデータを渡す
- この際同じ内容を、型とデータで二度書きしたくない
// 理想
const x: Tensor<[2,3]> = tensorBuilder<[2,3]>()
// or
const y: Tensor<[2,3]> = tensorBuilder([2,3])
結論
無理っぽい。TypeScriptにはTS上の型情報を実行時に取得する方法が一切用意されてなさそう
- コンパイル時に決定する値を定数値として使えるんじゃないかなーという期待はあった
- あるいは黒魔術リフレクション枠みたいなものがあるのではと期待はしていた
// これはできる
const x: Tensor<[2,3]> = tensorBuilder<[2,3]>([2,3])
全く同じ内容をコピペするのは微妙だけど、ひとまずこれで満足することにする
追記
Typescript Design Goalsには"式レベルの構文を追加しない" "型情報をランタイムに影響させない" が挙げられているので、皆さんの好きそうな機能は入らないことになっている
https://github.com/Microsoft/TypeScript/wiki/TypeScript-Design-Goals
はい
shape算
- せっかくLiteral Typeが使えるのだから、テンソル演算する関数の出力shapeを型推論で導出したくなった
- 出力型の計算に四則演算が入るタイプの実現は難しそう(例:Conv2D)
Conv2dのShapeの導出例。単純ながら四則演算が必須
解決法
- ユーザーに明示的に出力shapeを書かせて、条件の表明をさせる
- Babelプラグインなどを活用してマクロ的な何かができるようにする
いずれにしろ結構キツい
関連メモ
Rustにはマクロがあるので色々できそうだけど、Literal Type相当のものはないのでこちらも楽ではなさそう
TBD: 型レベルで個数可変のタプルを扱いたい
type NonNullListA = Exclude<number[], []>
const l1: NonNullListA = [] // passed...
type ListLike = []|[number]|[number, number]|[number, number, number] // ...
type NonNullListB = Exclude<ListLike, []>
const l2: NonNullListB = [] // ERROR, as expected
type NonNullList = [number, ...number[]];
const l3: NonNullList = [] // ERROR
type T = [...number[]]
type A = [number, ...number[]]
type U = A extends T ? true : false // true
const a: T = []
const b: A = [] // ERROR
メモ
Conditional types + infer で何かやりようある気がしているのでやる