基礎から実践 TensorFlow 重み共有

TensorFlow での重み共有と tf.variable_scope の説明を行います。
公式：Variables

はじめに

• ネットワークの一部の重みを共有することで効率的な学習を行わせたい
• RNN 等で、学習時と生成時に異なるネットワークを使うが重みは共通のものを使いたい

このような理由でネットワークの異なる部分の重み（より厳密には変数 Variable）を共有したい場合あがあります。
このようなときには２つの方法があります。

• Variable を共通のインスタンスとして作成し、ネットワークに渡す
  • この方法はシンプルですが、外から Variable や Layer を渡せないライブラリを使う場合には使えません。
• variable_scope を使う

この記事では後者を使って以下のようなコードで重み共有をする方法と、その裏側について説明します。

with tf.Session(graph=tf.Graph()) as sess:
    inputs_x = ...
    inputs_y = ...
    with tf.variable_scope('some_shared_weight_network') as scope:
        outputs_x = create_some_network(inputs_x)
        scope.reuse_variables()  # 重みの共有を ON
        outputs_y = create_some_network(inputs_y)  # 上の create_some_network と重みを共有

この記事のコードは Github にあげています。

基本： Variable, variable_scope, get_variable

Variable と get_variable

変数を作成するには tf.Variable と tf.get_variable があります。

• tf.Variable
  • 単純に変数を作成します。
  • 同じ名前の変数が既に存在する場合、変数名に _1 等数字をつけて登録します。
    • some_variable, some_variable_1, some_variable_2, ...
• tf.get_variable
  • 重みの共有を行う可能性がある場合に使います。（基本こちらを使うのが良いです。）
  • 同名の変数が無ければ作成、あれば既存の変数を取得します
  • 挙動の詳細は後述する variable_scope の reuse オプションによって変わります。

それでは、 Variable, get_variable を使って変数を作ってみましょう。

def show_variables():
    print('\n'.join([v.name for v in tf.global_variables()]))

# variable_scope 無しでの Variable/get_variable
with tf.Session(graph=tf.Graph()) as sess:
    v_a = tf.Variable(tf.random_uniform(shape=[2, 3]), name='variable')  # -> variable:0
    v_b = tf.Variable(tf.random_uniform(shape=[2, 3]), name='variable')  # -> variable_1:0
    v_c = tf.get_variable('variable', shape=[2, 3])  # -> variable_2:0
    # v_d = tf.get_variable('variable', shape=[2, 3])  # -> ERROR
    show_variables()

出力

variable:0
variable_1:0
variable_2:0

• tf.Variable: 同名の変数がある場合 _1 等数字の suffix をつけて作成
• tf.get_variable: get_variable によって作られた同名の変数がある場合、エラー
  • 後述するように、デフォルトでは get_variable は（get_variable によって宣言された）同名の変数が既にあるとエラーを出します。これは意図せず変数を共有してしまうバグを防ぐためです。
  • （ややこしいことに、 Variable によって作られた変数の場合、 get_variable 時に名前がかぶっていても怒られないようですね。）

get_variable は variable_scope と組み合わせることでその真価を発揮します。次の節でみていきます。

variable_scope で変数を階層化・共有する

以下では tf.Variable のことは置いておいて、 tf.get_variable と tf.variable_scope を見ていきます。

variable_scope は reuse という引数を持っており、この値によって、 get_variable 時に変数を新規作成するか、再利用（共有）するかを選択できます。

# variable_scope を使うことで名前を階層化できる
with tf.Session(graph=tf.Graph()) as sess:
    with tf.variable_scope('hoge'):
        v_a = tf.get_variable('variable', shape=[2, 3])  # -> hoge/variable_2:0
        # v_b = tf.get_variable('variable', shape=[2, 3])  #-> ERROR: reuse 指定なし（resuse=None）時は変数の再利用は不可
        
    with tf.variable_scope('hoge', reuse=True):  # 変数を共有
        v_c = tf.get_variable('variable', shape=[2, 3])  # -> hoge/variable_2:0 共有した！
        # v_d = tf.get_variable('variable_new', shape=[2, 3])  # -> ERROR: reuse=True 時は変数の新規作成は不可
        show_variables()

出力

hoge/variable:0

• 変数名が hoge/variable というように階層化
• reuse を指定しないと、同名の変数は作成できない
• reuse=True にすると同名の変数を再利用（共有）できるが、新規の変数は作成できない

このように variable_scope の reuse オプションによって get_variable の変数の作成・取得のモードを切り替えることができます。

（上記は tf.get_variable を使った場合のみです。 tf.Variable を使った場合、最初の例のように suffix をつけながら全く別の変数が作成されます。）

variable_scope の reuse option

variable_scope の reuse オプションには以下の3種類があります。

	新規の変数名	既存の変数名	備考
