feat: initial for IDF

integer_discrete_flows/utils/log_likelihood.py

@@ -0,0 +1,56 @@
+from __future__ import print_function
+import numpy as np
+from scipy.misc import logsumexp
+from optimization.loss import calculate_loss_array
+
+
+def calculate_likelihood(X, model, args, S=5000, MB=500):
+
+    # set auxiliary variables for number of training and test sets
+    N_test = X.size(0)
+
+    X = X.view(-1, *args.input_size)
+
+    likelihood_test = []
+
+    if S <= MB:
+        R = 1
+    else:
+        R = S // MB
+        S = MB
+
+    for j in range(N_test):
+        if j % 100 == 0:
+            print('Progress: {:.2f}%'.format(j / (1. * N_test) * 100))
+
+        x_single = X[j].unsqueeze(0)
+
+        a = []
+        for r in range(0, R):
+            # Repeat it for all training points
+            x = x_single.expand(S, *x_single.size()[1:]).contiguous()
+
+            x_mean, z_mu, z_var, ldj, z0, zk = model(x)
+
+            a_tmp = calculate_loss_array(x_mean, x, z_mu, z_var, z0, zk, ldj, args)
+
+            a.append(-a_tmp.cpu().data.numpy())
+
+        # calculate max
+        a = np.asarray(a)
+        a = np.reshape(a, (a.shape[0] * a.shape[1], 1))
+        likelihood_x = logsumexp(a)
+        likelihood_test.append(likelihood_x - np.log(len(a)))
+
+    likelihood_test = np.array(likelihood_test)
+
+    nll = -np.mean(likelihood_test)
+
+    if args.input_type == 'multinomial':
+        bpd = nll/(np.prod(args.input_size) * np.log(2.))
+    elif args.input_type == 'binary':
+        bpd = 0.
+    else:
+        raise ValueError('invalid input type!')
+
+    return nll, bpd