[PYTHON] Normalisierung der Strömungstheorie und -implementierung

Einführung

Ich habe Normalizing Flow implementiert, eine der Methoden zur variablen Inferenz, daher werde ich sie hier als Memorandum belassen. Wir beginnen mit der Bayes'schen linearen Regression, damit wir den gesamten Fluss abdecken können, und erklären dann Schritt für Schritt, warum eine variable Inferenz erforderlich ist und wie die Normalisierung des Flusses eine variable Inferenz durchführt. Obwohl es sich um einen normalisierenden Fluss handelt, kann er zusätzlich zur variablen Inferenz auch als Generierungsmodell verwendet werden. GAN und VAE sind als Generierungsmodelle bekannt, aber Normalizing Flow gehört zu einem anderen Flow-basierten Generierungsmodell. Das Papier zur Normalisierung des Flusses ist hier. Der implementierte Code wird auf GitHub veröffentlicht.

Bayesianische lineare Regression

Beginnen wir mit der Geschichte der Bayes'schen linearen Regression. Zusammenfassend wird bei normaler linearer Regression der vorhergesagte Wert $ \ boldsymbol y $ für die beobachteten Daten $ \ boldsymbol x $ ausgegeben, wie unten gezeigt. $ \ Boldsymbol \ theta $ ist ein Parameter und entspricht der Steigung einer geraden Linie bei einfacher Regression. $\boldsymbol y=\boldsymbol \theta^T\boldsymbol x \tag{1}$ Da der tatsächliche Wert jedoch Rauschen enthält, nehmen wir zur Modellierung einer solchen Mehrdeutigkeit eine Gaußsche Verteilung mit einem Durchschnitt von $ \ boldsymbol y $ und einer Kovarianz von $ I $ wie folgt an. (Wenn Sie beispielsweise die Temperatur messen, enthält der Wert, der durch den Gerätefehler des Thermometers erhalten wird, Rauschen. Es handelt sich also um ein Bild, das das Rauschen durch die Streuung der Gaußschen Verteilung ausdrückt.) $p(\boldsymbol y \mid \boldsymbol x，\boldsymbol \theta)=\mathcal{N}(\boldsymbol y\mid \boldsymbol \theta^T\boldsymbol x， I) \tag{2}$ Um hier eine Regressionsvorhersage für nicht beobachtete Daten unter Verwendung der obigen Gleichung zu erstellen, ist es notwendig, den optimalen Parameter $ \ boldsymbol \ theta $ zu bestimmen. Von hier aus verfolgen wir einen Bayes'schen Ansatz, aber Bayes betrachtet den Parameter $ \ boldsymbol \ theta $ als Wahrscheinlichkeitsverteilung. Das heißt, die Wahrscheinlichkeitsverteilung $ p (\ boldsymbol \ theta) des optimalen Parameters $ \ boldsymbol \ theta $ aus den bereits vorhandenen Daten $ \ boldsymbol X ^ {(train)} $ und $ \ boldsymbol y ^ {(train)} $ \ mid \ boldsymbol X ^ {(Zug)}, \ boldsymbol y ^ {(Zug)}) Bestimmen Sie $. (Der Nicht-Bayes-Ansatz findet den optimalen Parameter $ \ boldsymbol \ theta $ als einen deterministischen Wert anstelle einer Wahrscheinlichkeitsverteilung. In diesem Fall wird eine Technik verwendet, die als wahrscheinlichste Schätzung bezeichnet wird, was zu einer Methode mit dem kleinsten Quadrat führt. ) $ P (\ boldsymbol \ theta \ mid \ boldsymbol X ^ {(Zug)}, \ boldsymbol y ^ {(Zug)}) $ wird unter Verwendung des folgenden Bayes'schen Theorems berechnet. $ p(\boldsymbol \theta \mid \boldsymbol X^{(train)}，\boldsymbol y^{(train)})= \cfrac{p(\boldsymbol y^{(train)} \mid \boldsymbol X^{(train)}， \boldsymbol \theta)p(\boldsymbol \theta)}{p(\boldsymbol y^{(train)} \mid \boldsymbol X^{(train)})} \tag{3} $ $ p (\ boldsymbol \ theta \ mid \ boldsymbol X ^ {(Zug)}, \ boldsymbol y ^ {(Zug)}) $ wird als posteriore Verteilung bezeichnet, und sobald die posteriore Verteilung erhalten ist, wird die folgende Formel verwendet. Sie können damit die Regressionsvorhersage $ \ boldsymbol y ^ {(new)} $ für unbekannte Daten $ \ boldsymbol X ^ {(new)} $ berechnen. (Obwohl es kompliziert aussieht, werde ich diese Formeln dieses Mal nicht direkt verwenden, also lesen Sie es bitte leicht durch.) $ p(\boldsymbol y^{(new)}\mid \boldsymbol X^{(new)}， \boldsymbol y^{(train)}， \boldsymbol X^{(train)})= \int p(\boldsymbol y^{(new)}\mid \boldsymbol X^{(new)}，\boldsymbol \theta) p(\boldsymbol \theta \mid \boldsymbol y^{(train)}， \boldsymbol X^{(train)}) d\boldsymbol \theta $ Um auf die Geschichte zurückzukommen, schauen wir uns den Satz von Bayes (3) genauer an, um die hintere Verteilung zu erhalten. Zunächst wird das Molekül $ p (\ boldsymbol y ^ {(train)} \ mid \ boldsymbol X ^ {(train)}, \ boldsymbol \ theta) $ als Wahrscheinlichkeit bezeichnet und wurde bereits durch Gleichung (2) modelliert. Da dies bereits geschehen ist, kann es berechnet werden. In ähnlicher Weise wird $ p (\ boldsymbol \ theta) $, das im Molekül von Gleichung (3) erscheint, als vorherige Verteilung bezeichnet, die vom Designer auf der Grundlage vorheriger Kenntnisse über die Verteilung von $ \ boldsymbol \ theta $ modelliert wird. Ich werde. (Es gibt Regeln für die Auswahl, aber es ist kein Fehler, eine Wahrscheinlichkeitsverteilung wie die Gaußsche Verteilung oder die Gleichverteilung zu wählen.)