reuse=None	変数を新規作成	エラー	デフォルトはこれ
reuse=True	エラー	既存の変数を返す
reuse=tf.AUTO_REUSE	変数を新規作成	既存の変数を返す	便利だけど、バグを生みやすい

コードで確認しましょう。

with tf.Session(graph=tf.Graph()) as sess:
    with tf.variable_scope('reuse_none', reuse=None):  # 再利用しない。デフォルト。
        v_none_a = tf.get_variable('variable', shape=[2, 3])  # -> reuse_none/variable:0
        # v_none_b = tf.get_variable('variable', shape=[2, 3])  # -> ERROR

    with tf.variable_scope('reuse_true', reuse=None):  # reuse=True だと変数の作成ができないので予め reuse=None で作る
        v_true_a = tf.get_variable('variable', shape=[2, 3])  # -> reuse_true/variable:0
    with tf.variable_scope('reuse_true', reuse=True):
        v_true_a = tf.get_variable('variable', shape=[2, 3])  # -> reuse_true/variable:0
        # v_true_b = tf.get_variable('variable_b', shape=[2, 3])  # -> ERROR

    with tf.variable_scope('auto_reuse', reuse=tf.AUTO_REUSE):  # 無ければ作成、あれば再利用。便利だけど危険。
        v_none_a = tf.get_variable('variable', shape=[2, 3])  # -> auto_reuse/variable:0
        v_none_b = tf.get_variable('variable', shape=[2, 3])  # -> auto_reuse/variable:0

    show_variables()

出力

reuse_none/variable:0
reuse_true/variable:0
auto_reuse/variable:0

reuse したくなった時に毎回 variable_scope を書き直すのは面倒です。
その場合、以下のような書式で途中から reuse=True に変更することもできます。

with tf.Session(graph=tf.Graph()) as sess:
    with tf.variable_scope('reuse_true') as scope:
        # reuse=None (default)
        v_true_a = tf.get_variable('variable', shape=[2, 3])  # -> reuse_true/variable:0

        scope.reuse_variables()
        # reuse=True
        v_true_a = tf.get_variable('variable', shape=[2, 3])  # -> reuse_true/variable:0
        # v_true_b = tf.get_variable('variable_b', shape=[2, 3])  # -> ERROR
    show_variables()

reuse option の継承

公式のドキュメントに以下のように書いてあるとおり、 reuse option は内側のスコープに継承されます。

Note that the reuse flag is inherited: if we open a reusing scope, then all its sub-scopes become reusing as well.

この説明だとわかりにくいのですが、正確には以下の表ように、内側のスコープが reuse=None の場合には外側のスコープの reuse option が使われるようです。

	inner: None	inner: True	innner: AUTO_REUSE
outer: None	None	True	AUTO_REUSE
outer: True	True	True	AUTO_REUSE
outer: AUTO_REUSE	AUTO_REUSE	True	AUTO_REUSE

この表は以下のコードで確認しました。

def show_reuse(scope):
    print('{}: {}'.format(scope.name, scope.reuse))
    
with tf.variable_scope('reuse_none', reuse=None) as outer_scope:
    show_reuse(outer_scope)
    with tf.variable_scope('reuse_none', reuse=None) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('reuse_true', reuse=True) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('auto_reuse', reuse=tf.AUTO_REUSE) as inner_scope:
        show_reuse(inner_scope)

with tf.variable_scope('reuse_true', reuse=True) as outer_scope:
    show_reuse(outer_scope)
    with tf.variable_scope('reuse_none', reuse=None) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('reuse_true', reuse=True) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('auto_reuse', reuse=tf.AUTO_REUSE) as inner_scope:
        show_reuse(inner_scope)

with tf.variable_scope('auto_reuse', reuse=tf.AUTO_REUSE) as outer_scope:
    show_reuse(outer_scope)
    with tf.variable_scope('reuse_none', reuse=None) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('reuse_true', reuse=True) as inner_scope:
        show_reuse(inner_scope)
    with tf.variable_scope('auto_reuse', reuse=tf.AUTO_REUSE) as inner_scope:
        show_reuse(inner_scope)

出力

reuse_none: False
reuse_none/reuse_none: False
reuse_none/reuse_true: True
reuse_none/auto_reuse: _ReuseMode.AUTO_REUSE
reuse_true: True
reuse_true/reuse_none: True
reuse_true/reuse_true: True
reuse_true/auto_reuse: _ReuseMode.AUTO_REUSE
auto_reuse: _ReuseMode.AUTO_REUSE
auto_reuse/reuse_none: _ReuseMode.AUTO_REUSE
auto_reuse/reuse_true: True
auto_reuse/auto_reuse: _ReuseMode.AUTO_REUSE

実践重み共有

重みが共有されているかを確認するため、簡単なチェック用のプログラムを作っておきましょう。

