FFJORD

הצג באתר TensorFlow.org הפעל בגוגל קולאב צפה במקור ב-GitHub הורד מחברת

להכין

התקנה ראשונה של חבילות המשמשות בהדגמה זו.

pip install -q dm-sonnet

יבוא (tf, tfp עם טריק נלווה וכו')

import numpy as np
import tqdm as tqdm
import sklearn.datasets as skd

# visualization
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import kde

# tf and friends
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
import sonnet as snt
tf.enable_v2_behavior()

tfb = tfp.bijectors
tfd = tfp.distributions

def make_grid(xmin, xmax, ymin, ymax, gridlines, pts):
  xpts = np.linspace(xmin, xmax, pts)
  ypts = np.linspace(ymin, ymax, pts)
  xgrid = np.linspace(xmin, xmax, gridlines)
  ygrid = np.linspace(ymin, ymax, gridlines)
  xlines = np.stack([a.ravel() for a in np.meshgrid(xpts, ygrid)])
  ylines = np.stack([a.ravel() for a in np.meshgrid(xgrid, ypts)])
  return np.concatenate([xlines, ylines], 1).T

grid = make_grid(-3, 3, -3, 3, 4, 100)

/usr/local/lib/python3.6/dist-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.
  import pandas.util.testing as tm

פונקציות עוזר להדמיה

ביקטור FFJORD

בקולאב זה אנו מדגימים את FFJORD bijector, שהוצע במקור במאמר על ידי Grathwohl, Will, et al. arXiv קישור .

בשנים בקיצור את הרעיון מאחורי גישה כזו הוא להקים התכתבות בין חלוקת בסיס ידוע וחלוק הנתונים.

כדי ליצור את הקשר הזה, אנחנו צריכים

  1. גדר המפה bijective \(\mathcal{T}_{\theta}:\mathbf{x} \rightarrow \mathbf{y}\), \(\mathcal{T}_{\theta}^{1}:\mathbf{y} \rightarrow \mathbf{x}\) בין מרחב \(\mathcal{Y}\) שעליו הפצת בסיס מוגדרת ומרחב \(\mathcal{X}\) של תחום הנתונים.
  2. ביעילות לעקוב אחר דפורמציות שאנו מבצעים כדי להעביר את הרעיון של הסתברות על \(\mathcal{X}\).

התנאי השני מעוגן את הביטוי הבא עבור התפלגות ההסתברות מוגדרת על \(\mathcal{X}\):

\[ \log p_{\mathbf{x} }(\mathbf{x})=\log p_{\mathbf{y} }(\mathbf{y})-\log \operatorname{det}\left|\frac{\partial \mathcal{T}_{\theta}(\mathbf{y})}{\partial \mathbf{y} }\right| \]

FFJORD bijector משיג זאת על ידי הגדרת טרנספורמציה

\[ \mathcal{T_{\theta} }: \mathbf{x} = \mathbf{z}(t_{0}) \rightarrow \mathbf{y} = \mathbf{z}(t_{1}) \quad : \quad \frac{d \mathbf{z} }{dt} = \mathbf{f}(t, \mathbf{z}, \theta) \]

שינוי זה הוא להיפוך, כל עוד פונקציה \(\mathbf{f}\) המתאר את האבולוציה של המדינה \(\mathbf{z}\) הוא התנהג היטב ואת log_det_jacobian יכול להיות מחושב על ידי שילוב הביטוי הבא.

\[ \log \operatorname{det}\left|\frac{\partial \mathcal{T}_{\theta}(\mathbf{y})}{\partial \mathbf{y} }\right| = -\int_{t_{0} }^{t_{1} } \operatorname{Tr}\left(\frac{\partial \mathbf{f}(t, \mathbf{z}, \theta)}{\partial \mathbf{z}(t)}\right) d t \]

בהדגמה זו נוכל להכשיר bijector FFJORD עיוות הפצה גאוס על ההפצה שהוגדרה על ידי moons נתון. זה ייעשה ב-3 שלבים:

  • גדר הפצת בסיס
  • הגדר את FFJORD ביג'קטור
  • צמצם את הסבירות המדויקת ביומן של מערך הנתונים

ראשית, אנו טוענים את הנתונים

מערך נתונים

png

לאחר מכן, אנו מציגים התפלגות בסיס

base_loc = np.array([0.0, 0.0]).astype(np.float32)
base_sigma = np.array([0.8, 0.8]).astype(np.float32)
base_distribution = tfd.MultivariateNormalDiag(base_loc, base_sigma)

אנו משתמשים רב שכבתיים perceptron למודל state_derivative_fn .

אמנם לא הכרחי במערך זה, הוא לעתים קרובות benefitial לעשות state_derivative_fn תלוי זמן. כאן אנו להשיג זאת על ידי שרשור t לתשומות של הרשת שלנו.

class MLP_ODE(snt.Module):
  """Multi-layer NN ode_fn."""
  def __init__(self, num_hidden, num_layers, num_output, name='mlp_ode'):
    super(MLP_ODE, self).__init__(name=name)
    self._num_hidden = num_hidden
    self._num_output = num_output
    self._num_layers = num_layers
    self._modules = []
    for _ in range(self._num_layers - 1):
      self._modules.append(snt.Linear(self._num_hidden))
      self._modules.append(tf.math.tanh)
    self._modules.append(snt.Linear(self._num_output))
    self._model = snt.Sequential(self._modules)

  def __call__(self, t, inputs):
    inputs = tf.concat([tf.broadcast_to(t, inputs.shape), inputs], -1)
    return self._model(inputs)

פרמטרים של מודל והדרכה

כעת אנו בונים ערימה של FFJORD bictors. Bijector כל מסופק עם ode_solve_fn ו trace_augmentation_fn וזה עצמו state_derivative_fn המודל, כך שהם מייצגים רצף של טרנספורמציות שונות.

ביקטור בניין

עכשיו אנחנו יכולים להשתמש TransformedDistribution שהינו תוצאה של עיקום base_distribution עם stacked_ffjord bijector.

transformed_distribution = tfd.TransformedDistribution(
    distribution=base_distribution, bijector=stacked_ffjord)

כעת אנו מגדירים את הליך ההכשרה שלנו. אנו פשוט ממזערים את הסבירות היומן השלילי של הנתונים.

הַדְרָכָה

דגימות

ציירו דגימות מהתפלגות בסיס והתפלגות.

evaluation_samples = []
base_samples, transformed_samples = get_samples()
transformed_grid = get_transformed_grid()
evaluation_samples.append((base_samples, transformed_samples, transformed_grid))

WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.

panel_id = 0
panel_data = evaluation_samples[panel_id]
fig, axarray = plt.subplots(
  1, 4, figsize=(16, 6))
plot_panel(
    grid, panel_data[0], panel_data[2], panel_data[1], moons, axarray, False)
plt.tight_layout()

png

learning_rate = tf.Variable(LR, trainable=False)
optimizer = snt.optimizers.Adam(learning_rate)

for epoch in tqdm.trange(NUM_EPOCHS // 2):
  base_samples, transformed_samples = get_samples()
  transformed_grid = get_transformed_grid()
  evaluation_samples.append(
      (base_samples, transformed_samples, transformed_grid))
  for batch in moons_ds:
    _ = train_step(optimizer, batch)

0%|          | 0/40 [00:00<?, ?it/s]
WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/math/ode/base.py:350: calling while_loop_v2 (from tensorflow.python.ops.control_flow_ops) with back_prop=False is deprecated and will be removed in a future version.
Instructions for updating:
back_prop=False is deprecated. Consider using tf.stop_gradient instead.
Instead of:
results = tf.while_loop(c, b, vars, back_prop=False)
Use:
results = tf.nest.map_structure(tf.stop_gradient, tf.while_loop(c, b, vars))
100%|██████████| 40/40 [07:00<00:00, 10.52s/it]

panel_id = -1
panel_data = evaluation_samples[panel_id]
fig, axarray = plt.subplots(
  1, 4, figsize=(16, 6))
plot_panel(grid, panel_data[0], panel_data[2], panel_data[1], moons, axarray)
plt.tight_layout()

png

אימון זה לאורך זמן עם קצב למידה מביא לשיפורים נוספים.

לא נדון בדוגמה זו, FFJORD bictor תומך באומדן העקבות הסטוכסטיות של Hutchinson. אומד מסוים יכול להינתן באמצעות trace_augmentation_fn . באופן דומה אינטגרטורים אלטרנטיבה ניתן להשתמש בהגדרת מנהג ode_solve_fn .