This repository was archived by the owner on Jul 7, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 3.7k
Expand file tree
/
Copy pathproblem.py
More file actions
1086 lines (928 loc) · 36.8 KB
/
problem.py
File metadata and controls
1086 lines (928 loc) · 36.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# coding=utf-8
# Copyright 2023 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Base class for problem/dataset definitions."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import functools
import os
import random
import six
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import text_encoder
from tensor2tensor.utils import contrib
from tensor2tensor.utils import data_reader
from tensor2tensor.utils import hparam
from tensor2tensor.utils import metrics
from tensor2tensor.utils import mlperf_log
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator
# pylint: disable=g-import-not-at-top
try:
from tensorflow.contrib.tpu.python.tpu import tpu_config
except ImportError:
# TF 2.0 doesn't ship with contrib.
tpu_config = None
# pylint: enable=g-import-not-at-top
class DatasetSplit(object):
TRAIN = tf_estimator.ModeKeys.TRAIN
EVAL = tf_estimator.ModeKeys.EVAL
TEST = "test"
class SpaceID(object):
"""Input and target space ids. Add more as needed."""
# Generic / unknown output space (default)
GENERIC = 0
# Image labels
IMAGE_LABEL = 1
# English characters
EN_CHR = 2
# English tokens
EN_TOK = 3
# English bpe tokens
EN_BPE_TOK = 4
# French characters
FR_CHR = 5
# French tokens
FR_TOK = 6
# German characters
DE_CHR = 7
# German tokens
DE_TOK = 8
# German bpe tokens
DE_BPE_TOK = 9
# Digit cipher lexicon 0
DIGIT_0 = 10
# Digit cipher lexicon 1
DIGIT_1 = 11
# Audio waveform domain
AUDIO_WAV = 12
# Audio spectral domain
AUDIO_SPECTRAL = 13
# Parse characters
PARSE_CHR = 14
# Parse tokens
PARSE_TOK = 15
# Chinese tokens
ZH_TOK = 16
# Icelandic characters
ICE_CHAR = 17
# Icelandic tokens
ICE_TOK = 18
# Icelandic parse tokens
ICE_PARSE_TOK = 19
# Macedonian tokens
MK_TOK = 20
# Czech tokens
CS_TOK = 21
# Czech characters
CS_CHR = 22
# Genetic bases (ACTG)
DNA = 23
# Real numbers
REAL = 24
# Images
IMAGE = 25
# Peptide
PEPTIDE = 26
# Python
PY_TOK = 27
# C++
CPP_TOK = 28
# Strokes
STROKES = 29
# Pickled Python
PICKLED_PYTHON = 30
class TaskID(object):
"""Problem specific task ids. Add more as needed."""
# English characters
EN_CHR = 2
# English characters sentiment
EN_CHR_SENT = 3
# English Premise Hypothesis pair
EN_PR_HYP = 4
# English NLI
EN_NLI = 5
# COLA
COLA = 6
# Enligh Question Context pair
EN_Q_CONT = 7
# English similarity task
EN_SIM = 8
# English sentence pair
EN_SENT_PAIR = 9
# 3 class NLI
THREE_CL_NLI = 10
def default_model_hparams():
return hparam.HParams(
max_input_seq_length=0,
max_target_seq_length=0,
prepend_mode="none",
split_to_length=0,
data_dir=None)
def preprocess_example_common(example, mode, hparams):
"""Preprocessing steps common to all models."""
if "inputs" in example and hparams.max_input_seq_length > 0:
example["inputs"] = example["inputs"][:hparams.max_input_seq_length]
if hparams.prepend_mode != "none":
if mode == tf_estimator.ModeKeys.PREDICT:
example["partial_targets"] = tf.concat([example["inputs"], [0]], 0)
else:
example["targets"] = tf.concat(
[example["inputs"], [0], example["targets"]], 0)
if "targets" in example and hparams.max_target_seq_length > 0:
example["targets"] = example["targets"][:hparams.max_target_seq_length]
if hparams.split_to_length:
new_example = {}
for k, v in six.iteritems(example):
if k == "targets" or k == "inputs":
new_example[k] = tf.reshape(v, [-1, hparams.split_to_length, 1, 1])
else:
tf.logging.warning("Dropping feature %s" % k)
return tf.data.Dataset.from_tensor_slices(new_example)
return example
class Problem(object):
"""Problem base class. Specifies a T2T problem.
Problems unify the specification of a problem for data generation, training,
and inference.
New problems are specified by the following methods:
Data generation:
* generate_data(data_dir, tmp_dir)
- Generate training and dev datasets into data_dir.
- Additional files, e.g. vocabulary files, should also be written to
data_dir. Vocab files are newline-separated files with each line
containing a token. The standard convention for the filename is to
set it to be
${Problem.vocab_filename}.${Problem.targeted_vocab_size}
- Downloads and other files can be written to tmp_dir
- If you have a training and dev generator, you can generate the
training and dev datasets with
generator_utils.generate_dataset_and_shuffle.
- Use the self.training_filepaths and self.dev_filepaths functions to
get sharded filenames. If shuffled=False, the filenames will contain
an "unshuffled" suffix; you should then shuffle the data
shard-by-shard with generator_utils.shuffle_dataset.
- Allows to specify the number of shards, optionally (can be omitted).
- Subclasses must override
* dataset_filename()
- Base filename for problem.
- Defaults to registered name (self.name).
Training:
* hparams(defaults, model_hparams)
- Specify the problem hyperparameters (see _default_hparams)
- Mutate defaults as needed
* example_reading_spec
- Specify the names and types of the features on disk.
- Specify tf.contrib.slim.tfexample_decoder
* preprocess_example(example, mode, hparams)
- Preprocess the example feature dict from feature name to Tensor or
SparseTensor.
- Used in training, eval, and inference (specified by mode).
Eval:
* eval_metrics
- Specify the set of evaluation metrics for this problem.
* eval_hooks
- Specify the set of evalueation hooks for this problem.
Inference:
* feature_encoders(data_dir)
- Return a dict of <feature name, TextEncoder> for encoding and decoding
inference input/output.
- Defaults to TextEncoder for inputs and targets.
"""
# ============================================================================
# BEGIN SUBCLASS INTERFACE
# ============================================================================
def generate_data(self, data_dir, tmp_dir, task_id=-1):
raise NotImplementedError()
@property
def multiprocess_generate(self):
"""Whether to generate the data in multiple parallel processes."""
return False
@property
def num_generate_tasks(self):
"""Needed if multiprocess_generate is True."""
raise NotImplementedError()
@property
def num_training_examples(self):
"""Used when mixing problems - how many examples are in the dataset."""
raise NotImplementedError()
def prepare_to_generate(self, data_dir, tmp_dir):
"""Prepare to generate data in parallel on different processes.
This function is called if multiprocess_generate is True.
Some things that might need to be done once are downloading the data
if it is not yet downloaded, and building the vocabulary.
Args:
data_dir: a string
tmp_dir: a string
"""
raise NotImplementedError()
def hparams(self, defaults, model_hparams):
pass
def max_length(self, model_hparams):
"""Maximum sequence length.
Problems with fixed length should override.
Args:
model_hparams: model hyperparameters
Returns:
an integer
"""
return (model_hparams.split_to_length or model_hparams.max_length or
model_hparams.batch_size)
def tpu_batch_size_per_shard(self, model_hparams):
"""Batch size in examples per TPU core.
Args:
model_hparams: model hyperparameters
Returns:
an integer
"""
if self.batch_size_means_tokens and not model_hparams.use_fixed_batch_size:
return model_hparams.batch_size // self.max_length(model_hparams)
else:
return model_hparams.batch_size
@property
def batch_size_means_tokens(self):
"""Do we specify hparams.batch_size in tokens per datashard per batch.
This is generally done for text problems.
If False, we assume that batch sizes are specified in examples per
datashard per batch.
TODO(noam): we should be more explicit and replace the hyperparameter
batch size with two hyperparameters:
hparams.examples_per_batch_per_datashard
hparams.tokens_per_batch_per_datashard
Returns:
a boolean
"""
return False
@property
def skip_random_fraction_when_training(self):
"""Skip a random number of examples at the beginning of training."""
# Skip a random fraction at the beginning of the stream. The skip is
# essential for synchronous highly-parallel training to avoid multiple
# replicas reading the same data in lock-step. So keep this true unless
# you have a very specific setting in which it needs to be turned off.
return True
def dataset_filename(self):
return self.name
def feature_encoders(self, data_dir):
del data_dir
return {
"inputs": text_encoder.TextEncoder(),
"targets": text_encoder.TextEncoder()
}
def example_reading_spec(self):
"""Define how data is serialized to file and read back.
Returns:
data_fields: A dictionary mapping data names to its feature type.
data_items_to_decoders: A dictionary mapping data names to TF Example
decoders, to be used when reading back TF examples from disk.
"""
data_fields = {
"inputs": tf.VarLenFeature(tf.int64),
"targets": tf.VarLenFeature(tf.int64)
}
data_items_to_decoders = None
return (data_fields, data_items_to_decoders)
def preprocess_example(self, example, mode, hparams):
"""Runtime preprocessing.
Return a dict or a tf.data.Dataset.from_tensor_slices (if you want each
example to turn into multiple).
Args:
example: dict, features
mode: tf.estimator.ModeKeys
hparams: HParams, model hyperparameters
Returns:
dict or Dataset
"""
return preprocess_example_common(example, mode, hparams)
def eval_metrics(self):
return [
metrics.Metrics.ACC, metrics.Metrics.ACC_TOP5,
metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
]
@property
def all_metrics_fns(self):
return metrics.METRICS_FNS
def eval_metric_fns(self, model_hparams):
del model_hparams
metric_names = self.eval_metrics()
if not all([m in self.all_metrics_fns for m in metric_names]):
error_str = ("Unrecognized metric. Problem %s specified metrics "
"%s. Recognized metrics are %s.")
raise ValueError(error_str % (self.name,
metric_names,
list(self.all_metrics_fns.keys())))
return {
metric_name: self.all_metrics_fns[metric_name]
for metric_name in metric_names
}
def eval_hooks(self, features, logits, hparams):
del features, logits, hparams
return []
@property
def task_id(self):
if self._task_id == -1 and hasattr(self, "global_task_id"):
self._task_id = self.global_task_id()
return self._task_id
def set_task_id(self, new_task_id):
self._task_id = new_task_id
# ============================================================================
# END SUBCLASS INTERFACE
# ============================================================================
@tf.autograph.experimental.do_not_convert()
def preprocess(self, dataset, mode, hparams, interleave=True):
"""Runtime preprocessing on the whole dataset.
Return a tf.data.Dataset -- the preprocessed version of the given one.
By default this function calls preprocess_example.
Args:
dataset: the Dataset of already decoded but not yet preprocessed features.
mode: tf.estimator.ModeKeys
hparams: HParams, model hyperparameters
interleave: bool, whether to use parallel_interleave, which is faster
but will alter the order of samples non-deterministically, or flat_map,
which is slower but will preserve the sample order.
Returns:
a Dataset
"""
def _preprocess(example):
examples = self.preprocess_example(example, mode, hparams)
if not isinstance(examples, tf.data.Dataset):
examples = tf.data.Dataset.from_tensors(examples)
return examples
if interleave:
dataset = dataset.apply(
tf.data.experimental.parallel_interleave(
_preprocess, sloppy=True, cycle_length=8))
else:
dataset = dataset.flat_map(_preprocess)
return dataset
def training_filepaths(self, data_dir, num_shards, shuffled):
file_basename = self.dataset_filename()
if not shuffled:
file_basename += generator_utils.UNSHUFFLED_SUFFIX
return generator_utils.train_data_filenames(file_basename, data_dir,
num_shards)
def dev_filepaths(self, data_dir, num_shards, shuffled):
file_basename = self.dataset_filename()
if not shuffled:
file_basename += generator_utils.UNSHUFFLED_SUFFIX
return generator_utils.dev_data_filenames(file_basename, data_dir,
num_shards)
def test_filepaths(self, data_dir, num_shards, shuffled):
file_basename = self.dataset_filename()
if not shuffled:
file_basename += generator_utils.UNSHUFFLED_SUFFIX
return generator_utils.test_data_filenames(file_basename, data_dir,
num_shards)
def data_filepaths(self, split, output_dir, num_shards, shuffled):
if split == DatasetSplit.TRAIN:
return self.training_filepaths(output_dir, num_shards, shuffled)
elif split == DatasetSplit.EVAL:
return self.dev_filepaths(output_dir, num_shards, shuffled)
elif split == DatasetSplit.TEST:
return self.test_filepaths(output_dir, num_shards, shuffled)
else:
raise ValueError("Unknown value for split: %s" % split)
def filepattern(self, data_dir, mode, shard=None):
"""Get filepattern for data files for mode.
Matches mode to a suffix.
* DatasetSplit.TRAIN: train
* DatasetSplit.EVAL: dev
* DatasetSplit.TEST: test
* tf.estimator.ModeKeys.PREDICT: dev
Args:
data_dir: str, data directory.
mode: DatasetSplit
shard: int, if provided, will only read data from the specified shard.
Returns:
filepattern str
"""
path = os.path.join(data_dir, self.dataset_filename())
shard_str = "-%05d" % shard if shard is not None else ""
if mode == DatasetSplit.TRAIN:
suffix = "train"
elif mode in [DatasetSplit.EVAL, tf_estimator.ModeKeys.PREDICT]:
suffix = "dev"
else:
assert mode == DatasetSplit.TEST
suffix = "test"
return "%s-%s%s*" % (path, suffix, shard_str)
def __init__(self, was_reversed=False, was_copy=False):
"""Create a Problem.
Args:
was_reversed: bool, whether to reverse inputs and targets.
was_copy: bool, whether to copy inputs to targets. Can be composed with
was_reversed so that if both are true, the targets become the inputs,
which are then copied to targets so that the task is targets->targets.
"""
self._was_reversed = was_reversed
self._was_copy = was_copy
self._encoders = None
self._hparams = None
self._feature_info = None
self._task_id = -1
@property
def was_reversed(self):
"""Whether the problem was reversed."""
return self._was_reversed
def get_feature_encoders(self, data_dir=None):
if self._encoders is None:
self._encoders = self.feature_encoders(data_dir)
return self._encoders
def get_hparams(self, model_hparams=None):
"""Returns problem_hparams."""
if self._hparams is not None:
return self._hparams
if model_hparams is None:
model_hparams = default_model_hparams()
if self._encoders is None:
data_dir = (model_hparams and hasattr(model_hparams, "data_dir") and
model_hparams.data_dir) or None
self.get_feature_encoders(data_dir)
hp = _default_hparams()
ret = self.hparams(hp, model_hparams)
if ret is not None:
raise ValueError("The Problem subclass hparams function should mutate "
"the defaults passed in and return None.")
hp.add_hparam("vocabulary", self._encoders)
hp.add_hparam("was_reversed", self._was_reversed)
hp.add_hparam("was_copy", self._was_copy)
if self._was_reversed:
_reverse_problem_hparams(hp)
if self._was_copy:
_copy_problem_hparams(hp)
self._hparams = hp
return self._hparams
def maybe_reverse_features(self, feature_map):
"""Reverse features between inputs and targets if the problem is '_rev'."""
if not self._was_reversed:
return
inputs = feature_map.pop("inputs", None)
targets = feature_map.pop("targets", None)
inputs_seg = feature_map.pop("inputs_segmentation", None)
targets_seg = feature_map.pop("targets_segmentation", None)
inputs_pos = feature_map.pop("inputs_position", None)
targets_pos = feature_map.pop("targets_position", None)
if inputs is not None:
feature_map["targets"] = inputs
if targets is not None:
feature_map["inputs"] = targets
if inputs_seg is not None:
feature_map["targets_segmentation"] = inputs_seg
if targets_seg is not None:
feature_map["inputs_segmentation"] = targets_seg
if inputs_pos is not None:
feature_map["targets_position"] = inputs_pos
if targets_pos is not None:
feature_map["inputs_position"] = targets_pos
def maybe_copy_features(self, feature_map):
if not self._was_copy:
return
feature_map["targets"] = feature_map["inputs"]
if ("inputs_segmentation" in feature_map and
"targets_segmentation" not in feature_map):
feature_map["targets_segmentation"] = feature_map["inputs_segmentation"]
if ("inputs_position" in feature_map and
"targets_position" not in feature_map):
feature_map["targets_position"] = feature_map["inputs_position"]
def maybe_reverse_and_copy(self, example):
self.maybe_reverse_features(example)
self.maybe_copy_features(example)
return example
@tf.autograph.experimental.do_not_convert()
def dataset(self,
mode,
data_dir=None,
num_threads=None,
output_buffer_size=None,
shuffle_files=None,
hparams=None,
preprocess=True,
dataset_split=None,
shard=None,
partition_id=0,
num_partitions=1,
shuffle_buffer_size=1024,
max_records=-1):
"""Build a Dataset for this problem.
Args:
mode: tf.estimator.ModeKeys; determines which files to read from.
data_dir: directory that contains data files.
num_threads: int, number of threads to use for decode and preprocess
Dataset.map calls.
output_buffer_size: int, how many elements to prefetch at end of pipeline.
shuffle_files: whether to shuffle input files. Default behavior (i.e. when
shuffle_files=None) is to shuffle if mode == TRAIN.
hparams: HParams; hparams to be passed to
Problem.preprocess_example and Problem.hparams. If None, will use a
default set that is a no-op.
preprocess: bool, whether to map the Dataset through
Problem.preprocess_example.
dataset_split: DatasetSplit, which split to read data
from (TRAIN:"-train", EVAL:"-dev", "test":"-test"). Defaults to mode.
shard: int, if provided, will only read data from the specified shard.
partition_id: integer - which partition of the dataset to read from
num_partitions: how many partitions in the dataset
shuffle_buffer_size: if shuffle_files is True, this is the buffer size
used to shuffle records.
max_records: int, number of records to truncate to.
Returns:
Dataset containing dict<feature name, Tensor>.
Raises:
ValueError: if num_partitions is greater than the number of data files.
"""
is_training = mode == tf_estimator.ModeKeys.TRAIN
shuffle_files = shuffle_files or shuffle_files is None and is_training
dataset_split = dataset_split or mode
assert data_dir
if hparams is None:
hparams = default_model_hparams()
if not hasattr(hparams, "data_dir"):
hparams.add_hparam("data_dir", data_dir)
if not hparams.data_dir:
hparams.data_dir = data_dir
# Construct the Problem's hparams so that items within it are accessible
_ = self.get_hparams(hparams)
data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
tf.logging.info("Reading data files from %s", data_filepattern)
data_files = sorted(
contrib.slim().parallel_reader.get_data_files(data_filepattern))
# Functions used in dataset transforms below. `filenames` can be either a
# `tf.string` tensor or `tf.data.Dataset` containing one or more filenames.
def _load_records_and_preprocess(filenames):
"""Reads files from a string tensor or a dataset of filenames."""
# Load records from file(s) with an 8MiB read buffer.
dataset = tf.data.TFRecordDataset(filenames, buffer_size=8 * 1024 * 1024)
# Decode.
dataset = dataset.map(self.decode_example, num_parallel_calls=num_threads)
# Preprocess if requested.
# Note that preprocessing should happen per-file as order may matter.
if preprocess:
dataset = self.preprocess(dataset, mode, hparams,
interleave=shuffle_files)
return dataset
if len(data_files) < num_partitions:
raise ValueError(
"number of data files (%d) must be at least the number of hosts (%d)"
% (len(data_files), num_partitions))
data_files = [f for (i, f) in enumerate(data_files)
if i % num_partitions == partition_id]
tf.logging.info(
"partition: %d num_data_files: %d" % (partition_id, len(data_files)))
if shuffle_files:
mlperf_log.transformer_print(key=mlperf_log.INPUT_ORDER)
random.shuffle(data_files)
dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
# Create data-set from files by parsing, pre-processing and interleaving.
if shuffle_files:
dataset = dataset.apply(
tf.data.experimental.parallel_interleave(
_load_records_and_preprocess, sloppy=True, cycle_length=8))
else:
dataset = _load_records_and_preprocess(dataset)
dataset = dataset.map(
self.maybe_reverse_and_copy, num_parallel_calls=num_threads)
dataset = dataset.take(max_records)
## Shuffle records only for training examples.
if shuffle_files and is_training:
dataset = dataset.shuffle(shuffle_buffer_size)
if hparams.get("pack_dataset", False):
dataset = generator_utils.pack_dataset(
dataset, hparams.max_length, keys=["inputs", "targets"],
use_custom_ops=hparams.get("use_custom_ops", False))
if output_buffer_size:
dataset = dataset.prefetch(output_buffer_size)
return dataset
def decode_example(self, serialized_example):
"""Return a dict of Tensors from a serialized tensorflow.Example."""
data_fields, data_items_to_decoders = self.example_reading_spec()
# Necessary to rejoin examples in the correct order with the Cloud ML Engine
# batch prediction API.
data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
if getattr(self._hparams, "sampling_method", "") == "random_per_example":
data_fields["sampling_temp"] = tf.FixedLenFeature(
[1], tf.float32, getattr(self._hparams, "sampling_temp", 1.0))
data_fields["sampling_keep_top_k"] = tf.FixedLenFeature(
[1], tf.int64, getattr(self._hparams, "sampling_keep_top_k", -1))
if data_items_to_decoders is None:
data_items_to_decoders = {}
for field in data_fields:
if data_fields[field].dtype is tf.string:
default_value = b""
else:
default_value = 0
data_items_to_decoders[field] = contrib.slim().tfexample_decoder.Tensor(
field, default_value=default_value)
decoder = contrib.slim().tfexample_decoder.TFExampleDecoder(
data_fields, data_items_to_decoders)
decode_items = list(sorted(data_items_to_decoders))
decoded = decoder.decode(serialized_example, items=decode_items)
return dict(zip(decode_items, decoded))
@property
def decode_hooks(self):
"""List of functions to be run after full decodes have been produced.
Returns:
List of functions. Each function should expect a single argument, an
instance of decoding.DecodeHookArgs and optionally return a list of
tf.Summary.Value objects.
"""
return []
@property
def has_inputs(self):
return "inputs" in self.get_feature_encoders()
@property
def feature_info(self):
"""Retrieve dict<feature name, FeatureInfo>.
Must first call Problem.get_hparams or Problem.dataset to have the problem's
internal hparams already constructed.
Returns:
dict<feature name, FeatureInfo>
"""
if self._feature_info is not None:
return self._feature_info
assert self._hparams is not None
hp = self.get_hparams()
if self.has_inputs:
in_id = hp.input_space_id
out_id = hp.target_space_id
features = collections.defaultdict(FeatureInfo)
for feature_name, modality_cls in six.iteritems(hp.modality):
finfo = features[feature_name]
finfo.modality = modality_cls
finfo.vocab_size = hp.vocab_size[feature_name]
vocabs = hp.vocabulary
for name, encoder in six.iteritems(vocabs):
features[name].encoder = encoder
if self.has_inputs:
features["inputs"].space_id = in_id
features["targets"].space_id = out_id
self._feature_info = features
return features
def make_estimator_input_fn(self,
mode,
hparams,
data_dir=None,
force_repeat=False,
prevent_repeat=False,
dataset_kwargs=None):
"""Return input_fn wrapped for Estimator."""
def estimator_input_fn(params, config):
return self.input_fn(
mode,
hparams,
data_dir=data_dir,
params=params,
config=config,
force_repeat=force_repeat,
prevent_repeat=prevent_repeat,
dataset_kwargs=dataset_kwargs)
return estimator_input_fn
def _dataset_partition(self, mode, config, params):
"""Which part of the training data to read.
If there are multiple parallel calls to input_fn (multiple TPU hosts),
then we want each one to read from a separate partition of the training
data.
Args:
mode: tf.estimator.ModeKeys
config: RunConfig
params: A dict that contains parameters.
Returns:
partition_id: an integer
num_partitions: an integer
"""
if mode != tf_estimator.ModeKeys.TRAIN or not hasattr(config, "tpu_config"):
# Reset in the case when using TPU but alternating TRAIN and EVAL.
self._next_partition_id = 0
return 0, 1
phift = config.tpu_config.per_host_input_for_training
# This is the mesh-tensorflow case.
if (hasattr(tpu_config.InputPipelineConfig, "BROADCAST") and
phift == tpu_config.InputPipelineConfig.BROADCAST):
return 0, 1
if phift:
num_hosts = (params["context"].num_hosts if "context" in params
else config.tpu_config.num_shards // 8)
num_partitions = max(num_hosts, 1)
else:
num_partitions = config.tpu_config.num_shards
partition_id = getattr(self, "_next_partition_id", 0)
self._next_partition_id = partition_id + 1
tf.logging.info("num_partitions = %d partition_id = %d" %
(num_partitions, partition_id))
assert partition_id < num_partitions
return partition_id, num_partitions
def input_fn(self,
mode,
hparams,
data_dir=None,
params=None,
config=None,
force_repeat=False,
prevent_repeat=False,
dataset_kwargs=None):
"""Builds input pipeline for problem.
Args:
mode: tf.estimator.ModeKeys
hparams: HParams, model hparams
data_dir: str, data directory; if None, will use hparams.data_dir
params: dict, may include "batch_size"
config: RunConfig; should have the data_parallelism attribute if not using
TPU
force_repeat: bool, whether to repeat the data even if not training
prevent_repeat: bool, whether to not repeat when in training mode.
Overrides force_repeat.
dataset_kwargs: dict, if passed, will pass as kwargs to self.dataset
method when called
Returns:
(features_dict<str name, Tensor feature>, Tensor targets)
"""
partition_id, num_partitions = self._dataset_partition(mode, config, params)
is_training = mode == tf_estimator.ModeKeys.TRAIN
if config and config.use_tpu:
num_threads = 64
else:
num_threads = data_reader.cpu_count() if is_training else 1
data_dir = data_dir or (hasattr(hparams, "data_dir") and hparams.data_dir)
dataset_kwargs = dataset_kwargs or {}
dataset_kwargs.update({
"mode": mode,
"data_dir": data_dir,
"num_threads": num_threads,
"hparams": hparams,
"partition_id": partition_id,
"num_partitions": num_partitions,
})
return data_reader.input_fn(
self.dataset(**dataset_kwargs),
self.filepattern(data_dir, mode),
self.skip_random_fraction_when_training,
self.batch_size_means_tokens,
self.get_hparams().batch_size_multiplier,
self.max_length(hparams),
mode,
hparams,
data_dir=data_dir,
params=params,
config=config,
force_repeat=force_repeat,
prevent_repeat=prevent_repeat)
@property
def export_assets(self):
"""Assets to export with the model.
This property contains a dictionary of assets, such as vocabulary files,
that should be exported together with the model, or None if no assets
are needed.
"""
return None
def serving_input_fn(self, hparams, decode_hparams=None, use_tpu=False):
"""Input fn for serving export, starting from serialized example."""
self._hparams = hparams
mode = tf_estimator.ModeKeys.PREDICT
serialized_example = tf.placeholder(
dtype=tf.string, shape=[None], name="serialized_example")
dataset = tf.data.Dataset.from_tensor_slices(serialized_example)
dataset = dataset.map(self.decode_example)
dataset = dataset.map(lambda ex: self.preprocess_example(ex, mode, hparams))
dataset = dataset.map(data_reader.cast_ints_to_int32)
if use_tpu:
padded_shapes = data_reader.pad_for_tpu(dataset.output_shapes, hparams,
hparams.max_length)
batch_size = 1 if not decode_hparams else getattr(decode_hparams,
"batch_size", 1)
dataset = dataset.padded_batch(
batch_size, padded_shapes, drop_remainder=False)
dataset = dataset.map(
functools.partial(data_reader.pad_batch, batch_multiple=batch_size))
else:
dataset = dataset.padded_batch(
tf.shape(serialized_example, out_type=tf.int64)[0],
dataset.output_shapes)
dataset = dataset.map(data_reader.standardize_shapes)
features = tf.data.experimental.get_single_element(dataset)
if self.has_inputs:
features.pop("targets", None)
return tf_estimator.export.ServingInputReceiver(
features=features, receiver_tensors=serialized_example)
class FeatureInfo(object):
"""Encapsulates information about a feature."""
def __init__(self,
encoder=None,
modality=None,
vocab_size=None,
space_id=None):
self.encoder = encoder
self.modality = modality
self.vocab_size = vocab_size
self.space_id = space_id
def _copy_problem_hparams(p_hparams):
"""Use input modality, vocab, and space id for target."""
p = p_hparams
# Duplicate input modality.
p.modality["targets"] = p.modality["inputs"]
# Duplicate input vocab size.
p.vocab_size["targets"] = p.vocab_size["inputs"]
# Duplicate input vocabulary.
p.vocabulary["targets"] = p.vocabulary["inputs"]
# Duplicate input space ids.
p.target_space_id = p.input_space_id
# Mark that p was reversed.
p.was_copy = True
def _reverse_problem_hparams(p_hparams):
"""Swap input/output modalities, vocab, and space ids."""
p = p_hparams
# Swap modalities.
# TODO(trandustin): Note this assumes target modalities have feature name
# 'target', and each intended feature to swap has feature name 'input'.
# In the future, remove need for this behavior.
reversed_modality = {}
for feature_name in p.modality:
# Copy feature as-is.
if "target" not in feature_name and "input" not in feature_name:
reversed_modality[feature_name] = p.modality[feature_name]
else:
# Change "target" to "input" and vice-versa for this feature.
if "target" in feature_name:
reversed_feature_name = feature_name.replace("target", "input")