Pyro 1.4 : Examples : 深層マルコフモデル (翻訳)

翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 08/06/2020 (1.4.0)

* 本ページは、Pyro の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

Deep Markov Model

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

★ 無料セミナー開催中 ★ クラスキャット主催人工知能 & ビジネス Web セミナー

人工知能とビジネスをテーマにウェビナー (WEB セミナー) を定期的に開催しています。スケジュールは弊社公式 Web サイトでご確認頂けます。

お住まいの地域に関係なく Web ブラウザからご参加頂けます。事前登録 が必要ですのでご注意ください。
Windows PC のブラウザからご参加が可能です。スマートデバイスもご利用可能です。

◆ お問合せ : 本件に関するお問い合わせ先は下記までお願いいたします。

株式会社クラスキャット セールス・マーケティング本部セールス・インフォメーション

E-Mail：sales-info@classcat.com ; WebSite: https://www.classcat.com/

Facebook: https://www.facebook.com/ClassCatJP/

深層マルコフモデル

イントロダクション

シーケンシャルデータのための深層確率モデルを構築していきます: 深層マルコフモデル です。モデル化することを望む特定のデータセットは多声音楽 (= polyphonic music) のスニペットから成ります。シークエンスの各時間スライスは四分音符に渡りそしてその時間ステップで音符をエンコードする 88-次元二値ベクトルで表されます。

音楽は (明らかに) 時間的に首尾一貫しているので、私達は観測データで複雑な時間依存を表せるモデルが必要です。例えば、特定の時間ステップの音符が前の時間ステップの音符から独立であるようなモデルを考えることは適切ではないでしょう。これを行なう一つの方法は潜在変数モデルを構築することです、そこでは観測の変動性と時間構造は潜在変数のダイナミクスにより制御されます。

このアイデアの一つの特定の具現化はマルコフモデルです、そこでは潜在変数の連鎖を持ち、連鎖の各潜在変数は前の潜在変数上で条件付けられます。これはパワフルなアプローチですが、複雑なデータを複雑な (そしてこの場合未知の) ダイナミクスで表すことを望む場合、潜在的に非常に非線形であるダイナミクスを提供するためにモデルが十分に柔軟であって欲しいです。こうして深層マルコフモデルは : 潜在変数のダイナミクスを支配する (= govern) 遷移確率 そして (非線形) ニューラルネットワークによりパラメータ化される潜在ダイナミクスにより観測がどのように生成されるかを支配する 放出 (= emission) 確率 を考慮します。

実装していく特定のモデルは次のリファレンスに基づきます :

Structured Inference Networks for Nonlinear State Space Models, Rahul G. Krishnan, Uri Shalit, David Sontag

このチュートリアルの読者がリファレンスを読んでいることは仮定しない一方で、他の時系列モデルのコンテキストで深層マルコフモデルのより包括的な議論を求めるに間違いなく良い場所であることに注意してください。

モデルを説明しましたが、それの訓練にどのように取り組みますか？利用していく推論ストラテジーは 変分推論 で、これは、潜在確率変数に渡る事後分布を近似するために使用できるパラメータ化された分布の族を指定する ことを必要とします。私達のモデルとデータに固有な非線形と複雑な時間依存性が与えられたとき、正確な事後分布が非常に自明でないことを想定します。そこで良いモデルを学習することを希望すれば変分分布の柔軟な族を必要とします。幸い、PyTorch と Pyro は一緒に総ての必要な構成要素を提供しています。見ていくように、それらを集めることは簡単です。作業に進みましょう。

モデル

モデルの高位構造を記述する便利な方法はグラフィカルモデルによります。

Figure 1: T=3 時間ステップのために展開されたモデル

ここで、観測のシークエンスが長さ 3 : $\{{\bf x}_1, {\bf x}_2, {\bf x}_3\}$ と仮定してモデルを展開しました。観測のシークエンスを反映して潜在確率変数のシークエンス : $\{{\bf z}_1, {\bf z}_2, {\bf z}_3\}$ も持ちます。図はモデルの構造をエンコードしています。対応する同時分布は :

$p({\bf x}_{123} , {\bf z}_{123})=p({\bf x}_1|{\bf z}_1)p({\bf x}_2|{\bf z}_2)p({\bf x}_3|{\bf z}_3)p({\bf z}_1)p({\bf z}_2|{\bf z}_1)p({\bf z}_3|{\bf z}_2)$

${\bf z}_t$ に条件付けられて、各観測 ${\bf x}_t$ は他の観測から独立です。これは、下向きの矢印により示されるように、各観測 ${\bf x}_t$ が対応する潜在変数 ${\bf z}_t$ にだけ依拠するという事実から読み取れます。モデルのマルコフ性もまた読み取れます : 各潜在変数 ${\bf z}_t$ は、前の潜在変数 ${\bf z}_{t-1}$ に条件付けられるとき、総ての前の潜在変数 $\{ {\bf z}_{t-2}, {\bf z}_{t-3}, …\}$ から独立です。これは、時間 $t$ における系の状態について知る必要がある総ては潜在変数 ${\bf z}_{t}$ によりカプセル化されていることを効果的に言っています。

観測を制御する観測尤度, i.e. 確率分布 $p({{\bf x}_t}|{{\bf z}_t})$ はベルヌーイ分布により与えられることを仮定します。これは適切な選択です、何故ならば私達の観測は総て 0 か 1 であるからです。潜在的ダイナミクスを制御する確率分布 $p({\bf z}_t|{\bf z}_{t-1})$ に対しては、対角共分散を持つ (条件付き) ガウス分布を選択します。これは合理的です、何故ならば潜在変数空間は連続であるからです。

無地の黒い四角はニューラルネットワークによりパラメータ化された非線形関数を表します。これはこれを深層マルコフモデルとするものです。黒い四角は 2 つの異なる場所に現れることに注意してください : 潜在変数のペアの間と潜在変数と観測の間です。潜在変数を接続する非線形関数 (Fig. 1 の ‘Trans’) は潜在変数のダイナミクスを制御します。 ${\bf z}_{t}$ の条件付き確率分布に複雑な方法で ${\bf z}_{t-1}$ に依拠することを許容しますので、モデルで複雑なダイナミクスを捕捉できるでしょう。同様に、潜在変数を観測に接続する非線形関数 (Fig. 1 の ‘Emit’) は観測がどのように潜在変数に依拠するかを制御します。

幾つかの追加ノート :- 手元の問題に適合するために潜在空間の次元を自由に選択できます : 単純な問題のためには小さい潜在空間をそして複雑なダイナミクスを持つ問題のためにはより大きな空間を – Fig. 1 のパラメータ ${\bf z}_0$ に注意してください。コードからより明らかになるように、これは最初の時間ステップのための確率分布 $p({\bf z}_1)$ をパラメータ化する単に便利な方法で、そこではその上で条件付ける前の潜在変数はありません。

Gated Transition と Emitter

