FFJORD

সেটআপ

প্রথমে এই ডেমোতে ব্যবহৃত প্যাকেজগুলি ইনস্টল করুন।

pip install -q dm-sonnet

আমদানি (tf, adjoint trick সহ 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

ভিজ্যুয়ালাইজেশনের জন্য হেল্পার ফাংশন

def plot_density(data, axis):
  x
, y = np.squeeze(np.split(data, 2, axis=1))
  levels
= np.linspace(0.0, 0.75, 10)
  kwargs
= {'levels': levels}
 
return sns.kdeplot(x, y, cmap="viridis", shade=True,
                     shade_lowest
=True, ax=axis, **kwargs)


def plot_points(data, axis, s=10, color='b', label=''):
  x
, y = np.squeeze(np.split(data, 2, axis=1))
  axis
.scatter(x, y, c=color, s=s, label=label)


def plot_panel(
    grid
, samples, transformed_grid, transformed_samples,
    dataset
, axarray, limits=True):
 
if len(axarray) != 4:
   
raise ValueError('Expected 4 axes for the panel')
  ax1
, ax2, ax3, ax4 = axarray
  plot_points
(data=grid, axis=ax1, s=20, color='black', label='grid')
  plot_points
(samples, ax1, s=30, color='blue', label='samples')
  plot_points
(transformed_grid, ax2, s=20, color='black', label='ode(grid)')
  plot_points
(transformed_samples, ax2, s=30, color='blue', label='ode(samples)')
  ax3
= plot_density(transformed_samples, ax3)
  ax4
= plot_density(dataset, ax4)
 
if limits:
    set_limits
([ax1], -3.0, 3.0, -3.0, 3.0)
    set_limits
([ax2], -2.0, 3.0, -2.0, 3.0)
    set_limits
([ax3, ax4], -1.5, 2.5, -0.75, 1.25)


def set_limits(axes, min_x, max_x, min_y, max_y):
 
if isinstance(axes, list):
   
for axis in axes:
      set_limits
(axis, min_x, max_x, min_y, max_y)
 
else:
    axes
.set_xlim(min_x, max_x)
    axes
.set_ylim(min_y, max_y)

FFJORD বিজেক্টর

এই কোল্যাবে আমরা FFJORD বাইজেক্টর প্রদর্শন করি, যা মূলত Grathwohl, Will, et al দ্বারা পেপারে প্রস্তাবিত। লিংক arXiv

সংক্ষেপে এই ধরনের পদ্ধতির পিছনে ধারণা একটি পরিচিত বেস বন্টন এবং তথ্য বিতরণ মধ্যে একটি সাদৃশ্য স্থাপন করা হয়।

এই সংযোগ স্থাপন করতে, আমাদের প্রয়োজন

  1. একটি bijective মানচিত্র নির্ধারণ Tθ:xy, Tθ1:yx স্থান মধ্যে Y যার উপর ভিত্তি বন্টন সংজ্ঞায়িত করা হয় এবং স্থান X তথ্য ডোমেনের।
  2. দক্ষতার পরিগ্রহ করে আমরা সম্মুখের দিকে সম্ভাব্যতা ধারণা হস্তান্তর সঞ্চালন ট্র্যাক রাখতে X

দ্বিতীয় শর্ত সম্ভাব্যতা বিতরণের সংজ্ঞাসমূহ জন্য নিম্নলিখিত এক্সপ্রেশনে বিধিবদ্ধ করা হয় X:

logpx(x)=logpy(y)logdet|Tθ(y)y|

FFJORD বিজেক্টর একটি রূপান্তর সংজ্ঞায়িত করে এটি সম্পন্ন করে

Tθ:x=z(t0)y=z(t1):dzdt=f(t,z,θ)

যতদিন ফাংশন হিসাবে এই রূপান্তর, বিপরীত হয় f রাষ্ট্র বিবর্তন বর্ণনা z ভাল ভদ্র হয় এবং log_det_jacobian নিম্নলিখিত অভিব্যক্তি একীভূত নির্ণিত করা যেতে পারে।

logdet|Tθ(y)y|=t0t1Tr(f(t,z,θ)z(t))dt

এই ডেমো আমরা একটি FFJORD bijector প্রশিক্ষণ বন্টন দ্বারা সংজ্ঞায়িত সম্মুখের দিকে একটি গসিয়ান বন্টন দোমড়ানো হবে moons ডেটা সেটটি। এটি 3টি ধাপে করা হবে:

  • বেস বন্টন নির্ধারণ
  • FFJORD দ্বিজেক্টর সংজ্ঞায়িত করুন
  • ডেটাসেটের সঠিক লগ-সম্ভাবনা কমিয়ে দিন

প্রথমত, আমরা ডেটা লোড করি

ডেটাসেট

DATASET_SIZE = 1024 * 8 
BATCH_SIZE
= 256
SAMPLE_SIZE
= DATASET_SIZE

moons
= skd.make_moons(n_samples=DATASET_SIZE, noise=.06)[0]

moons_ds
= tf.data.Dataset.from_tensor_slices(moons.astype(np.float32))
moons_ds
= moons_ds.prefetch(tf.data.experimental.AUTOTUNE)
moons_ds
= moons_ds.cache()
moons_ds
= moons_ds.shuffle(DATASET_SIZE)
moons_ds
= moons_ds.batch(BATCH_SIZE)

plt
.figure(figsize=[8, 8])
plt
.scatter(moons[:, 0], moons[:, 1])
plt
.show()

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 সময় উপর নির্ভরশীল। এখানে আমরা concatenating দ্বারা এই অর্জন 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)

