example.py

from minifox import MinFoxSolver
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

# toy task: random projections
# first, we sample p-dimensional matrix M
# A is constructed as a random projection of M
# B is a random projection of first 4 components of M
# hence the "right answer" is to extract 5-th component of M from A
n, p = 1000, 5

M = np.random.randn(n, p)
A = M.dot(np.random.randn(p, p))
B = M[:, :-1].dot(np.random.randn(p - 1, p))

fox = MinFoxSolver(p=5, max_iters=10 ** 5, tolerance=1e-4, verbose=True).fit(A, B)

unpredictable_variable = fox.predict(A=A)
unknown_component = M[:, -1]

plt.scatter(unpredictable_variable, unknown_component, alpha=0.1)
plt.xlabel('generated variable')
plt.ylabel('component that is missing from B')

minifox.py

import numpy as np
import tensorflow as tf
L = tf.keras.layers

class MinFoxSolver:
    def __init__(self, p, p_b=None, pred_steps=5, gen_steps=1, max_iters=10 ** 5, tolerance=1e-3, 
                 optimizer=tf.train.AdamOptimizer(5e-4),
                 make_generator=lambda: L.Dense(1, name='he_who_generates_unpredictable'),
                 make_predictor=lambda: L.Dense(1, name='he_who_predicts_generated_variable'),
                 sess=None, verbose=False,
                ):
        """
        Given two matrices A and B, predict a variable f(A) that is impossible to predict from matrix B
        :param p: last dimension of A
        :param p_b: dimension of B, default p_b = p
        :param pred_steps: predictor g(B) training iterations per one training step
        :param gen_steps: generator f(A) training iterations per one training step
        :param max_iters: maximum number of optimization steps till termination
        :param tolerance: terminates if loss difference between 10-iteration cycles reaches this value
            set to 0 to iterate for max_steps
        :param optimizer: tf optimizer to be used on both generator and discriminator
        :param make_generator: callback to create a keras model for target variable generator given A
        :param make_predictor: callback to create a keras model for target variable predictor given B
        /* Маааленькая лисёнка */
        """
        self.session = sess = sess or tf.get_default_session() \
            or tf.Session(config=tf.ConfigProto(device_count={'GPU': 0}))
        self.pred_steps, self.gen_steps = pred_steps, gen_steps
        self.max_iters, self.tolerance = max_iters, tolerance
        self.verbose = verbose
        
        with sess.as_default(), sess.graph.as_default():
            A = self.A = tf.placeholder(tf.float32, [None, p])
            B = self.B = tf.placeholder(tf.float32, [None, p_b or p])
            self.generator = make_generator()
            self.predictor = make_predictor()
            
            prediction = self.predictor(B)
            target = self.generator(A)
            
            # normalize target to prevent generator from scaling it infinitely wide
            mu, var = tf.nn.moments(target[:, 0], axes=0)
            target = (target - mu) / tf.sqrt(var)

            self.loss = self.compute_loss(target, prediction)
            
            self.update_pred = optimizer.minimize(self.loss, var_list=self.predictor.trainable_variables)
            self.update_gen = optimizer.minimize(-self.loss, var_list=self.generator.trainable_variables)
            self.prediction, self.target = prediction, target
            
    def compute_loss(self, target, prediction):
        return tf.reduce_mean(tf.squared_difference(target, prediction))
    
    def fit(self, A, B):
        sess = self.session
        with sess.as_default(), sess.graph.as_default():
            sess.run(tf.global_variables_initializer())
            
            feed = {self.A: A, self.B: B}
            prev_loss = sess.run(self.loss, feed)
            
            for i in range(1, self.max_iters + 1):
                for j in range(self.pred_steps):
                    sess.run(self.update_pred, feed)
                for j in range(self.gen_steps):
                    sess.run(self.update_gen, feed)
                if i % 100 == 0:
                    loss_i = sess.run(self.loss, feed)
                    if self.verbose:
                        print("step %i; loss=%.3f; delta=%.3f" % (i, loss_i, abs(prev_loss - loss_i)))
                    if abs(prev_loss - loss_i) < self.tolerance:
                        if self.verbose: print("Done: reached target tolerance")
                        break
                    prev_loss = loss_i
            else:
                if self.verbose:
                    print("Done: reached max steps")
            return self
    
    def predict(self, A=None, B=None):
        assert (A is None) != (B is None), "Please use either predict(A=...) or predict(B=...)"
        sess = self.session
        with sess.as_default(), sess.graph.as_default():
            if A is not None:
                return sess.run(self.target, {self.A: A})
            else:
                return sess.run(self.prediction, {self.B: B})
    
    def get_weights(self):
        return self.session.run({'generator': self.generator.trainable_variables, 
                                 'predictor': self.predictor.trainable_variables})