さっさと幾つかのコードを書き始めましょう。最初に 2 つの PyTorch モジュールを定義します、これらは Fig.1 の黒い四角に対応します。最初に emission (放出) 関数は :

class Emitter(nn.Module):
    """
    Parameterizes the bernoulli observation likelihood p(x_t | z_t)
    """
    def __init__(self, input_dim, z_dim, emission_dim):
        super().__init__()
        # initialize the three linear transformations used in the neural network
        self.lin_z_to_hidden = nn.Linear(z_dim, emission_dim)
        self.lin_hidden_to_hidden = nn.Linear(emission_dim, emission_dim)
        self.lin_hidden_to_input = nn.Linear(emission_dim, input_dim)
        # initialize the two non-linearities used in the neural network
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
 
    def forward(self, z_t):
        """
        Given the latent z at a particular time step t we return the vector of
        probabilities `ps` that parameterizes the bernoulli distribution p(x_t|z_t)
        """
        h1 = self.relu(self.lin_z_to_hidden(z_t))
        h2 = self.relu(self.lin_hidden_to_hidden(h1))
        ps = self.sigmoid(self.lin_hidden_to_input(h2))
        return ps

コンストラクタで emission 関数で使用される線形変換を定義します。emission_dim はニューラルネットワークの隠れユニットの数であることに注意してください。私達が使用していく非線形も定義します。forward 呼び出しは関数の計算フローを定義します。入力として潜在変数 ${\bf z}_{t}$ を取りそして長さ 88 のベクトルを得るまで変換のシークエンスを行ないます、これはベルヌーイ尤度の emission 確率を定義します。sigmoid ですから、ps の各要素は 0 と 1 の間で正当な確率を定義します。ps の要素が一緒に取られて ( ${\bf z}_{t}$ でエンコードされるような) 系の状態が与えられたとき時間 $t$ にどの音符を観測することを期待するかエンコードします。

今は gated transition (遷移) 関数を定義します :

class GatedTransition(nn.Module):
    """
    Parameterizes the gaussian latent transition probability p(z_t | z_{t-1})
    See section 5 in the reference for comparison.
    """
    def __init__(self, z_dim, transition_dim):
        super().__init__()
        # initialize the six linear transformations used in the neural network
        self.lin_gate_z_to_hidden = nn.Linear(z_dim, transition_dim)
        self.lin_gate_hidden_to_z = nn.Linear(transition_dim, z_dim)
        self.lin_proposed_mean_z_to_hidden = nn.Linear(z_dim, transition_dim)
        self.lin_proposed_mean_hidden_to_z = nn.Linear(transition_dim, z_dim)
        self.lin_sig = nn.Linear(z_dim, z_dim)
        self.lin_z_to_loc = nn.Linear(z_dim, z_dim)
        # modify the default initialization of lin_z_to_loc
        # so that it's starts out as the identity function
        self.lin_z_to_loc.weight.data = torch.eye(z_dim)
        self.lin_z_to_loc.bias.data = torch.zeros(z_dim)
        # initialize the three non-linearities used in the neural network
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.softplus = nn.Softplus()
 
    def forward(self, z_t_1):
        """
        Given the latent z_{t-1} corresponding to the time step t-1
        we return the mean and scale vectors that parameterize the
        (diagonal) gaussian distribution p(z_t | z_{t-1})
        """
        # compute the gating function
        _gate = self.relu(self.lin_gate_z_to_hidden(z_t_1))
        gate = self.sigmoid(self.lin_gate_hidden_to_z(_gate))
        # compute the 'proposed mean'
        _proposed_mean = self.relu(self.lin_proposed_mean_z_to_hidden(z_t_1))
        proposed_mean = self.lin_proposed_mean_hidden_to_z(_proposed_mean)
        # assemble the actual mean used to sample z_t, which mixes
        # a linear transformation of z_{t-1} with the proposed mean
        # modulated by the gating function
        loc = (1 - gate) * self.lin_z_to_loc(z_t_1) + gate * proposed_mean
        # compute the scale used to sample z_t, using the proposed
        # mean from above as input. the softplus ensures that scale is positive
        scale = self.softplus(self.lin_sig(self.relu(proposed_mean)))
        # return loc, scale which can be fed into Normal
        return loc, scale

これは上の Emitter の構造を反映し、計算フローが少し複雑であるという違いを伴います。これは 2 つの理由のためです。最初に、GatedTransition の出力は正当な (対角) ガウス分布を定義する必要があります。そこで 2 つのパラメータを出力する必要があります : mean loc、そして (平方根) 共分散 scale です。これら両者は潜在空間と同じ次元を持つ必要があります。2 番目に、ダイナミクスを非線型を強制することを望みません。こうして私達の mean loc は 2 つの項の合計になります、その一つだけが入力 z_t_1 上で非線形性に依拠します。この方法で線形と非線型ダイナミクスの両者をサポートできます (あるいは実際には潜在空間の一部のダイナミクスを線形に、一方でダイナミクスの残りは非線型にします)。

モデル – Pyro 確率関数

ここまで私達が行なったことの総ては純粋な PyTorch です。モデルをコードに翻訳することを終えるため Pyro を図式に持ち込む必要があります。基本的には Fig. 1 の確率的ノード (i.e. 円) を実装する必要があります。これを行なうため呼び出し可能な model() を導入します、これは Pyro プリミティブ pyro.sample を含みます。sample ステートメントは潜在変数 ${\bf z}_{1:T}$ に渡る同時分布を指定するために使用されます。更に、観測 ${\bf x}_{1:T}$ が潜在変数にどのように依拠するかを指定するため obj 引数が sample ステートメントで使用できます。model() のための完全なコードを見る前に、主ロジックを含む必要最小限のバージョンを見ましょう :

def model(...):
    z_prev = self.z_0
 
    # sample the latents z and observed x's one time step at a time
    for t in range(1, T_max + 1):
        # the next two lines of code sample z_t ~ p(z_t | z_{t-1}).
        # first compute the parameters of the diagonal gaussian
        # distribution p(z_t | z_{t-1})
        z_loc, z_scale = self.trans(z_prev)
        # then sample z_t according to dist.Normal(z_loc, z_scale)
        z_t = pyro.sample("z_%d" % t, dist.Normal(z_loc, z_scale))
 
        # compute the probabilities that parameterize the bernoulli likelihood
        emission_probs_t = self.emitter(z_t)
        # the next statement instructs pyro to observe x_t according to the
        # bernoulli distribution p(x_t|z_t)
        pyro.sample("obs_x_%d" % t,
                    dist.Bernoulli(emission_probs_t),
                    obs=mini_batch[:, t - 1, :])
        # the latent sampled at this time step will be conditioned upon
        # in the next time step so keep track of it
        z_prev = z_t

