Merge 452df3d into 22b819d

convoys/regression.py

@@ -17,25 +17,32 @@ def _get_placeholders(n, k):
     )
 
 
-def _optimize(sess, target, feed_dict):
+def _optimize(sess, target, feed_dict, variables):
     learning_rate_input = tf.placeholder(tf.float32, [])
     optimizer = tf.train.AdamOptimizer(learning_rate_input).minimize(-target)
 
     # TODO(erikbern): this is going to add more and more variables every time we run this
     sess.run(tf.global_variables_initializer())
 
-    best_cost, best_step, step = float('-inf'), 0, 0
-    learning_rate = 0.1
+    best_step, step = 0, 0
+    learning_rate = 1.0
+    best_state = sess.run(variables)
+    best_cost = sess.run(target, feed_dict=feed_dict)
+
     while True:
         feed_dict[learning_rate_input] = learning_rate
         sess.run(optimizer, feed_dict=feed_dict)
         cost = sess.run(target, feed_dict=feed_dict)
         if cost > best_cost:
             best_cost, best_step = cost, step
+            best_state = sess.run(variables)
         if step - best_step > 40:
             learning_rate /= 10
-            best_cost = float('-inf')
+            best_step = step
+            for variable, value in zip(variables, best_state):
+                sess.run(tf.assign(variable, value))
         if learning_rate < 1e-6:
+            sys.stdout.write('\n')
             break
         step += 1
         sys.stdout.write('step %6d (lr %6.6f): %12.4f' % (step, learning_rate, cost))
@@ -100,18 +107,18 @@ def fit(self, X, B, T):
         lambd = tf.exp(X_prod_alpha)
         c = tf.sigmoid(X_prod_beta)
 
-        log_pdf = lambda T: tf.log(lambd) - T*lambd
-        cdf = lambda T: 1 - tf.exp(-(T * lambd))
+        log_pdf = tf.log(lambd) - T_input*lambd
+        cdf = 1 - tf.exp(-(T_input * lambd))
 
-        LL_observed = tf.log(c) + log_pdf(T_input)
-        LL_censored = tf.log((1-c) + c * (1 - cdf(T_input)))
+        LL_observed = tf.log(c) + log_pdf
+        LL_censored = tf.log((1-c) + c * (1 - cdf))
 
         LL = tf.reduce_sum(B_input * LL_observed + (1 - B_input) * LL_censored, 0)
         LL_penalized = LL - self._L2_reg * tf.reduce_sum(beta * beta, 0)
 
         with tf.Session() as sess:
             feed_dict = {X_input: X, B_input: B, T_input: T}
-            _optimize(sess, LL_penalized, feed_dict)
+            _optimize(sess, LL_penalized, feed_dict, (alpha, beta))
             self.params = _get_params(sess, {'beta': beta, 'alpha': alpha})
             self.params['alpha_hessian'] = _get_hessian(sess, LL_penalized, alpha, feed_dict)
             self.params['beta_hessian'] = _get_hessian(sess, LL_penalized, beta, feed_dict)
@@ -142,19 +149,19 @@ def fit(self, X, B, T):
         c = tf.sigmoid(X_prod_beta)
 
         # PDF of Weibull: k * lambda * (x * lambda)^(k-1) * exp(-(t * lambda)^k)
-        log_pdf = lambda T: tf.log(k) + tf.log(lambd) + (k-1)*(tf.log(T) + tf.log(lambd)) - (T*lambd)**k
+        log_pdf = tf.log(k) + tf.log(lambd) + (k-1)*(tf.log(T_input) + tf.log(lambd)) - (T_input*lambd)**k
         # CDF of Weibull: 1 - exp(-(t * lambda)^k)
-        cdf = lambda T: 1 - tf.exp(-(T * lambd)**k)
+        cdf = 1 - tf.exp(-(T_input * lambd)**k)
 
-        LL_observed = tf.log(c) + log_pdf(T_input)
-        LL_censored = tf.log((1-c) + c * (1 - cdf(T_input)))
+        LL_observed = tf.log(c) + log_pdf
+        LL_censored = tf.log((1-c) + c * (1 - cdf))
 
         LL = tf.reduce_sum(B_input * LL_observed + (1 - B_input) * LL_censored, 0)
         LL_penalized = LL - self._L2_reg * tf.reduce_sum(beta * beta, 0)
 
         with tf.Session() as sess:
             feed_dict = {X_input: X, B_input: B, T_input: T}
-            _optimize(sess, LL_penalized, feed_dict)
+            _optimize(sess, LL_penalized, feed_dict, (alpha, beta, log_k_var))
             self.params = _get_params(sess, {'beta': beta, 'alpha': alpha, 'k': k})
             self.params['alpha_hessian'] = _get_hessian(sess, LL_penalized, alpha, feed_dict)
             self.params['beta_hessian'] = _get_hessian(sess, LL_penalized, beta, feed_dict)
@@ -185,19 +192,19 @@ def fit(self, X, B, T):
         c = tf.sigmoid(X_prod_beta)
 
         # PDF of gamma: 1.0 / gamma(k) * lambda ^ k * t^(k-1) * exp(-t * lambda)
-        log_pdf = lambda T: -tf.lgamma(k) + k*tf.log(lambd) + (k-1)*tf.log(T) - lambd*T
+        log_pdf = -tf.lgamma(k) + k*tf.log(lambd) + (k-1)*tf.log(T_input) - lambd*T_input
         # CDF of gamma: gammainc(k, lambda * t)
-        cdf = lambda T: tf.igamma(k, lambd * T)
+        cdf = tf.igamma(k, lambd * T_input)
 
-        LL_observed = tf.log(c) + log_pdf(T_input)
-        LL_censored = tf.log((1-c) + c * (1 - cdf(T_input)))
+        LL_observed = tf.log(c) + log_pdf
+        LL_censored = tf.log((1-c) + c * (1 - cdf))
 
         LL = tf.reduce_sum(B_input * LL_observed + (1 - B_input) * LL_censored, 0)
         LL_penalized = LL - self._L2_reg * tf.reduce_sum(beta * beta, 0)
 
         with tf.Session() as sess:
             feed_dict = {X_input: X, B_input: B, T_input: T}
-            _optimize(sess, LL_penalized, feed_dict)
+            _optimize(sess, LL_penalized, feed_dict, (alpha, beta, log_k_var))
             self.params = _get_params(sess, {'beta': beta, 'alpha': alpha, 'k': k})
             self.params['alpha_hessian'] = _get_hessian(sess, LL_penalized, alpha, feed_dict)
             self.params['beta_hessian'] = _get_hessian(sess, LL_penalized, beta, feed_dict)