FFJORD

مشاهده در TensorFlow.org در Google Colab اجرا شود مشاهده منبع در 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 را نشان می‌دهیم که در ابتدا توسط Grathwohl، Will و همکاران در مقاله پیشنهاد شده بود. آرشیو لینک .

در خلاصه ایده پشت چنین رویکردی است که برای ایجاد مکاتبات بین یک توزیع پایه شناخته شده و توزیع داده ها.

برای ایجاد این ارتباط، ما نیاز داریم

  1. تعریف یک دوسویی نقشه \(\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)

ما با استفاده از یک لایه پرسپترون چند مدل 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 می سازیم. هر 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 از تخمین ردیابی تصادفی هاچینسون پشتیبانی می کند. برآوردگر خاص را می توان از طریق ارائه trace_augmentation_fn . به طور مشابه انتگرال جایگزین می تواند با تعریف سفارشی استفاده می شود ode_solve_fn .