私達が行なう必要がある最初のことは ${\bf z}_1$ をサンプリングすることです。ひとたび ${\bf z}_1$ をサンプリングしたのであれば、 ${\bf z}_2 \sim p({\bf z}_2|{\bf z}_1)$ 等もサンプリングできます。これが for ループで実装されたロジックです。確率分布 $p({\bf z}_t|{\bf z}_{t-1})$ を定義するパラメータ z_loc と z_scale は self.trans を使用して計算されます、これは上で定義された GatedTransition モジュールの単なるインスタンスです。 $t=1$ の最初の時間ステップに対して self.z_0 上で条件付けます、これは (訓練可能な) Parameter です、一方で続く時間ステップに対しては前にドローされた潜在変数上で条件付けします。各確率変数 z_t はユーザにより一意な名前が割り当てられることに注意してください。

一度与えられた時間ステップで ${\bf z}_t$ をサンプリングしたのであればデータポイント ${\bf x}_t$ を観測する必要があります。そこで emission_probs_t を得るために上で定義された Emitter モジュールのインスタンスである self.emitter に z_t を渡します。sample ステートメントの引数 dist.Bernoulli() と共に、これらの確率は観測尤度を完全に指定します。最後に、sample への obs 引数を使用して観測データ ${\bf x}_t$ のスライス : mini_batch[:, t – 1, :] も指定します。

これはモデルを完全に指定してそれを Pyro に渡せる callable にカプセル化します。先へ進む前に、model() の完全バージョンを見て、最初のパスで注釈をつけた詳細の幾つかを通り抜けましょう。

def model(self, mini_batch, mini_batch_reversed, mini_batch_mask,
          mini_batch_seq_lengths, annealing_factor=1.0):
 
    # this is the number of time steps we need to process in the mini-batch
    T_max = mini_batch.size(1)
 
    # register all PyTorch (sub)modules with pyro
    # this needs to happen in both the model and guide
    pyro.module("dmm", self)
 
    # set z_prev = z_0 to setup the recursive conditioning in p(z_t | z_{t-1})
    z_prev = self.z_0.expand(mini_batch.size(0), self.z_0.size(0))
 
    # we enclose all the sample statements in the model in a plate.
    # this marks that each datapoint is conditionally independent of the others
    with pyro.plate("z_minibatch", len(mini_batch)):
        # sample the latents z and observed x's one time step at a time
        for t in range(1, T_max + 1):
            # the next chunk of code samples z_t ~ p(z_t | z_{t-1})
            # note that (both here and elsewhere) we use poutine.scale to take care
            # of KL annealing. we use the mask() method to deal with raggedness
            # in the observed data (i.e. different sequences in the mini-batch
            # have different lengths)
 
            # first compute the parameters of the diagonal gaussian
            # distribution p(z_t | z_{t-1})
            z_loc, z_scale = self.trans(z_prev)
 
            # then sample z_t according to dist.Normal(z_loc, z_scale).
            # note that we use the reshape method so that the univariate
            # Normal distribution is treated as a multivariate Normal
            # distribution with a diagonal covariance.
            with poutine.scale(None, annealing_factor):
                z_t = pyro.sample("z_%d" % t,
                                  dist.Normal(z_loc, z_scale)
                                      .mask(mini_batch_mask[:, t - 1:t])
                                      .to_event(1))
 
            # compute the probabilities that parameterize the bernoulli likelihood
            emission_probs_t = self.emitter(z_t)
            # the next statement instructs pyro to observe x_t according to the
            # bernoulli distribution p(x_t|z_t)
            pyro.sample("obs_x_%d" % t,
                        dist.Bernoulli(emission_probs_t)
                            .mask(mini_batch_mask[:, t - 1:t])
                            .to_event(1),
                        obs=mini_batch[:, t - 1, :])
            # the latent sampled at this time step will be conditioned upon
            # in the next time step so keep track of it
            z_prev = z_t

注意すべき最初のことは model() は多くの引数を取ることです。当面は単に mini_batch と mini_batch_mask を見ましょう。mini_batch は 3 次元 tensor で、最初の次元はバッチ次元で、2 番目の次元は時間次元で、そして最後の次元は特徴 (私達のケースでは 88-次元) です。コードをスピードアップするため、モデルを実行するときはいつでもシークエンスの mini-batch 全体を処理していきます (i.e. ベクトル化を活用していきます)。

これは思慮深いです、何故ならば私達のモデルは単一の観測シークエンスに渡り暗黙的に定義されているからです。シークエンスのセットの確率は個々のシークエンス確率の積により単に与えられます。換言すれば、モデルのパラメータが与えられたときシークエンスは条件的に独立です。

このベクトル化は何某かの複雑さを導入します、何故ならばシークエンスは異なる長さである可能性があるからです。これが mini_batch_mask が入ってくるところです。mini_batch_mask は次元 mini_batch_size x T_max の 2 次元 0/1 マスクで、ここで T_max は mini-batch の任意のシークエンスの最大長です。これは mini_batch のどの部分が有効な観測であるかをエンコードします。

そこで私達が行なう最初のことは T_max を把握することです : モデルを少なくともこの多くの時間ステップのために展開しなければなりません。これは多くの「無駄な (= wasted)」計算という結果になることに注意してください、何故ならばシークエンスの幾つかは T_max よりも短いからです、しかしこれはベクトル化に伴う大きなスピードアップのために払う小さな price です。「無駄な」計算がモデル計算を「汚染 (= pollute)」しないことを確実にする必要があるだけです。時間ステップ $t$ に適切なマスクを mask メソッド (これはマスキングを必要とする分布上で作用します) に渡すことによりこれを成します。

最後に、行 pyro.module(“dmm”, self) はモデルの各パラメータのための多くの pyro.param ステートメントと同値です。これは Pyro にどのパラメータがモデルの一部であるかを知らしめます。丁度 sample ステートメントのようにモジュールに一意な名前を与えます。この名前はモデルの Parameter の名前に組み込まれます。KL annealing 因子の議論は後に委ねます。

推論

この時点でモデルを完全に指定しました。次のステップは推論のためにセットアップすることです。イントロダクションで述べたように、私達の推論ストラテジーは変分推論になります (イントロダクションのためには SVI パート I 参照)。そして私達の次のタスクは深層マルコフモデルで推論を行なうために適切な変分分布の族を構築することです。けれども、model() を実装した方法について何も私達を変分推論と結び付けていない ことは強調する価値があります。原理的には Pyro で利用可能な任意のストラテジーを利用できるでしょう。例えば、この特定のコンテキストで逐次モンテカルロのある亜種を利用することを想像できるでしょう (これは現在 Pyro でサポートされていませんが)。

ガイド

