模型调试

train/train_service_sk.py

 import tensorflow as tf
 tf.logging.set_verbosity(tf.logging.INFO)
 import sys
+import os
+os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
 
 BASE_DIR = '/data/files/wideAndDeep/'
 
@@ -79,9 +81,9 @@ def input_fn(csv_path, epoch, shuffle, batch_size):
     dataset = dataset.map(parse_line, num_parallel_calls=8)
     dataset = dataset.padded_batch(batch_size, padded_shapes, padding_values=padding_values)
     if shuffle:
-        dataset = dataset.shuffle(1000).prefetch(512 * 10).repeat(epoch)
+        dataset = dataset.shuffle(1000).prefetch(512 * 100).repeat(epoch)
     else:
-        dataset = dataset.prefetch(512 * 10).repeat(epoch)
+        dataset = dataset.prefetch(512 * 100).repeat(epoch)
 
     return dataset