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)