Schließlich wird $ p (\ boldsymbol y ^ {(Zug)} \ mid \ boldsymbol X ^ {(Zug)}) , das als Beweis oder Grenzwahrscheinlichkeit bezeichnet wird, aus der folgenden Gleichung berechnet. $ p(\boldsymbol y^{(train)} \mid \boldsymbol X^{(train)})= \int p(\boldsymbol y^{(train)} \mid \boldsymbol X^{(train)}， \boldsymbol \theta)p(\boldsymbol \theta)d\boldsymbol \theta \tag{4} $$ Es ist möglich, diese Gleichung mit einem einfachen Modell wie der linearen Regression analytisch zu berechnen, aber es wird schwierig, das Integral zu berechnen, wenn das Modell kompliziert wird. Die dort verwendete Methode ist die Differentialinferenz. (Es gibt andere Approximationsmethoden wie die Monte-Carlo-Methode.)

Variantenüberlegung

Vermeiden Sie die genaue analytische Analyse der Grenzwahrscheinlichkeit und versuchen Sie, die posteriore Verteilung $ p (\ boldsymbol z \ mid \ boldsymbol x) $, die Sie durch die Wahrscheinlichkeitsverteilung $ q (\ boldsymbol z) $ finden möchten, direkt zu approximieren. Die Idee ist variables Denken. (Um die Notation mit dem Papier abzugleichen, wird der bisher aufgetretene Parameter $ \ boldsymbol \ theta $ als $ \ boldsymbol z $ beschrieben.) Als Beispiel für eine variable Inferenz gibt es ein Verfahren zum Zerlegen der Wahrscheinlichkeitsverteilung, indem die Unabhängigkeit der Verteilung für die ungefähre hintere Verteilung $ q (\ boldsymbol z) $ angenommen wird. Diese Methode wird als mittlere Feldnäherung bezeichnet und nimmt die Form einer ungefähren posterioren Verteilung an, wie in der folgenden Gleichung gezeigt. $q(\boldsymbol z)=\prod_{i} q(\boldsymbol z_i) \tag{5}$ Dann wird eine Optimierung durchgeführt, um diese Wahrscheinlichkeitsverteilung $ q (\ boldsymbol z) $ nahe an die wahre hintere Verteilung $ p (\ boldsymbol z \ mid \ boldsymbol x) $ zu bringen. Die Optimierung der Differentialinferenz beinhaltet eine etwas schwierige Ausdruckstransformation. Um zunächst zu schließen, maximieren Sie die Verlustfunktion von ELBO (Evidence Lower Bound) $ \ mathcal {L} (\ boldsymbol {x}) $ in der folgenden Gleichung. Es optimiert dabei. $ \mathcal{L}(\boldsymbol{x})=\int q(\boldsymbol z)\ln \frac{p(\boldsymbol x， \boldsymbol z)}{q(\boldsymbol z)}d\boldsymbol z \tag{6} $ Von hier wird ELBO abgeleitet, aber Sie können es überspringen. Erstens besteht der ursprüngliche Zweck darin, eine ungefähre hintere Verteilung $ q (\ boldsymbol z) $ zu finden, die den Beweis $ p (\ boldsymbol x) $ maximiert, der die Wahrscheinlichkeit der Datenerzeugung darstellt. Das Maximieren von $ p (\ boldsymbol x) $ entspricht dem Maximieren von $ \ ln p (\ boldsymbol x) $, dem Logarithmus. Betrachten Sie letzteres der Einfachheit halber. Die folgende Formel wird erhalten, indem mit der Berechnung für $ \ ln p (\ boldsymbol x) $ fortgefahren wird.