ガイド (i.e. 変分分布) の目的は正確な事後分布 $p({\bf z}_{1:T}|{\bf x}_{1:T})$ への (パラメータ化された) 近似を提供することです。実際には、明示的にすべき暗黙的な仮定がここにありますので、一歩戻りましょう。私達のデータセット $\mathcal{D}$ が $N$ シークエンス $\{ {\bf x}_{1:T_1}^1, {\bf x}_{1:T_2}^2, …, {\bf x}_{1:T_N}^N \}$ から成ると仮定します。そして私達が実際に関心がある事後分布は $p({\bf z}_{1:T_1}^1, {\bf z}_{1:T_2}^2, …, {\bf z}_{1:T_N}^N | \mathcal{D})$ で与えられます、i.e 総ての $N$ シークエンスのための潜在変数を推論することを望みます。小さい $N$ に対してでさえもこれは非常に高次元な分布で、これは指定する非常に多くのパラメータを必要とします。特にこの形式で事後分布を直接パラメータ化するのであれば、必要なパラメータの数は (少なくとも) $N$ の線形で増大します。データセットのサイズによるこの扱いにくい増大を回避する一つの方法は amortization です (SVI パート II の類似の議論参照)。

余談 : Amortization

これは次のように動作します。データセットの各シークエンスのために変分パラメータを導入する代わりに、単一のパラメトリック関数 $f({\bf x}_{1:T})$ を学習して、形式 $\prod_{n=1}^N q({\bf z}_{1:T_n}^n | f({\bf x}_{1:T_n}^n))$ を持つ変分分布で作業していきます。関数 $f(\cdot)$ — これは基本的には与えられた観測シークエンスをそのシークエンスにあつらえた変分パラメータのセットにマップします — は事後分布を正確に捕捉するために十分にリッチである必要がありますが、今は非常識な数の変分パラメータを導入しなければならないことなく巨大なデータセットを扱えます。

そして私達のタスクは関数 $f(\cdot)$ を構築することです。私達のケースでは可変長シークエンスをサポートする必要がありますので、ループで $f(\cdot)$ が RNN を持つことは自然であるだけです。 $f(\cdot)$ を構成する様々なコンポーネント部分を詳細に見る前に、基本構造をエンコードする計算グラフを見ましょう :

Figure 2: T=3 時間ステップのために展開されたガイド

図のボトムで 3 つの観測のシークエンスを持ちます。これらの観測は RNN により消費されます、これは観測を右から左へと読みそして 3 つの隠れ状態 $\{ {\bf h}_1, {\bf h}_2,{\bf h}_3\}$ を出力します。この計算は任意の潜在変数からサンプリングする前に行なわれることに注意してください。次に、隠れ状態の各々は Combiner モジュールに供給されます、そのジョブは条件付き分布 $q({\bf z}_t | {\bf z}_{t-1}, {\bf x}_{t:T})$ の平均と (対角ガウス分布により与えられる) 共分散を出力することです。(丁度モデルでのように、ガイドの ${\bf z}_{1:T}$ の条件付き構造は ${\bf z}_t$ から前向きに時間を合わせてサンプリングするようなものです。) RNN 隠れ状態に加えて、Combiner はまた入力として前の時間ステップからの潜在確率変数を取ります、 $t=1$ を除いて、そこではそれは代わりに訓練可能な (変分) パラメータ ${\bf z}_0^{\rm{q}}$ を取ります。

余談: ガイド構造

何故 RNN を右から左へ観測を消費するようにセットアップするのでしょう？何故左から右ではないのでしょう？この選択により条件分布 $q({\bf z}_t |…)$ は 2 つのことに依拠します :

前の時間ステップからの潜在変数 ${\bf z}_{t-1}$ ; そして
観測 ${\bf x}_{t:T}$ , i.e. 総ての未来の観測と一緒の現在の観測

自由に他の選択ができます ; 要求されることの総てはガイドが autograd と上手くやる正しく正規化された分布であることです。この特定の選択は真の事後分布の依存性構造により動機付けられています : 詳細な議論についてはリファレンス [1] を見てください。簡潔に言えば、例えば、観測のシークエンス全体の上で条件付けられるかもしれない一方で、モデルのマルコフ構造ゆえに前の観測 ${\bf x}_{1:t-1}$ について知る必要がある総ては ${\bf z}_{t-1}$ によりカプセル化しています。より多くのものの上で条件付けられるかもしれませんが、必要ありません ; そしてそれを行なうことは多分学習信号を薄める傾向にあります。従って RNN を右から左へと実行することはこの特定のモデルのためには最も自然な選択です。

コンポーネント部分を詳細に見ましょう。最初に、Combiner モジュールです :

class Combiner(nn.Module):
    """
    Parameterizes q(z_t | z_{t-1}, x_{t:T}), which is the basic building block
    of the guide (i.e. the variational distribution). The dependence on x_{t:T} is
    through the hidden state of the RNN (see the pytorch module `rnn` below)
    """
    def __init__(self, z_dim, rnn_dim):
        super().__init__()
        # initialize the three linear transformations used in the neural network
        self.lin_z_to_hidden = nn.Linear(z_dim, rnn_dim)
        self.lin_hidden_to_loc = nn.Linear(rnn_dim, z_dim)
        self.lin_hidden_to_scale = nn.Linear(rnn_dim, z_dim)
        # initialize the two non-linearities used in the neural network
        self.tanh = nn.Tanh()
        self.softplus = nn.Softplus()
 
    def forward(self, z_t_1, h_rnn):
        """
        Given the latent z at at a particular time step t-1 as well as the hidden
        state of the RNN h(x_{t:T}) we return the mean and scale vectors that
        parameterize the (diagonal) gaussian distribution q(z_t | z_{t-1}, x_{t:T})
        """
        # combine the rnn hidden state with a transformed version of z_t_1
        h_combined = 0.5 * (self.tanh(self.lin_z_to_hidden(z_t_1)) + h_rnn)
        # use the combined hidden state to compute the mean used to sample z_t
        loc = self.lin_hidden_to_loc(h_combined)
        # use the combined hidden state to compute the scale used to sample z_t
        scale = self.softplus(self.lin_hidden_to_scale(h_combined))
        # return loc, scale which can be fed into Normal
        return loc, scale

このモジュールはモデルの Emitter と GatedTransition と同じ一般的な構造を持ちます。注意すべき唯一のことは Combiner は各時間ステップで 2 つの入力を消費する必要があるので、出力を計算する前にそれは入力を単一の結合された隠れ状態 h_combined に変換することです。

RNN は別にして、今はガイド分布を構築するために必要な構成要素の総てを持ちます。幸い、PyTorch は素晴らしい組込み RNN モジュールを持ちますので、ここで行なう多くの作業は持ちません。後で RNN をどこでインスタンス化するかを見ます。代わりに確率関数 guide() の定義に真っ直ぐに飛び込みましょう。

