update condig

src/main.py

@@ -21,9 +21,9 @@ from models.esmm.model import esmm_model_fn, model_export, model_predict_diary
 def main():
     time_begin = time.time()
 
-    tf.logging.set_verbosity(tf.logging.INFO)
+    tf.logging.set_verbosity(tf.compat.v1.logging.INFO)
 
-    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
     device_df, diary_df, click_df, conversion_df = read_csv_data(Path("~/data/cvr_data/"))
     # print(diary_df.sample(1))
@@ -43,7 +43,12 @@ def main():
     # if os.path.exists(model_path):
     #     shutil.rmtree(model_path)
 
-    model = tf.estimator.Estimator(model_fn=esmm_model_fn, params=params, model_dir=model_path)
+    session_config = tf.ConfigProto()
+    session_config.gpu_options.allow_growth = True
+    session_config.gpu_options.per_process_gpu_memory_fraction = 0.9
+    estimator_config = tf.estimator.RunConfig(session_config=session_config)
+
+    model = tf.estimator.Estimator(model_fn=esmm_model_fn, params=params, model_dir=model_path, config=estimator_config)
     train_spec = tf.estimator.TrainSpec(input_fn=lambda: esmm_input_fn(train_df, shuffle=True), max_steps=50000)
     eval_spec = tf.estimator.EvalSpec(input_fn=lambda: esmm_input_fn(val_df, shuffle=False))
     tf.estimator.train_and_evaluate(model, train_spec, eval_spec)