spectral_normalization.py

def spectral_normalization(W, num_iters=1, name='spectral_normalization'):
    with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
        
        def l2_normalize(v, eps=1e-12):
            return v / (tf.reduce_sum(v ** 2) ** 0.5 + eps)

        # Usually num_iters = 1 will be enough
        W_shape    = W.shape.as_list()
        W_reshaped = tf.reshape(W, [-1, W_shape[-1]])
        
        u = tf.get_variable("u", [1, W_shape[-1]], initializer=tf.truncated_normal_initializer(), trainable=False)

        def power_iteration(i, u_i, v_i):
            v_ip1 = l2_normalize(tf.matmul(u_i, tf.transpose(W_reshaped)))
            u_ip1 = l2_normalize(tf.matmul(v_ip1, W_reshaped))
            return i + 1, u_ip1, v_ip1

        _, u_final, v_final = tf.while_loop(
            cond      = lambda i, _1, _2: i < num_iters,
            body      = power_iteration,
            loop_vars = (tf.constant(0, dtype=tf.int32),
            u, tf.zeros(dtype=tf.float32, shape=[1, W_reshaped.shape.as_list()[0]]))
        )

        sigma = tf.matmul(tf.matmul(v_final, W_reshaped), tf.transpose(u_final))[0, 0]
        # sigma = tf.reduce_sum(tf.matmul(u_final, tf.transpose(W_reshaped)) * v_final)
        W_bar = W_reshaped / sigma

        with tf.control_dependencies([u.assign(u_final)]):
            W_bar = tf.reshape(W_bar, W_shape)
            
        return W_bar