def guide(self, mini_batch, mini_batch_reversed, mini_batch_mask,
          mini_batch_seq_lengths, annealing_factor=1.0):
 
    # this is the number of time steps we need to process in the mini-batch
    T_max = mini_batch.size(1)
    # register all PyTorch (sub)modules with pyro
    pyro.module("dmm", self)
 
    # if on gpu we need the fully broadcast view of the rnn initial state
    # to be in contiguous gpu memory
    h_0_contig = self.h_0.expand(1, mini_batch.size(0),
                                 self.rnn.hidden_size).contiguous()
    # push the observed x's through the rnn;
    # rnn_output contains the hidden state at each time step
    rnn_output, _ = self.rnn(mini_batch_reversed, h_0_contig)
    # reverse the time-ordering in the hidden state and un-pack it
    rnn_output = poly.pad_and_reverse(rnn_output, mini_batch_seq_lengths)
    # set z_prev = z_q_0 to setup the recursive conditioning in q(z_t |...)
    z_prev = self.z_q_0.expand(mini_batch.size(0), self.z_q_0.size(0))
 
    # we enclose all the sample statements in the guide in a plate.
    # this marks that each datapoint is conditionally independent of the others.
    with pyro.plate("z_minibatch", len(mini_batch)):
        # sample the latents z one time step at a time
        for t in range(1, T_max + 1):
            # the next two lines assemble the distribution q(z_t | z_{t-1}, x_{t:T})
            z_loc, z_scale = self.combiner(z_prev, rnn_output[:, t - 1, :])
            z_dist = dist.Normal(z_loc, z_scale)
 
            # sample z_t from the distribution z_dist
            with pyro.poutine.scale(None, annealing_factor):
                z_t = pyro.sample("z_%d" % t,
                                  z_dist.mask(mini_batch_mask[:, t - 1:t])
                                        .to_event(1))
            # the latent sampled at this time step will be conditioned
            # upon in the next time step so keep track of it
            z_prev = z_t

guide() の高位構造は model() に非常に良く似ています。最初にモデルとガイドが同じ引数を取ることに注意してください : これは Pyro のモデル/ガイドのための一般的な要件です。モデルでのように、Pyro で総てのパラメータを登録する pyro.module への呼び出しがあります。また、for ループは model() のものと同じ構造を持ちますが、ガイドは潜在変数をサンプリングする必要があるだけ という違いがあります (obs キーワードを持つ sample ステートメントはありません)。最後に、ガイドの潜在変数の名前はモデルのそれらに正確に一致することに注意してください。これは Pyro が確率変数を正しくアラインすることをどのように知るかです。

RNN ロジックは PyTorch ユーザに馴染みがあるはずですが、それを素早く通り抜けましょう。最初に、RNN の初期状態 h_0 を準備します。それから forward 呼び出しを通して RNN を起動します ; 結果としての tensor rnn_output はミニバッチ全体のための隠れ状態を含みます。RNN が右から左へ観測を消費することを望みますので、RNN への入力は mini_batch_reversed です、これは総てのシークエンスが反対の時間的順序で動作する mini_batch のコピーです。更に、RNN が可変長シークエンスを扱えるように mini_batch_reversed は PyTorch rnn.pack_padded_sequence にラップされています。潜在空間のサンプリングを通常の時間的順序で行ないますので、Combiner に正しくアラインされて順序付けられた RNN 隠れ状態を供給できるように rnn_output の隠れ状態シークエンスを反対にするヘルパー関数 pad_and_reverse を使用します。このヘルパー関数はまた rnn_output をアンパックしますのでそれはもはや PyTorch rnn.pack_padded_sequence の形式ではありません。

モデルとガイドを PyTorch モジュールとしてパッケージ化する

この段階で、推論へ進む準備ができました。しかしそれを行なう前にモデルとガイドを単一の PyTorch モジュールとしてどのようにパッケージ化したかを素早く調べましょう。これは一般に良い実践です、特により巨大なモデルに対しては。

class DMM(nn.Module):
    """
    This PyTorch Module encapsulates the model as well as the
    variational distribution (the guide) for the Deep Markov Model
    """
    def __init__(self, input_dim=88, z_dim=100, emission_dim=100,
                 transition_dim=200, rnn_dim=600, rnn_dropout_rate=0.0,
                 num_iafs=0, iaf_dim=50, use_cuda=False):
        super().__init__()
        # instantiate pytorch modules used in the model and guide below
        self.emitter = Emitter(input_dim, z_dim, emission_dim)
        self.trans = GatedTransition(z_dim, transition_dim)
        self.combiner = Combiner(z_dim, rnn_dim)
        self.rnn = nn.RNN(input_size=input_dim, hidden_size=rnn_dim,
                          nonlinearity='relu', batch_first=True,
                          bidirectional=False, num_layers=1, dropout=rnn_dropout_rate)
 
        # define a (trainable) parameters z_0 and z_q_0 that help define
        # the probability distributions p(z_1) and q(z_1)
        # (since for t = 1 there are no previous latents to condition on)
        self.z_0 = nn.Parameter(torch.zeros(z_dim))
        self.z_q_0 = nn.Parameter(torch.zeros(z_dim))
        # define a (trainable) parameter for the initial hidden state of the rnn
        self.h_0 = nn.Parameter(torch.zeros(1, 1, rnn_dim))
 
        self.use_cuda = use_cuda
        # if on gpu cuda-ize all pytorch (sub)modules
        if use_cuda:
            self.cuda()
 
    # the model p(x_{1:T} | z_{1:T}) p(z_{1:T})
    def model(...):
 
        # ... as above ...
 
    # the guide q(z_{1:T} | x_{1:T}) (i.e. the variational distribution)
    def guide(...):
 
        # ... as above ...

既にモデルとガイドを調べましたので、ここでの焦点はコンストラクタ上にあります。最初に私達のモデルとガイドで使用する 4 つの PyTorch モジュールをインスタンス化します。モデル側: Emitter と GatedTransition。ガイド側: Combiner と RNN。

次にRNN の初期状態そして z_0 と z_q_0 のための PyTorch パラメータを定義します、これらは非存在確率変数 $\bf z_0$ の代わりにそれぞれ self.trans と self.combiner に供給されます。

ここで強調する重要な点はこれらのモジュールとパラメータの総ては DMM (それ自身は nn.Module から継承しています) の属性であることです。これはこれらがモジュールに属すように自動的に登録されるという結論を持ちます。従って、例えば、DMM のインスタンス上で parameters() を呼び出すとき、PyTorch は総ての関連パラメータを返すことを知ります。それは、model() と guide() で pyro.module(“dmm”, self) を起動するとき、モデルとガイドの両者の総てのパラメータは Pyro で登録されることを意味します。最後に、それは GPU 上で実行している場合、cuda() への呼び出しは総てのパラメータを GPU メモリに移動することを意味します。

確率的変分推論

手元のモデルとガイドで、最終的に推論を行なう準備ができました。完全な実験スクリプトに関与する full ロジックを見る前に、最初に単一の勾配ステップをどのように取るかを見ましょう。最初に DMM のインスタンスをインスタンス化して optimizer をセットアップします。