\begin{eqnarray}
\ln p(\boldsymbol x)=\ln \int p(\boldsymbol x， \boldsymbol z)d\boldsymbol z\\
=\ln \int q(\boldsymbol z)\frac{p(\boldsymbol x， \boldsymbol z)}{q(\boldsymbol z)}d\boldsymbol z\\
\geq \int q(\boldsymbol z)\ln \frac{p(\boldsymbol x， \boldsymbol z)}{q(\boldsymbol z)}d\boldsymbol z \\
= \mathcal{L}(\boldsymbol{x})
\end{eqnarray}

Hier ist die Transformation der Formel in der (3) -Stufe, die jedoch unter Verwendung der Jensen-Ungleichung berechnet wurde. Wie Sie dieser Formel entnehmen können, ist die Untergrenze von $ \ ln p (\ boldsymbol x) $ ELBO $ \ mathcal {L} (\ boldsymbol {x}) $. Ich werde nicht auf Details eingehen, aber das Maximieren des ELBO entspricht dem Minimieren der KL-Divergenz der posterioren Ziel- und ungefähren posterioren Verteilungen, sodass diese Aufgabe auf die ELBO-Maximierung reduziert werden kann.

Zurück zur Geschichte: Nehmen wir an, dass die ungefähre hintere Verteilung $ q (\ boldsymbol z) $ durch Lösen des Optimierungsproblems erhalten werden kann. Erstens gibt es jedoch ein Problem mit der Annahme von Gleichung (5). Durch die Annahme einer Unabhängigkeit wie in Gleichung (5) kann die zwischen den Variablen bestehende Korrelation nicht berücksichtigt werden, und die Wahrscheinlichkeitsverteilung wird schlecht ausdrucksstark, und die wahre hintere Verteilung kann nicht ausgedrückt werden. .. Die Normalisierung des Flusses ist eine der Methoden der variablen Inferenz, die dieses Problem der mangelnden Ausdruckskraft löst.

Normalizing Flow Das Normalisieren des Flusses ist eine komplexe Verteilung $ q_k (\ boldsymbol z_k) durch Überlagerung einer nichtlinearen Transformation $ f $ mit einer Wahrscheinlichkeitsvariablen $ \ boldsymbol z $, die einer einfachen Wahrscheinlichkeitsverteilung (Gaußsche Verteilung usw.) $ q (\ boldsymbol z) $ folgt. ) Die Idee ist, $ zu bekommen. <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/302831/4cfcd80e-18ee-f050-df70-5f59a8e3e8fd.png ", width="100%"> Quelle: L.Weng. "Flow-based Deep Generative Model