def has_same_outputs(tensor1, tensor2, sess):
    # ２つのテンソルの値が同じかを返します
    sess.run(tf.global_variables_initializer())
    result1, result2 = sess.run([tensor1, tensor2])
    return result1.tolist() == result2.tolist()

単純なフィードフォワードネットワーク

# ネットワークの作成を部品化
def feed_forward(inputs):
    d1 = tf.layers.dense(inputs, 20, name='dense_a')
    d2 = tf.layers.dense(d1, 30, name='dense_b')
    return d2

重み共有をしない場合

# 独立したネットワークを２つ作る例
with tf.Session(graph=tf.Graph()) as sess:
    inputs = tf.ones(shape=[2, 3])
    with tf.variable_scope('feed_forward_x'):
        outputs_x = feed_forward(inputs)
    with tf.variable_scope('feed_forward_y'):
        outputs_y = feed_forward(inputs)
    print('shared: {}'.format(has_same_outputs(outputs_x, outputs_y, sess)))
    show_variables()

出力

shared: False
feed_forward_x/dense_a/kernel:0
feed_forward_x/dense_a/bias:0
feed_forward_x/dense_b/kernel:0
feed_forward_x/dense_b/bias:0
feed_forward_y/dense_a/kernel:0
feed_forward_y/dense_a/bias:0
feed_forward_y/dense_b/kernel:0
feed_forward_y/dense_b/bias:0

重み共有をする場合

# 重みを共有した実質一つのネットワークを２つ作る例
with tf.Session(graph=tf.Graph()) as sess:
    inputs = tf.ones(shape=[2, 3])
    with tf.variable_scope('feed_forward_shared') as scope:
        outputs_x = feed_forward(inputs)
        scope.reuse_variables()
        outputs_y = feed_forward(inputs)
    print('shared: {}'.format(has_same_outputs(outputs_x, outputs_y, sess)))
    show_variables()

出力

shared: True
feed_forward_shared/dense_a/kernel:0
feed_forward_shared/dense_a/bias:0
feed_forward_shared/dense_b/kernel:0
feed_forward_shared/dense_b/bias:0

tf.layers.dense の実装をみてみるとわかりますが、 self.add_variable で重みの作成もしくは読み込みを行っています。
add_variable の定義 をみてみると、 add_variable は自動的に variable_scope を切ってその中で get_variable しているのがわかります。
このように、 tf.layers の各レイヤーは variable_scope + get_variable ベースで変数を定義する作法になっているようです。（全ての layer で確認をしているわけではありません。）

変数名を出力させた結果も、 (自分の作ったスコープ)/(name属性で渡した名前)/(kernel or bias) となっていて見通しが良いですね。

Seq2Seq のデコーダを学習時と生成時で分ける

最後に、より実践的な例としてシンプルな Seq2Seq のデコーダを重み共有を用いて実装してみましょう。

Seq2Seq は時系列を別の時系列に変換するネットワークで、主に翻訳などに使われます。
Seq2Seq は、時系列を ThoughtVector と呼ばれるその時系列を表すベクトルに変換する Encoder 部分と、その ThoughtVector から別の時系列に変換する Decoder 部分からなります。

この Decoder は学習時と生成時で処理が変わります。

• 学習時：時刻 t-1 の教師信号から時刻 t の出力を出す
• 生成時：時刻 t-1 の出力から時刻 t の出力を出す

これを簡単に行うための関数が tensorflow には組み込まれていますが、学習時・生成時で別の関数となっています。
従って学習時と生成時のネットワークは別にする必要があります。しかし当然、学習で得られた重みを生成時にも使う必要があります。

Decoder に含まれる変数は、 GRU などのセルの重みと、 Decoder の隠れ層から出力層への写像のレイヤー（output_layer）の二種類です。（更にアテンションメカニズムを使ったりすると他にも変数が増えますが。）
これらを variable_scope の中で作ることで変数を共有します。
（厳密には入力の ID 列を分散表現にする embedding も変数ですが、この例では encoder も含めて共有する変数にするため外から渡す方式でやっています。）

import numpy as np
vocab_size = 3000
embedding_dim = 256
hidden_dim = 256
batch_size = 2
beam_width = 2
max_len = 8
decoder_scope = 'decoder'
cell_name = 'cell'
out_layer_name = 'out_layer'

def create_encoder(inputs, inputs_length, embedding):
    inputs_embedded = tf.nn.embedding_lookup(embedding, inputs)
    cell = tf.nn.rnn_cell.GRUCell(hidden_dim, name=cell_name)
    outputs, final_state = tf.nn.dynamic_rnn(
        cell=cell,
        inputs=inputs_embedded,
        sequence_length=inputs_length,
        dtype=tf.float32,
        scope='encoder'
    )
    return final_state