মডেল এবং প্রশিক্ষণ পরামিতি

LR = 1e-2 
NUM_EPOCHS
= 80
STACKED_FFJORDS
= 4
NUM_HIDDEN
= 8
NUM_LAYERS
= 3
NUM_OUTPUT
= 2

এখন আমরা FFJORD বিজেক্টরগুলির একটি স্ট্যাক তৈরি করি। প্রতিটি bijector সাথে উপলব্ধ করা হয় ode_solve_fn এবং trace_augmentation_fn এবং এটি নিজের ব্যাপার state_derivative_fn মডেল, যাতে তারা বিভিন্ন রূপান্তরের একটা ক্রম প্রতিনিধিত্ব করে।

বিল্ডিং দ্বিজেক্টর

solver = tfp.math.ode.DormandPrince(atol=1e-5)
ode_solve_fn
= solver.solve
trace_augmentation_fn
= tfb.ffjord.trace_jacobian_exact

bijectors
= []
for _ in range(STACKED_FFJORDS):
  mlp_model
= MLP_ODE(NUM_HIDDEN, NUM_LAYERS, NUM_OUTPUT)
  next_ffjord
= tfb.FFJORD(
      state_time_derivative_fn
=mlp_model,ode_solve_fn=ode_solve_fn,
      trace_augmentation_fn
=trace_augmentation_fn)
  bijectors
.append(next_ffjord)

stacked_ffjord
= tfb.Chain(bijectors[::-1])

এখন আমরা ব্যবহার করতে পারেন TransformedDistribution যা warping ফল base_distribution সঙ্গে stacked_ffjord bijector।

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

এখন আমরা আমাদের প্রশিক্ষণ পদ্ধতি সংজ্ঞায়িত. আমরা কেবল ডেটার নেতিবাচক লগ-সম্ভাবনা কমিয়ে দিই।

প্রশিক্ষণ

@tf.function
def train_step(optimizer, target_sample):
 
with tf.GradientTape() as tape:
    loss
= -tf.reduce_mean(transformed_distribution.log_prob(target_sample))
  variables
= tape.watched_variables()
  gradients
= tape.gradient(loss, variables)
  optimizer
.apply(gradients, variables)
 
return loss

নমুনা

@tf.function
def get_samples():
  base_distribution_samples
= base_distribution.sample(SAMPLE_SIZE)
  transformed_samples
= transformed_distribution.sample(SAMPLE_SIZE)
 
return base_distribution_samples, transformed_samples


@tf.function
def get_transformed_grid():
  transformed_grid
= stacked_ffjord.forward(grid)
 
return transformed_grid

ভিত্তি এবং রূপান্তরিত বিতরণ থেকে প্লট নমুনা.

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 । একইভাবে বিকল্প integrators কাস্টম নির্ধারণ করে ব্যবহার করা যেতে পারে ode_solve_fn