# instantiate the dmm
dmm = DMM(input_dim, z_dim, emission_dim, transition_dim, rnn_dim,
          args.rnn_dropout_rate, args.num_iafs, args.iaf_dim, args.cuda)
 
# setup optimizer
adam_params = {"lr": args.learning_rate, "betas": (args.beta1, args.beta2),
               "clip_norm": args.clip_norm, "lrd": args.lr_decay,
               "weight_decay": args.weight_decay}
optimizer = ClippedAdam(adam_params)

ここで勾配クリッピングを含む Adam optimizer を使用しています。これはリカレント・ニューラルネットワークを訓練するときに起きるかもしれない問題の幾つかを移動します (= migrate) (e.g. 勾配消失/爆発)。次に推論アルゴリズムをセットアップします。

# setup inference algorithm
svi = SVI(dmm.model, dmm.guide, optimizer, Trace_ELBO())

推論アルゴリズム SVI は、この場合 ELBO で与えられる目的関数上で勾配ステップを取るために確率的勾配推定器を利用します。名前が示すように、ELBO は対数エビデンスへの下限です : $\log p(\mathcal{D})$ 。ELBO を最大化する勾配ステップを取るとき、ガイド $q(\cdot)$ を正確な事後分布に近付けます。引数 Trace_ELBO() は勾配推定器のバージョンをコンストラクトします、これはモデルとガイドの依存性構造へのアクセスを必要としません。モデルの総ての潜在変数は再パラメータ可能なので、これは私達のユースケースのために適切な勾配推定器です。(それはまたデフォルトオプションです。)

dmm.model と dmm.guide の幾つかの引数を準備したと仮定して、勾配ステップを取ることは次を呼び出すことで成されます :

1	`svi.step(mini_batch, ...)`

That’s all there is to it!

そうですね、完全ではありません。これは私達の推論アルゴリズムの主要ステップですが、依然としてミニバッチ、評価等の準備を伴う完全な訓練ループを実装する必要があります。この類のロジックは深層学習者には良く知られているでしょうが、それが PyTorch/Pyro でどのようなものかを見ましょう。

最適化の黒魔術

実際には、訓練の中心部に到達する前に、少し時間を取り (私達がセットアップした) 最適化問題について少し考えましょう。

私達は非線形モデルの Bayesian 推論を特定の最適化問題のための — 困難な問題 — 高次元潜在空間と交換しました。私達自身をからかわないようにしましょう、この最適化問題もまた非常に困難です。何故でしょう？幾つかの理由を調べましょう :

(それに渡り) 最適化しているパラメータの空間は非常に高次元です (それは定義した総てのニューラルネットワークの総ての重みを含みます)。
目的関数 (ELBO) は解析的に計算できません。そのためパラメータ更新は noisy モンテカルロ勾配推定に従っていきます。
データ・サブサンプリングは偶然性の追加ソースとしてサーブします : 私達が望んだ場合でさえも、一般にはデータセット全体に渡り定義された ELBO 上で勾配ステップを取れません (実際には私達の特定のケースではデータセットはそれほど大きくありませんが、それは無視しましょう)。
ループで私達が持つ総てのニューラルネットワークと非線形が与えられたとき、(確率的) 損失面は非常に非自明です。

結論は、ELBO の合理的な (局所) 最適化条件 (= optima) を見つけていく場合、どのように最適化を行なうかを決定することに少し注意した方が良いことです。これは採用するかもしれない幾つかのストラテジーの総てを議論する時間や場所ではありませんが、学習ハイパーパラメータ (学習率、ミニバッチ・サイズ等) の良い or 悪い選択がどれほど決定的であり得るかを強調することは重要です。

先に進む前に、私達が利用している特定の最適化ストラテジーをより詳しく議論しましょう : KL アニーリングです。私達のケースでは ELBO は 2 つの項の合計です : 期待対数尤度項 (これはモデル fit を測定します) と KL ダイバージェンス項 (これは近似事後分布を正則化するために役立ちます) の合計です :

$\rm{ELBO} = \mathbb{E}_{q({\bf z}_{1:T})}[\log p({\bf x}_{1:T}|{\bf z}_{1:T})] – \mathbb{E}_{q({\bf z}_{1:T})}[ \log q({\bf z}_{1:T}) – \log p({\bf z}_{1:T})]$

後者の項は非常に強い regularizer で、訓練の初期にはそれは、多くの悪い局所最適化条件を含む損失面の領域に有利に働く傾向を持ちます。これらの悪い局所最適化条件を回避する一つのストラテジーは、これはリファレンス [1] でも採用されていますが、KL ダイバージェンス項を 0 と 1 の範囲のスカラー annealing_factor で乗算することにより焼きなますことです :

$\mathbb{E}_{q({\bf z}_{1:T})}[\log p({\bf x}_{1:T}|{\bf z}_{1:T})] – \rm{annealing\_factor} \times \mathbb{E}_{q({\bf z}_{1:T})}[ \log q({\bf z}_{1:T}) – \log p({\bf z}_{1:T})]$

アイデアは、訓練の過程の間に annealing_factor が 0 かその近くの初期値から 1.0 の最終値にゆっくりと上昇することです。アニーリング・スケジュールは任意です ; 下では単純な線形スケジュールを使用します。コードの視点からは、対数尤度を適切なアニーリング因子でスケールするためにモデルとガイドの潜在 sample ステートメントの各々を pyro.poutine.scale で囲みます。

最後に、ここで説明される DMM 実装とリファレンス [1] で使用されるものの間に主な違いは、(私達がモンテカルロ推定に依拠するところを) それらは 2 つのガウス分布の間の KL ダイバージェンスのための解析的式 (= analytic formula) を利用していることを述べておくべきでしょう。これは ELBO のより低い分散勾配推定に繋がり、これは訓練を少し容易にします。私達は依然としてこの解析的代用を行なうことなくモデルを訓練できますが、訓練は高い分散ゆえに多分幾分長くかかります。Pyro での解析的 KL ダイバージェンスのサポートは将来的に追加することを計画しています。

データローディング、訓練そして評価

最初にデータをロードします。訓練データセットには 229 シークエンスがあり、各々は ~60 時間ステップの平均長を持ちます。

jsb_file_loc = "./data/jsb_processed.pkl"
data = pickle.load(open(jsb_file_loc, "rb"))
training_seq_lengths = data['train']['sequence_lengths']
training_data_sequences = data['train']['sequences']
test_seq_lengths = data['test']['sequence_lengths']
test_data_sequences = data['test']['sequences']
val_seq_lengths = data['valid']['sequence_lengths']
val_data_sequences = data['valid']['sequences']
N_train_data = len(training_seq_lengths)
N_train_time_slices = np.sum(training_seq_lengths)
N_mini_batches = int(N_train_data / args.mini_batch_size +
                     int(N_train_data % args.mini_batch_size > 0))