def create_trainer_decoder(thought_vector, embedding, inputs, inputs_length):
    cell = tf.nn.rnn_cell.GRUCell(hidden_dim, name=cell_name)
    output_layer = tf.layers.Dense(vocab_size, use_bias=False, name=out_layer_name)
    inputs_embedded = tf.nn.embedding_lookup(embedding, inputs)
    helper = tf.contrib.seq2seq.TrainingHelper(
        inputs=inputs_embedded,
        sequence_length=inputs_length,
    )
    decoder = tf.contrib.seq2seq.BasicDecoder(cell, helper, thought_vector, output_layer=output_layer)
    outputs, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder, swap_memory=True, scope=decoder_scope)
    return outputs.rnn_output

def create_generation_decoder(thought_vector, embedding):
    cell = tf.nn.rnn_cell.GRUCell(hidden_dim, name=cell_name)
    output_layer = tf.layers.Dense(vocab_size, use_bias=False, name=out_layer_name)
    start_tokens = tf.ones([batch_size], tf.int32)  # BOS==1
    end_token = 2  # EOS==2
    tiled_thought_vector = tf.contrib.seq2seq.tile_batch(thought_vector, multiplier=beam_width)

    decoder = tf.contrib.seq2seq.BeamSearchDecoder(
        cell=cell,
        embedding=embedding,
        start_tokens=start_tokens,
        end_token=end_token,
        initial_state=tiled_thought_vector,
        beam_width=beam_width,
        output_layer=output_layer,
    )

    decoder_outputs, _, _ = tf.contrib.seq2seq.dynamic_decode(
        decoder=decoder, maximum_iterations=max_len, scope=decoder_scope
    )
    return decoder_outputs.predicted_ids

with tf.Session(graph=tf.Graph()) as sess:
    encoder_inputs = tf.ones(shape=[batch_size, max_len], dtype=tf.int32)
    encoder_inputs_length = tf.ones(shape=[batch_size], dtype=tf.int32) * max_len
    decoder_inputs = tf.ones(shape=[batch_size, max_len], dtype=tf.int32)
    decoder_inputs_length = tf.ones(shape=[batch_size], dtype=tf.int32) * max_len

    embedding = tf.Variable(tf.random_uniform([vocab_size, embedding_dim], -1.0, 1.0), dtype=tf.float32, name='embedding')
    thought_vector = create_encoder(encoder_inputs, encoder_inputs_length, embedding)
    
    with tf.variable_scope('decoder') as scope:
        train_outputs = create_trainer_decoder(
            thought_vector,
            embedding,
            decoder_inputs,
            decoder_inputs_length,
        )
        scope.reuse_variables()
        generation_outputs = create_generation_decoder(
            thought_vector,
            embedding,
        )
    
    sess.run(tf.global_variables_initializer())
    train_result, generation_result = sess.run([train_outputs, generation_outputs])
    train_ids = np.argmax(train_result, axis=-1)
    generation_ids = generation_result[:,:,0]
    print(train_ids)
    print(generation_ids)
    show_variables()

出力

[[1260 1260 1260 1526 1526 1526  367  367]
 [1260 1260 1260 1526 1526 1526  367  367]]
[[1260  222  321 1377 1377 2386 2686 2686]
 [1260  222  321 1377 1377 2386 2686 2686]]
embedding:0
encoder/cell/gates/kernel:0
encoder/cell/gates/bias:0
encoder/cell/candidate/kernel:0
encoder/cell/candidate/bias:0
decoder/decoder/cell/gates/kernel:0
decoder/decoder/cell/gates/bias:0
decoder/decoder/cell/candidate/kernel:0
decoder/decoder/cell/candidate/bias:0
decoder/decoder/out_layer/kernel:0

デコーダの処理の違いによって時刻2以降は違う値になりますが、最初の時刻の出力は同じであることがわかります。（つまり重みが共有されています。 BeamSearch を行っているので、厳密には最初の時刻でも異なる出力が出る可能性はあります。）
また、変数一覧を見ても余分な変数ができていないのがわかると思います。

（因みにこのコードは BeamSearchDecoder のサンプルとしても有用だと思っています(ΦωΦ)

まとめ

• tf.variable_scope で囲うことで、変数の名前を階層化できる
• tf.variable_scope の reuse option によって、変数を共有したり新規作成したりを制御できる
• tf.variable_scope と合わせて階層的なネットワーク部品を作りたい場合 tf.Variable ではなく
  tf.get_variable を使うべし
• tf.layers 等公式の主なレイヤーは tf.variable_scope/tf.get_variable を使っているので tf.variable_scope ベースの重みの階層化・共有ができる

手元で動かしてみたい人は Github からコードを clone して動かしてみて下さい。
動作の確認は tensorflow 1.5 で行っています。