dense matrix factorization

mf.py

import tensorflow as tf
import tqdm

def tfmf(D, k, 
         params_dict={"iters":2**14, "p":2, "h":1E-2, "init":"kmeans"}, 
         regularizer_dict={"lambda_C":1E-3, "lambda_B":1E-3, "p":1}):
  
   m, n = D.shape[0], D.shape[1]
 
   D_con = tf.constant(D)
   C_var = tf.Variable(np.random.normal(0, 1, (m, k, 1)).astype(D.dtype))
   
   if params_dict["init"] == "kmeans":
       from sklearn.cluster import KMeans
       kmeans  = KMeans(k, n_jobs=-1).fit(D)
       centers = np.expand_dims(kmeans.cluster_centers_, 0)
       B_var = tf.Variable(centers.astype(D.dtype))
   else:
       B_var = tf.Variable(np.random.normal(0, 1, (1, k, n)).astype(D.dtype))
  
   BdotC = tf.reduce_sum(B_var*C_var, 1) 
   raw_loss = tf.reduce_mean(tf.abs(BdotC-D_con)**params_dict["p"])
   all_loss = raw_loss   

   if regularizer_dict:
       p_r   = regularizer_dict["p"]
       L_C = regularizer_dict["lambda_C"]
       L_B = regularizer_dict["lambda_B"]
       reg_loss  = L_C*tf.reduce_mean(tf.abs(C_var)**p_r) + \
                   L_B*tf.reduce_mean(tf.abs(B_var)**p_r)
       all_loss += reg_loss

   h = params_dict["h"]
   optimizer_C = tf.train.AdamOptimizer(h).minimize(all_loss, var_list=[C_var])
   optimizer_B = tf.train.AdamOptimizer(h).minimize(all_loss, var_list=[B_var])
   
   init = tf.global_variables_initializer()
   sess = tf.Session()
   sess.run(init) 

   if params_dict["init"] == "kmeans":
       for _ in range(10):
           sess.run(optimizer_C)

   progress_bar = tqdm.tqdm(range(params_dict["iters"]), desc=str(sess.run(all_loss)))
   for _ in progress_bar:
       sess.run(optimizer_B)
       sess.run(optimizer_C)
       
       if _ % max(params_dict["iters"]/1024, 1) == 0:
           progress_bar.desc=str(sess.run(all_loss))
  
   C_mat = sess.run(C_var)
   B_mat = sess.run(B_var)
   L_raw = sess.run(raw_loss)
   L_all = sess.run(all_loss)
   
   return (C_mat[:,:,0], B_mat[0]), (L_raw, L_all)