Repository navigation
Expand file tree
/
Copy pathpreprocessing.py
More file actions
207 lines (177 loc) · 7.03 KB
/
Copy pathpreprocessing.py
File metadata and controls
207 lines (177 loc) · 7.03 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
# Copyright 2018 Google LLC
#
# 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.
'''Utilities to create, read, write tf.Examples.'''
import functools
import numpy as np
import tensorflow as tf
import random
import coords
import features as features_lib
import go
import sgf_wrapper
TF_RECORD_CONFIG = tf.python_io.TFRecordOptions(
tf.python_io.TFRecordCompressionType.ZLIB)
SHUFFLE_BUFFER_SIZE = int(4 * 1e6)
# Constructing tf.Examples
def _one_hot(index):
* go.N + 1], dtype=np.float32)
onehot[index] = 1
return onehot
def make_tf_example(features, pi, value):
'''
Args:
features: [N, N, FEATURE_DIM] nparray of uint8
pi: [N * N + 1] nparray of float32
value: float
'''
return tf.train.Example(features=tf.train.Features(feature={
'x': tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[features.tostring()])),
'pi': tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[pi.tostring()])),
'outcome': tf.train.Feature(
float_list=tf.train.FloatList(
value=[value]))}))
# Write tf.Example to files
def write_tf_examples(filename, tf_examples, serialize=True):
'''
Args:
filename: Where to write tf.records
tf_examples: An iterable of tf.Example
serialize: whether to serialize the examples.
'''
with tf.python_io.TFRecordWriter(
filename, options=TF_RECORD_CONFIG) as writer:
for ex in tf_examples:
if serialize:
writer.write(ex.SerializeToString())
else:
writer.write(ex)
# Read tf.Example from files
def batch_parse_tf_example(batch_size, example_batch):
'''
Args:
example_batch: a batch of tf.Example
Returns:
A dict of batched tensors
'''
features = {
'x': tf.FixedLenFeature([], tf.string),
'pi': tf.FixedLenFeature([], tf.string),
'outcome': tf.FixedLenFeature([], tf.float32),
}
parsed = tf.parse_example(example_batch, features)
x = tf.decode_raw(parsed['x'], tf.uint8)
x = tf.cast(x, tf.float32)
x = tf.reshape(x, [batch_size, go.N, go.N,
features_lib.NEW_FEATURES_PLANES])
pi = tf.decode_raw(parsed['pi'], tf.float32)
pi = tf.reshape(pi, [batch_size, go.N * go.N + 1])
outcome = parsed['outcome']
outcome.set_shape([batch_size])
return {
'pos_tensor': x,
'pi_tensor': pi,
'value_tensor': outcome,
}
def read_tf_records(batch_size, tf_records, num_repeats=None,
shuffle_records=True, shuffle_examples=True,
filter_amount=1.0):
'''
Args:
batch_size: batch size to return
tf_records: a list of tf_record filenames
num_repeats: how many times the data should be read (default: infinite)
shuffle_records: whether to shuffle the order of files read
shuffle_examples: whether to shuffle the tf.Examples
filter_amount: what fraction of records to keep
Returns:
a tf dataset of batched tensors
'''
# compression_type here must agree with write_tf_examples
# cycle_length = how many tfrecord files are read in parallel
# block_length = how many tf.Examples are read from each file before
# moving to the next file
# The idea is to shuffle both the order of the files being read,
# and the examples being read from the files.
if shuffle_records:
random.shuffle(tf_records)
record_list = tf.data.Dataset.from_tensor_slices(tf_records)
dataset = record_list.interleave(lambda x:
tf.data.TFRecordDataset(
x, compression_type='ZLIB'),
cycle_length=64, block_length=16)
dataset = dataset.filter(lambda x: tf.less(
tf.random_uniform([1]), filter_amount)[0])
# TODO(amj): apply py_func for transforms here.
if num_repeats is not None:
dataset = dataset.repeat(num_repeats)
else:
dataset = dataset.repeat()
if shuffle_examples:
dataset = dataset.shuffle(buffer_size=SHUFFLE_BUFFER_SIZE)
dataset = dataset.batch(batch_size)
return dataset
def get_input_tensors(batch_size, tf_records, num_repeats=None,
shuffle_records=True, shuffle_examples=True,
filter_amount=0.05):
'''Read tf.Records and prepare them for ingestion by dual_net
Returns a dict of tensors (see return value of batch_parse_tf_example)
'''
dataset = read_tf_records(batch_size, tf_records, num_repeats=num_repeats,
shuffle_records=shuffle_records,
shuffle_examples=shuffle_examples,
filter_amount=filter_amount)
dataset = dataset.filter(lambda t: tf.equal(tf.shape(t)[0], batch_size))
dataset = dataset.map(functools.partial(
batch_parse_tf_example, batch_size))
return dataset.make_one_shot_iterator().get_next()
# End-to-end utility functions
def make_dataset_from_selfplay(data_extracts):
'''
Returns an iterable of tf.Examples.
Args:
data_extracts: An iterable of (position, pi, result) tuples
'''
tf_examples = (make_tf_example(features_lib.extract_features(pos), pi, result)
for pos, pi, result in data_extracts)
return tf_examples
def make_dataset_from_sgf(sgf_filename, tf_record):
pwcs = sgf_wrapper.replay_sgf_file(sgf_filename)
tf_examples = map(_make_tf_example_from_pwc, pwcs)
write_tf_examples(tf_record, tf_examples)
def _make_tf_example_from_pwc(position_w_context):
features = features_lib.extract_features(position_w_context.position)
pi = _one_hot(coords.flatten_coords(position_w_context.next_move))
value = position_w_context.result
return make_tf_example(features, pi, value)
def shuffle_tf_examples(gather_size, records_to_shuffle):
'''Read through tf.Record and yield shuffled, but unparsed tf.Examples
Args:
gather_size: The number of tf.Examples to be gathered together
records_to_shuffle: A list of filenames
Returns:
An iterator yielding lists of bytes, which are serialized tf.Examples.
'''
dataset = read_tf_records(gather_size, records_to_shuffle, num_repeats=1)
batch = dataset.make_one_shot_iterator().get_next()
sess = tf.Session()
while True:
try:
result = sess.run(batch)
yield list(result)
except tf.errors.OutOfRangeError:
break