def test_normalizing_flow():
flow_layers = [
tfb.RealNVP(num_masked=1,
shift_and_log_scale_fn=tfb.real_nvp_default_template(
hidden_layers=[8, 8]))
]
# 3D Multivariate Gaussian base distribution
nf = NormalizingFlow(
base_distribution=tfd.MultivariateNormalDiag(loc=[0., 0.]),
flow_layers=flow_layers)
nfm = NormalizingFlowModel(nf)
# Must be float32 for RealNVP
target_distribution = tfd.MultivariateNormalDiag(loc=[1., 0.])
dataset = NumpyDataset(X=target_distribution.sample(96))
# Tests a simple flow of one RealNVP layer.
X = nfm.flow.sample()
x1 = tf.zeros([2])
x2 = dataset.X[0]
# log likelihoods should be negative
assert nfm.flow.log_prob(X).numpy() < 0
assert nfm.flow.log_prob(x1).numpy() < 0
assert nfm.flow.log_prob(x2).numpy() < 0
# # Fit model
final = nfm.fit(dataset, nb_epoch=5)
print(final)
assert final > 0
def test_normalizing_flow_model_reload():
"""Test that NormalizingFlowModel can be reloaded correctly."""
from deepchem.models.normalizing_flows import NormalizingFlow, NormalizingFlowModel
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
tfk = tf.keras
model_dir = tempfile.mkdtemp()
Made = tfb.AutoregressiveNetwork(params=2,
hidden_units=[512, 512],
activation='relu',
dtype='float64')
flow_layers = [tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn=Made)]
# 3D Multivariate Gaussian base distribution
nf = NormalizingFlow(base_distribution=tfd.MultivariateNormalDiag(
loc=np.zeros(2), scale_diag=np.ones(2)),
flow_layers=flow_layers)
nfm = NormalizingFlowModel(nf, model_dir=model_dir)
target_distribution = tfd.MultivariateNormalDiag(loc=np.array([1., 0.]))
dataset = dc.data.NumpyDataset(X=target_distribution.sample(96))
final = nfm.fit(dataset, nb_epoch=1)
x = np.zeros(2)
lp1 = nfm.flow.log_prob(x).numpy()
assert nfm.flow.sample().numpy().shape == (2, )
reloaded_model = NormalizingFlowModel(nf, model_dir=model_dir)
reloaded_model.restore()
# Check that reloaded model can sample from the distribution
assert reloaded_model.flow.sample().numpy().shape == (2, )
lp2 = reloaded_model.flow.log_prob(x).numpy()
# Check that density estimation is same for reloaded model
assert np.all(lp1 == lp2)