Further TF 1.15 updates

Add progress bar for training (tqdm dependency)

Config.py

@@ -8,14 +8,14 @@ def cfg():
     # Base configuration
     model_config = {"enst_path" : "/mnt/windaten/Datasets/ENST_Drums", # SET MUSDB PATH HERE, AND SET CCMIXTER PATH IN CCMixter.xml
                     "estimates_path" : "/mnt/windaten/Source_Estimates", # SET THIS PATH TO WHERE YOU WANT SOURCE ESTIMATES PRODUCED BY THE TRAINED MODEL TO BE SAVED. Folder itself must exist!
-                    "data_path" : "data", # Set this to where the preprocessed dataset should be saved
+                    "data_path" : "/mnt/windaten/Mix-U-Net Data", # Set this to where the preprocessed dataset should be saved
 
                     "model_base_dir" : "checkpoints", # Base folder for model checkpoints
                     "log_dir" : "logs", # Base folder for logs files
                     "batch_size" : 16, # Batch size
                     "init_sup_sep_lr" : 1e-4, # Supervised separator learning rate
                     "epoch_it" : 2000, # Number of supervised separator steps per epoch
-                    'cache_size': 4000, # Number of audio snippets buffered in the random shuffle queue. Larger is better, since workers put multiple examples of one song into this queue. The number of different songs that is sampled from with each batch equals cache_size / num_snippets_per_track. Set as high as your RAM allows.
+                    'cache_size': 1000, # Number of audio snippets buffered in the random shuffle queue. Larger is better, since workers put multiple examples of one song into this queue. The number of different songs that is sampled from with each batch equals cache_size / num_snippets_per_track. Set as high as your RAM allows.
                     'num_workers' : 4, # Number of processes used for each TF map operation used when loading the dataset
                     "num_snippets_per_track" : 100, # Number of snippets that should be extracted from each song at a time after loading it. Higher values make data loading faster, but can reduce the batches song diversity
                     'num_layers' : 10, # How many U-Net layers
@@ -34,7 +34,7 @@ def cfg():
                     'upsampling' : 'linear', # Type of technique used for upsampling the feature maps in a unet architecture, either 'linear' interpolation or 'learned' filling in of extra samples
                     'task' : 'dry', # Type of separation task. 'voice' : Separate music into voice and accompaniment. 'multi_instrument': Separate music into guitar, bass, vocals, drums and other (Sisec)
                     'augmentation' : False, # Random attenuation of source signals to improve generalisation performance (data augmentation)
-                    'raw_audio_loss' : True, # Only active for unet_spectrogram network. True: L2 loss on audio. False: L1 loss on spectrogram magnitudes for training and validation and test loss
+                    'raw_audio_loss' : True, # Only active for unet_spectrogram network. True: L1 loss on audio. False: L1 loss on spectrogram magnitudes for training and validation and test loss
                     'worse_epochs' : 20, # Patience for early stoppping on validation set
                     }
     experiment_id = np.random.randint(0,1000000)

Datasets.py

@@ -173,7 +173,7 @@ def get_dataset(model_config, input_shape, output_shape, partition):
         # The dataset structure is a dictionary with "train", "valid", "test" keys, whose entries are lists, where each element represents a song.
         # Each song is represented as a dictionary containing elements mix, acc, vocal or mix, bass, drums, other, vocal depending on the task.
 
-        num_cores = 1
+        num_cores = 8
 
         for curr_partition in ["train", "val", "test"]:
             print("Writing " + curr_partition + " partition...")
@@ -224,7 +224,7 @@ def get_dataset(model_config, input_shape, output_shape, partition):
         dataset = dataset.repeat()
         dataset = dataset.shuffle(buffer_size=model_config["cache_size"])
 
-    dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(model_config["batch_size"]))
+    dataset = dataset.batch(model_config["batch_size"], drop_remainder=True)
     dataset = dataset.prefetch(1)
 
     return dataset

Evaluate.py

@@ -45,7 +45,7 @@ def predict(audio, model_config, load_model):
     frame_pred = separator_func(tracks_ph, training=False, return_spectrogram=False, reuse=False)
 
     # Start session and queue input threads
-    sess = tf.Session()
+    sess = tf.compat.v1.Session()
     sess.run(tf.compat.v1.global_variables_initializer())
 
     # Load model

Test.py

@@ -37,7 +37,7 @@ def test(model_config, partition, model_folder, load_model):
     global_step = tf.compat.v1.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False, dtype=tf.int64)
 
     # Start session and queue input threads
-    sess = tf.Session()
+    sess = tf.compat.v1.Session()
     sess.run(tf.compat.v1.global_variables_initializer())
     writer = tf.compat.v1.summary.FileWriter(model_config["log_dir"] + os.path.sep +  model_folder, graph=sess.graph)
 

Training.py

@@ -1,4 +1,6 @@
 from sacred import Experiment
+from tqdm import tqdm
+
 from Config import config_ingredient
 import tensorflow as tf
 import numpy as np
@@ -74,7 +76,7 @@ def train(model_config, experiment_id, load_model=None):
     print("Sep_Vars: " + str(Utils.getNumParams(separator_vars)))
     print("Num of variables" + str(len(tf.compat.v1.global_variables())))
 
-    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
     with tf.control_dependencies(update_ops):
         with tf.compat.v1.variable_scope("separator_solver"):
             separator_solver = tf.compat.v1.train.AdamOptimizer(learning_rate=model_config["init_sup_sep_lr"]).minimize(separator_loss, var_list=separator_vars)
@@ -84,9 +86,9 @@ def train(model_config, experiment_id, load_model=None):
     sup_summaries = tf.compat.v1.summary.merge_all(key='sup')
 
     # Start session and queue input threads
-    config = tf.ConfigProto()
+    config = tf.compat.v1.ConfigProto()
     config.gpu_options.allow_growth=True
-    sess = tf.Session(config=config)
+    sess = tf.compat.v1.Session(config=config)
     sess.run(tf.compat.v1.global_variables_initializer())
     writer = tf.compat.v1.summary.FileWriter(model_config["log_dir"] + os.path.sep + str(experiment_id),graph=sess.graph)
 
@@ -98,12 +100,12 @@ def train(model_config, experiment_id, load_model=None):
         restorer.restore(sess, load_model)
         print('Pre-trained model restored from file ' + load_model)
 
-    saver = tf.train.Saver(tf.compat.v1.global_variables(), write_version=tf.compat.v1.train.SaverDef.V2)
+    saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), write_version=tf.compat.v1.train.SaverDef.V2)
 
     # Start training loop
     _global_step = sess.run(global_step)
     _init_step = _global_step
-    for _ in range(model_config["epoch_it"]):
+    for _ in tqdm(range(model_config["epoch_it"])):
         # TRAIN SEPARATOR
         #try:
         _, _sup_summaries = sess.run([separator_solver, sup_summaries])

Utils.py

@@ -3,7 +3,7 @@
 import librosa
 
 def getTrainableVariables(tag=""):
-    return [v for v in tf.trainable_variables() if tag in v.name]
+    return [v for v in tf.compat.v1.trainable_variables() if tag in v.name]
 
 def getNumParams(tensors):
     return np.sum([np.prod(t.get_shape().as_list()) for t in tensors])