このデータセットのため典型的には 20 の mini_batch_size を使用しますので、エポック毎に 12 ミニバッチがあります。次に関数 process_minibatch を定義します、これは訓練のためにミニバッチを準備して勾配ステップを取ります :

def process_minibatch(epoch, which_mini_batch, shuffled_indices):
    if args.annealing_epochs > 0 and epoch < args.annealing_epochs:
        # compute the KL annealing factor appropriate
        # for the current mini-batch in the current epoch
        min_af = args.minimum_annealing_factor
        annealing_factor = min_af + (1.0 - min_af) * \
            (float(which_mini_batch + epoch * N_mini_batches + 1) /
             float(args.annealing_epochs * N_mini_batches))
    else:
        # by default the KL annealing factor is unity
        annealing_factor = 1.0
 
    # compute which sequences in the training set we should grab
    mini_batch_start = (which_mini_batch * args.mini_batch_size)
    mini_batch_end = np.min([(which_mini_batch + 1) * args.mini_batch_size,
                             N_train_data])
    mini_batch_indices = shuffled_indices[mini_batch_start:mini_batch_end]
    # grab the fully prepped mini-batch using the helper function in the data loader
    mini_batch, mini_batch_reversed, mini_batch_mask, mini_batch_seq_lengths \
        = poly.get_mini_batch(mini_batch_indices, training_data_sequences,
                              training_seq_lengths, cuda=args.cuda)
    # do an actual gradient step
    loss = svi.step(mini_batch, mini_batch_reversed, mini_batch_mask,
                     mini_batch_seq_lengths, annealing_factor)
    # keep track of the training loss
    return loss

最初に (前に説明した線形スケジュールに従って) ミニバッチに適切な KL アニーリング因子を計算します。それからミニバッチ・インデックスを計算します、これをヘルパー関数 get_mini_batch に渡します。このヘルパー関数は多くの異なることに対応します :

それはシークエンス長で各ミニバッチをソートします。
それは逆時間順序のミニバッチのコピーを得るために他のヘルパー関数を呼び出します。
それは rnn.pack_padded_sequence の各逆のミニバッチをパックします、これはそれから RNN により摂取される準備ができます。
それは GPU 上であれば総ての tensor を cuda 化します。
それはミニバッチのための適切な 0/1 マスクを得るために他のヘルパー関数を呼び出します。

それから get_mini_batch() の総ての戻り値を elbo.step(...) にパイプします。elbo の勾配推定器の構築の間これらの引数は更に model(...) と guide(...) にパイプされることを思い出してください。最後に、float を返します、これはそのミニバッチのための損失の noisy 推定です。

今は訓練ループの主要部分のために必要な構成要素の総てを持ちます :

times = [time.time()]
for epoch in range(args.num_epochs):
    # accumulator for our estimate of the negative log likelihood
    # (or rather -elbo) for this epoch
    epoch_nll = 0.0
    # prepare mini-batch subsampling indices for this epoch
    shuffled_indices = np.arange(N_train_data)
    np.random.shuffle(shuffled_indices)
 
    # process each mini-batch; this is where we take gradient steps
    for which_mini_batch in range(N_mini_batches):
        epoch_nll += process_minibatch(epoch, which_mini_batch, shuffled_indices)
 
    # report training diagnostics
    times.append(time.time())
    epoch_time = times[-1] - times[-2]
    log("[training epoch %04d]  %.4f \t\t\t\t(dt = %.3f sec)" %
        (epoch, epoch_nll / N_train_time_slices, epoch_time))

各エポックの最初で訓練データをポイントするインデックスをシャッフルします。それから進みにつれて訓練損失を累積し、各ミニバッチを訓練セット全体を通り抜けるまで処理します。最後に幾つかの診断情報を報告します。

損失を訓練セットの時間スライスの総数で正規化することに注意してください (これはリファレンス [1] と比較することを可能にします)。

評価

この訓練ループは依然としてどのような種類の評価診断も欠いています。Let’s fix that. 最初に評価のための検証とテストデータを準備する必要があります。検証とデータセットはそれらをメモリに容易に収められるに十分小さいので、各データセットをバッチ wise に処理していきます (i.e. データセットをミニバッチに分解していきません)。[余談: この時点で読者は訓練セットのためと同じことを何故行なわないのか尋ねるかもしれません。理由は、データ・サブサンプリングによる偶然性はしばしば最適化の間に有利であるからです : 特に局所最適化条件を回避する助けとなれます。]そして実際に、ELBO の noisy 推定を得るため、マルチサンプル推定を計算していきます。これを行なう最も単純は方法は以下になるでしょう :

1	`val_loss` `=` `svi.evaluate_loss(val_batch, ..., num_particles=5)`

これはけれども、5 つの反復を伴う明示的な for ループを含みます。私達の特定のモデルのためには、より良くやれて計算全体をベクトル化できます。Pyro で現在これを行なう唯一の方法はデータ n_eval_samples を明示的に何回も複製することです。これが私達が従うストラテジーです :

# package repeated copies of val/test data for faster evaluation
# (i.e. set us up for vectorization)
def rep(x):
    return np.repeat(x, n_eval_samples, axis=0)
 
# get the validation/test data ready for the dmm: pack into sequences, etc.
val_seq_lengths = rep(val_seq_lengths)
test_seq_lengths = rep(test_seq_lengths)
val_batch, val_batch_reversed, val_batch_mask, val_seq_lengths = poly.get_mini_batch(
    np.arange(n_eval_samples * val_data_sequences.shape[0]), rep(val_data_sequences),
    val_seq_lengths, cuda=args.cuda)
test_batch, test_batch_reversed, test_batch_mask, test_seq_lengths = \
    poly.get_mini_batch(np.arange(n_eval_samples * test_data_sequences.shape[0]),
                        rep(test_data_sequences),
                        test_seq_lengths, cuda=args.cuda)

テストと検証データが完全に今準備され、評価を行なうヘルパー関数を定義します :

def do_evaluation():
    # put the RNN into evaluation mode (i.e. turn off drop-out if applicable)
    dmm.rnn.eval()
 
    # compute the validation and test loss
    val_nll = svi.evaluate_loss(val_batch, val_batch_reversed, val_batch_mask,
                                 val_seq_lengths) / np.sum(val_seq_lengths)
    test_nll = svi.evaluate_loss(test_batch, test_batch_reversed, test_batch_mask,
                                  test_seq_lengths) / np.sum(test_seq_lengths)
 
    # put the RNN back into training mode (i.e. turn on drop-out if applicable)
    dmm.rnn.train()
    return val_nll, test_nll

elbo の evaluate_loss メソッドを単純に呼び出します、これは step() と同じ引数を取ります、つまりモデルとガイドに渡される引数です。dropout を構成するために RNN を評価モードに入れて出さなければならないことに注意してください。今は do_evaluation() を訓練ループに突き刺すことができます ; 詳細はソースコードを見てください。

結果

