模型调试

0713f4c7 · 宋柯 · 71064fbb · 0713f4c7
Commit 0713f4c7 authored 3 years ago by 宋柯
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 3 deletions

train_service_sk_tfrecord.py train/train_service_sk_tfrecord.py +3 -3

No files found.
--- a/train/train_service_sk_tfrecord.py
+++ b/train/train_service_sk_tfrecord.py
@@ -7,7 +7,7 @@ os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
 BASE_DIR = '/data/files/wideAndDeep/'

 def input_fn(csv_path, epoch, shuffle, batch_size):
-    dataset = tf.data.TFRecordDataset(csv_path, buffer_size = 1024, num_parallel_reads = 2)
+    dataset = tf.data.TFRecordDataset(csv_path)

    dics = {
        'ITEM_CATEGORY_card_id': tf.FixedLenFeature((), tf.string, default_value='-1'),
@@ -105,9 +105,9 @@ def input_fn(csv_path, epoch, shuffle, batch_size):

    dataset = dataset.padded_batch(batch_size, padded_shapes, padding_values = padding_values)

-    dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    dataset = dataset.repeat(epoch)

-    dataset.repeat(epoch)
+    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

    return dataset