Lass uns genauer hinschauen. Die Gleichungen (7) und (8) werden erhalten, indem die nichtlineare Transformation $ f $ auf die Wahrscheinlichkeitsvariable $ \ boldsymbol z $ angewendet wird, die der Wahrscheinlichkeitsverteilung $ q (\ boldsymbol z) $ folgt. Das Differential, das hier auf der rechten Seite erscheint, ist Jacobian. det bedeutet einen Matrixausdruck. Am Ende können Sie jedoch Jacobian und Det manuell berechnen, sodass Sie nicht lange nachdenken müssen. $\boldsymbol z'=f(\boldsymbol z)\tag{7}$ $q'(\boldsymbol z')=q(\boldsymbol z)\left| \det \frac{\partial f}{\partial \boldsymbol z'}\right|^{-1} \tag{8}$ Durch mehrmaliges Wiederholen dieser nichtlinearen Transformation wird die Wahrscheinlichkeitsverteilung komplexer und aussagekräftiger. $ \boldsymbol z_k=f_k\circ \cdots \circ f_2 \circ f_1(\boldsymbol z_0) \tag{9} $ $ q_k(\boldsymbol z_k)=q_0(\boldsymbol z_0)\prod_{k=1}^{K} \left| \det \frac{\partial f_k}{\partial \boldsymbol z_{k-1}}\right|^{-1} \tag{10} $ In Bezug auf die Auswahl der nichtlinearen Funktion $ f_k $ ist es in Ordnung, wenn es sich um eine vollständig einmalige Funktion handelt, dh um eine nichtlineare Funktion, mit der die Umkehrfunktion berechnet werden kann. In der Arbeit werden zwei Arten von nichtlinearen Funktionen $ f_k $ vorgestellt: planare Strömung und radiale Strömung. Hier erklären wir nur die planare Strömung.

Planare Strömung ist die folgende Funktion. $ \ boldsymbol u, \ boldsymbol w und b $ sind Parameter, und ihre Werte werden durch Lernen aktualisiert. $ h $ ist eine Aktivierungsfunktion und verwendet tanh. $ f(\boldsymbol z)=\boldsymbol z + \boldsymbol u h(\boldsymbol w^T \boldsymbol z + b) \tag{11} $ Unter Verwendung der Funktion der obigen Gleichung kann die Jacobi-Matrixgleichung wie folgt berechnet werden. $ \left| \det \frac{\partial f(\boldsymbol z)}{\partial \boldsymbol z}\right| =\left| 1+\boldsymbol u^T \psi(\boldsymbol z)\right| \tag{12} $ $ \ Psi (\ boldsymbol z) $ wird jedoch aus der folgenden Gleichung berechnet. $ \psi(\boldsymbol z)=h'(\boldsymbol w^T \boldsymbol z+b)\boldsymbol w \tag{13} $ Dies ermöglicht es, Gleichung (10) zu berechnen, und die Wahrscheinlichkeitsverteilung nach der Umwandlung kann erhalten werden. $ q_k(\boldsymbol z_k)=q_0(\boldsymbol z_0)\prod_{k=1}^{K} \left| 1+\boldsymbol u_k^T \psi_k(\boldsymbol z_k)\right|^{-1} \tag{13} $ Gleichung (13) ist in planarer Strömung implementiert, aber Sie können sehen, dass die Berechnungskosten sehr gering sind (berechnen Sie einfach das innere Produkt der Vektoren und nehmen Sie die Umkehrung des Absolutwerts).

Als nächstes optimieren wir die Wahrscheinlichkeitsverteilung, die durch Gleichung (13) erhalten wird, um sie näher an die wahre hintere Verteilung zu bringen. Wie bereits erwähnt, wird bei der Differentialinferenz die Optimierung durch Maximieren von ELBO durchgeführt. Wir werden die ELBO-Formel für die Normalisierung des Durchflusses manuell berechnen.

\mathcal{L}(\boldsymbol{x})=
\int q(\boldsymbol z)\ln \frac{p(\boldsymbol x， \boldsymbol z)}{q(\boldsymbol z)}d\boldsymbol z\\
=\mathbb{E}_{q(\boldsymbol z)}[\ln p(\boldsymbol x， \boldsymbol z) - \ln q(\boldsymbol z)]\\
=\mathbb{E}_{q_k(\boldsymbol z_k)}[\ln p(\boldsymbol x， \boldsymbol z_k) - \ln q(\boldsymbol z_k)]\\
\sim \sum_{l=1}^{L}[\ln p(\boldsymbol{x}， \boldsymbol{z_k^{(l)}})-\ln q(\boldsymbol{z_k^{(l)}})]
\tag{14}

