import datetime import os from pathlib import Path import numpy as np import tensorflow as tf from sklearn.metrics import roc_auc_score def print_curtime(note=None): now = datetime.datetime.now() current_time = now.strftime("%H:%M:%S") if note is not None: print(f"{note}: {current_time}") else: print(f"Current time: {current_time}") def tf_allow_growth(): gpus = tf.config.experimental.list_physical_devices("GPU") for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) def create_logdir(root="logs/", args=None): log_dir = root + datetime.datetime.now().strftime("%Y%m%d-%H%M%S") Path(log_dir).mkdir(parents=True, exist_ok=True) with open(os.path.join(log_dir, "config.txt"), "w") as f: print(args, file=f) print(f"LOG_DIR: {log_dir}") summary_writer = tf.summary.create_file_writer( os.path.join(log_dir, "train")) summary_writer.set_as_default() return log_dir def auc_score(y_true, y_pred): if len(np.unique(y_true[:, 0])) == 1: return 0.5 else: return roc_auc_score(y_true, y_pred) def auc(y_true, y_pred): return tf.numpy_function(auc_score, (y_true, y_pred), tf.double) def num_params(model): total_parameters = 0 embed_parameters = 0 dense_parameters = 0 for variable in model.trainable_variables: shape = variable.get_shape() variable_parameters = 1 for dim in shape: variable_parameters *= dim total_parameters += variable_parameters if 'embedding' in variable.name: embed_parameters += variable_parameters else: dense_parameters += variable_parameters print(f"Total Params: {total_parameters}") print(f"Dense Params: {dense_parameters}") print(f"Embed Params: {embed_parameters}") return total_parameters