私達の実装が合理的な結果を与えることを確かめましょう。サニティチェックとしてリファレンス [1] で報告されている数値を利用できます。同じデータセットそして類似のモデル/ガイドのセットアップ (潜在空間の次元、RNN の隠れユニット数等) についてテストセット上 6.93 の正規化された負対数尤度 (NLL) を報告しています (より低いほうがより良いです)。これは 6.87 の私達の結果と比較されます。これらの数値は非常に同じボールパークにあり、安心を与えてくれます。少なくともこのデータセットについて、KL ダイバージェンスのための解析的式の使用は学習モデルの品質を劣化させないようです (上で議論したように、訓練は多分幾分長くかかりますが)。

Figure 3: サンプル訓練実行について訓練が進むときのテストセット NLL の進捗

図では (保守的な学習率で) 単一のサンプル実行について訓練の間にテスト NLL がどのように進捗するかを示しています。進捗の殆どは最初の 3000 エポックかそこらの間のもので、訓練を更に長く続ければ何某かの周辺的なゲインを持つでしょう。GeForce GTX 1080 上、5000 エポックはおよそ 20 時間かかります。

num_iafs	test NLL
0	6.87
1	6.82
2	6.80

最後にガイドのための結果を報告します (詳細は次のセクションで見つけられます)。

Finally, we also report results for guides with normalizing flows in the mix (details to be found in the next section).

実際には、それらは同じモデル/ガイドに対して 2 つの数値 — 6.93 と 7.03 — を報告するようです、そして 2 つの報告された数字がどのように異なるのかは全く明瞭ではありません。

おまけと他の改良

自己回帰 (= autoregressive) フローを反対にする

確率的プログラミング言語について素晴らしいことの一つはそれはモジュール化を奨励することです。DMM のコンテキストで例を示しましょう。mix に normalizing フローを追加することにより私達の変分分布をリッチにしていきます (議論のためにはリファレンス [2] 参照)。これはコードの追加 4 行のコストがかかるだけです！

最初に、DMM コンストラクタで以下を追加します :

iafs = [AffineAutoregressive(AutoRegressiveNN(z_dim, [iaf_dim])) for _ in range(num_iafs)]
self.iafs = nn.ModuleList(iafs)

これは AffineAutoregressive 型の num_iafs だけ多くの bijective 変換をインスタンス化します (リファレンス [3,4] 参照) ; 各 normalizing フローは iaf_dim だけ多くの隠れユニットを持ちます。それから nn.ModuleList で normalizing フローをバンドルします ; これは単に nn.Modules のリストをパッケージ化する PyTorch の方法です。次に、ガイドで以下の行を追加します :

if self.iafs.__len__() > 0:
    z_dist = TransformedDistribution(z_dist, self.iafs)

ここでベース分布 z_dist を取っています、これは私達のケースでは条件付きガウス分布で、TransformedDistribution コンストラクタを使用してそれを非ガウシアン分布に変換します、これはコンストラクションによりベース分布よりもリッチです。Voila!

チェックポインティング

訓練ループでの壊滅的な失敗からリカバーすることを望む場合、追跡する必要がある 2 つの種類の状態があります。最初はモデルとガイドの幾つかのパラメータです。2 番目は optimizer の状態です (e.g. Adam ではこれは各パラメータのための最近の勾配推定の実行平均を含みます)。

Pyro では、パラメータは ParamStore で総て見つけられます。けれども、PyTorch はまた nn.Module の parameters() メソッドを通してそれらを追跡します。従ってモデルとガイドのパラメータをセーブできる一つの単純な方法は torch.save() と連携して dmm の state_dict() メソッドを利用することです ; 下を見てください。ループで AffineAutoregressive のものを持つ場合にはこれは実際には私達が使える唯一の選択肢です。これは AffineAutoregressive モジュールが PyTorch 用語の「persistent バッファ」と呼ばれるものを含むからです。これらはパラメータではない状態を運ぶものです。nn.Module の state_dict() と load_state_dict() メソッドはバッファをどのように正しく扱うかを知っています。

optimizer の状態をセーブするには、pyro.optim.PyroOptim の内部の機能を利用しなければなりません。典型的なユーザは Pyro を利用するとき PyTorch Optimizer と直接には決して相互作用しないことを思い出してください ; 何故ならば任意の確率的プログラムでパラメータは動的に作成できるからで、Pyro は Optimizer を管理する必要があります。私達のケースでは optimizer 状態をセーブすることは optimizer.save() を呼び出すように容易です。ローディング・ロジックは全体的に類推的です。従ってチェックポイントのセーブとロードのためのロジック全体は数行かかるだけです :

# saves the model and optimizer states to disk
def save_checkpoint():
    log("saving model to %s..." % args.save_model)
    torch.save(dmm.state_dict(), args.save_model)
    log("saving optimizer states to %s..." % args.save_opt)
    optimizer.save(args.save_opt)
    log("done saving model and optimizer checkpoints to disk.")
 
# loads the model and optimizer states from disk
def load_checkpoint():
    assert exists(args.load_opt) and exists(args.load_model), \
        "--load-model and/or --load-opt misspecified"
    log("loading model from %s..." % args.load_model)
    dmm.load_state_dict(torch.load(args.load_model))
    log("loading optimizer states from %s..." % args.load_opt)
    optimizer.load(args.load_opt)
    log("done loading model and optimizer states.")

幾つかの最後のコメント

深層マルコフモデルは比較的複雑なモデルです。多声音楽データセットにあつらえた深層学習モデルのバージョンを実装する努力をした今、他に何ができるかを私達自身に問うべきです。異なるシーケンシャルなデータセットを手渡された場合には？最初からやり直さなければならないでしょうか？

Not at all! 確率的プログラミングの美点はそれはモデリングと推論のためにモジュール化アプローチを可能に — そして奨励 — することです。多声音楽モデルを連続値の観測を持つデータセットに適応させるのは観測尤度を変更するほどに単純です。コードの大部分は変更なしで引き継げるかもしれません。これは、ほんの少しの余分の作業で、このチュートリアルのコードが非常に様々な異なるモデルを有効にするために再目的化できかもしれないことを意味します。

Github の完全なコードを見てください。

リファレンス

[1] Structured Inference Networks for Nonlinear State Space Models, Rahul G. Krishnan, Uri Shalit, David Sontag
[2] Variational Inference with Normalizing Flows, Danilo Jimenez Rezende, Shakir Mohamed
[3] Improving Variational Inference with Inverse Autoregressive Flow, Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling
[4] MADE: Masked Autoencoder for Distribution Estimation, Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle
[5] Modeling Temporal Dependencies in High-Dimensional Sequences: Application to Polyphonic Music Generation and Transcription, Boulanger-Lewandowski, N., Bengio, Y. and Vincent, P.

以上

月	火	水	木	金	土	日
					1	2
3	4	5	6	7	8	9
10	11	12	13	14	15	16
17	18	19	20	21	22	23
24	25	26	27	28	29	30
31