Die Summe, die am Ende des Ausdrucks erscheint, ist die Summe der probabilistischen Variablen der L-Stichprobe, die im Mini-Batch enthalten sind (Index $ l). $ Bedeutet die $ l $ -te Stichprobe). Da es bei der Implementierung bequemer ist, zu minimieren als zu maximieren, minimieren Sie die folgende Gleichung, indem Sie ELBO mit Minus multiplizieren. $ -\mathcal{L}(\boldsymbol{x})=\sum_{l=1}^{L}[\ln q(\boldsymbol{z_k^{(l)}}) - \ln p(\boldsymbol{x}， \boldsymbol{z_k^{(l)}})]　\tag{15} $ Dies ist das Ende der Erklärung der Theorie. Der Rest ist die Implementierung.

Ausführungsumgebung

Python 3.7.5 TensorFlow 1.15.0

Implementierung

Implementieren von TensorFlow. Wie bei der Arbeit ist die Problemeinstellung die Gaußsche Verteilung als Anfangsverteilung, und die Zielverteilung wird durch Überlagerung nichtlinearer Transformationen durch planare Strömung reproduziert. (Die linke ist die Gaußsche Verteilung, die die Anfangsverteilung ist, und die Mitte und die rechte sind die Zielverteilungen).

Zuerst werden wir es aus dem planaren Fluss implementieren. Bei der Berechnung der Verlustfunktion wird $ \ ln q_k (\ boldsymbol z_k) $ benötigt, daher wird es hier berechnet.

normalizing_flow.py

class PlanarFlow:
    def __init__(self, dim):
        self.dim = dim
        self.h = lambda x: tf.tanh(x)
        self.h_prime = lambda x: 1 - tf.tanh(x)**2
        self.w = tf.Variable(tf.random.truncated_normal(shape=(1, self.dim)))
        self.b = tf.Variable(tf.zeros(shape=(1)))
        self.u = tf.Variable(tf.random.truncated_normal(shape=(1, self.dim)))
        

    def __call__(self, z, log_q):
        z = z + self.u*self.h(tf.expand_dims(tf.reduce_sum(z*self.w, -1), -1) + self.b)
        psi = self.h_prime(tf.expand_dims(tf.reduce_sum(z*self.w, -1), -1) + self.b)*self.w
        det_jacob = tf.abs(1 + tf.reduce_sum(psi*self.u, -1))
        log_q = log_q -  tf.log(1e-7 + det_jacob)
        return z, log_q

Als nächstes wird ein Normalisierungsfluss konstruiert, indem K planare Flüsse gestapelt werden.

normalizing_flow.py

class NormalizingFlow:
    def __init__(self, K, dim):
        self.K = K
        self.dim = dim
        self.planar_flow = [PlanarFlow(self.dim) for i in range(self.K)]
        

    def __call__(self, z_0,log_q_0):
        z, log_q = self.planar_flow[0](z_0,log_q_0)
        for pf in self.planar_flow[1:]:
            z, log_q = pf(z, log_q)
        return z, log_q

Dann wird die Verlustfunktion (15) wie folgt berechnet. Hier berechnet target_density die Zielwahrscheinlichkeitsverteilung. Hier wird auch die Optimierungsmethode festgelegt. Verwenden Sie Adam, um die Optimierung durchzuführen und die Verlustfunktion zu minimieren. Definieren Sie auch eine Funktion, um den Platzhalter zu erhalten.

normalizing_flow.py

def calc_loss(z_k, log_q_k, target_density):
    log_p = tf.log(target_density.calc_prob_tf(z_k)+1e-7)
    loss = tf.reduce_mean(log_q_k - log_p, -1)
    return loss


def get_train(loss):
    return tf.train.AdamOptimizer().minimize(loss)

def get_placeholder():
    z_0 = tf.placeholder(tf.float32, shape=[None, 2])
    log_q_0 = tf.placeholder(tf.float32, shape=[None])
    return z_0,log_q_0

Erstellen Sie ein Berechnungsdiagramm mit den oben genannten Klassen und Funktionen.

main.py

normalizing_flow = NormalizingFlow(K=16, dim=2)

z_0,log_q_0 = get_placeholder()
z_k, log_q_k = normalizing_flow(z_0,log_q_0)
loss = calc_loss(z_k, log_q_k, target_density)
train = get_train(loss)

Das Lernen wird mit der stochastischen Gradientenabstiegsmethode durchgeführt. Die Mini-Charge hat 1000 Proben und die Anzahl der Lernvorgänge beträgt 100.000.

main.py

with tf.Session() as sess:
    invisible_axis = True
    sess.run(tf.global_variables_initializer())

    for iteration in range(100000+1):
        z_0_batch = normal_distribution.sample(1000)
        log_q_0_batch = np.log(normal_distribution.calc_prob(z_0_batch))
        _, loss_value = sess.run([train, loss], {z_0:z_0_batch, log_q_0:log_q_0_batch})

Das Ergebnis ist hier. Sie können sehen, dass die komplexe Verteilung ausgehend von der Gaußschen Verteilung korrekt abgetastet wird. <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/302831/4ab8c588-1d49-e110-c3c9-deb72f6e8eac.png "， width=33%> <img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/302831/4d21dd7d-6aa1-fd70-57c0-e689f7525111.png "， width=33%><img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/302831/4dc89cef-d8ac-a2b6-4897-4224d1de1fae.png "， width=33%>

Alle Implementierungen

Im Folgenden finden Sie eine Zusammenfassung aller Implementierungen. Es gibt drei Typen: main.py, normalizing_flow.py und Distribution.py. main.py ist eine Datei zum Erstellen eines Berechnungsgraphen und zum Ausführen von Schulungen, normalizing_flow.py ist eine Datei zum Definieren eines Modells und einer Verlustfunktion zum Erstellen eines Berechnungsgraphen, und schließlich ist Distribution.py eine Stichprobe der Wahrscheinlichkeitsverteilung. Es ist eine Datei, die eine Funktion zur Berechnung der Wahrscheinlichkeit hat. Die Datei wird auch auf GitHub veröffentlicht.

main.py

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from distribution import *
from normalizing_flow import *

normal_distribution = NormalDistribution2D()
target_density = TargetDistribution1()

normalizing_flow = NormalizingFlow(K=16, dim=2)

z_0,log_q_0 = get_placeholder()
z_k, log_q_k = normalizing_flow(z_0,log_q_0)
loss = calc_loss(z_k, log_q_k, target_density)
train = get_train(loss)

with tf.Session() as sess:
    invisible_axis = True
    sess.run(tf.global_variables_initializer())

    for iteration in range(100000+1):
        z_0_batch = normal_distribution.sample(1000)
        log_q_0_batch = np.log(normal_distribution.calc_prob(z_0_batch))
        _, loss_value = sess.run([train, loss], {z_0:z_0_batch, log_q_0:log_q_0_batch})
        
        if iteration % 100 == 0:
            print('Iteration : {}   Loss : {}'.format(iteration, loss_value))
            
        if iteration % 10000 == 0:
            z_k_value = sess.run(z_k, {z_0:z_0_batch, log_q_0:log_q_0_batch})
            plt.figure(figsize=(6, 6))
            plt.scatter(z_k_value[:, 0], z_k_value[:, 1], alpha=0.7)
            if invisible_axis:
                plt.tick_params(bottom=False,left=False,right=False,top=False)
                plt.tick_params(labelbottom=False,labelleft=False,labelright=False,labeltop=False)
            plt.show()

normalizing_flow.py

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt


def get_placeholder():
    z_0 = tf.placeholder(tf.float32, shape=[None, 2])
    log_q_0 = tf.placeholder(tf.float32, shape=[None])
    return z_0,log_q_0


def calc_loss(z_k, log_q_k, target_density):
    log_p = tf.log(target_density.calc_prob_tf(z_k)+1e-7)
    loss = tf.reduce_mean(log_q_k - log_p, -1)
    return loss


def get_train(loss):
    return tf.train.AdamOptimizer().minimize(loss)


class PlanarFlow:
    def __init__(self, dim):
        self.dim = dim
        self.h = lambda x: tf.tanh(x)
        self.h_prime = lambda x: 1 - tf.tanh(x)**2
        self.w = tf.Variable(tf.random.truncated_normal(shape=(1, self.dim)))
        self.b = tf.Variable(tf.zeros(shape=(1)))
        self.u = tf.Variable(tf.random.truncated_normal(shape=(1, self.dim)))
        

    def __call__(self, z, log_q):
        z = z + self.u*self.h(tf.expand_dims(tf.reduce_sum(z*self.w, -1), -1) + self.b)
        psi = self.h_prime(tf.expand_dims(tf.reduce_sum(z*self.w, -1), -1) + self.b)*self.w
        det_jacob = tf.abs(1 + tf.reduce_sum(psi*self.u, -1))
        log_q = log_q -  tf.log(1e-7 + det_jacob)
        return z, log_q


class NormalizingFlow:
    def __init__(self, K, dim):
        self.K = K
        self.dim = dim
        self.planar_flow = [PlanarFlow(self.dim) for i in range(self.K)]
        
        
    def __call__(self, z_0,log_q_0):
        z, log_q = self.planar_flow[0](z_0,log_q_0)
        for pf in self.planar_flow[1:]:
            z, log_q = pf(z, log_q)
        return z, log_q

distribution.py

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

class Distribution:    
    def calc_prob(self, z):
        p = np.zeros(z.shape[0])
        return p
    

    def plot(self, size=5):
        side = np.linspace(-size, size, 1000)
        z1, z2 = np.meshgrid(side, side)
        shape = z1.shape
        z1 = z1.ravel()
        z2 = z2.ravel()
        z = np.c_[z1, z2]
        probability = self.calc_prob(z).reshape(shape)
        plt.figure(figsize=(6, 6))
        plt.imshow(probability)
        plt.tick_params(bottom=False,left=False,right=False,top=False)
        plt.tick_params(labelbottom=False,labelleft=False,labelright=False,labeltop=False)
        plt.show()


class NormalDistribution2D(Distribution):
    def sample(self, sample_num):
        z = np.random.randn(sample_num, 2)
        return z
    

    def sample_tf(self, sample_num):
        z = tf.random_normal([sample_num, 2])
        return z
    

    def calc_prob(self, z):
        p = np.exp(-(z[:, 0]**2+z[:, 1]**2)/2)/(2*np.pi) 
        return p
    

    def calc_prob_tf(self, z):
        p = tf.exp(-(z[:, 0]**2+z[:, 1]**2)/2)/(2*np.pi) 
        return p


class TargetDistribution1(Distribution):
    def calc_prob(self, z):
        z1, z2 = z[:, 0], z[:, 1]
        norm = np.sqrt(z1**2+z2**2)
        exp1 = np.exp(-0.5*((z1-2)/0.6)**2)
        exp2 = np.exp(-0.5*((z1+2)/0.6)**2)
        p = 0.5*((norm - 2)/0.4)**2 - np.log(exp1 + exp2)
        return np.exp(-p)
    

    def calc_prob_tf(self, z):
        z1, z2 = z[:, 0], z[:, 1]
        norm = tf.sqrt(z1**2+z2**2)
        exp1 = tf.exp(-0.5*((z1-2)/0.6)**2)
        exp2 = tf.exp(-0.5*((z1+2)/0.6)**2)
        p = 0.5*((norm - 2)/0.4)**2 - tf.log(exp1 + exp2)
        return tf.exp(-p)


class TargetDistribution2(Distribution):
    def calc_prob(self, z):
        z1, z2 = z[:, 0], z[:, 1]
        w1 = np.sin(0.5*np.pi*z1)
        p = 0.5*((z2 - w1)/0.4)**2
        return np.exp(-p)
    
    
    def calc_prob_tf(self, z):
        z1, z2 = z[:, 0], z[:, 1]
        w1 = tf.sin(0.5*np.pi*z1)
        p = 0.5*((z2 - w1)/0.4)**2
        return tf.exp(-p)

Verweise

[1] Ian.Goodfellow et al., "Deep Learning". [2] C. M. Bishop, "Pattern Recognition and Machine Learning". [3] C. M. Bishop, "Pattern Recognition and Machine Learning". [4] Atsushi Suyama, "Einführung in das maschinelle Lernen durch Bayes'sche Inferenz".