"awq/vscode:/vscode.git/clone" did not exist on "661047f119c82324fcda3dbe11eaf5b7e233a359"
Commit 68612afb authored by Eli Bixby's avatar Eli Bixby
Browse files

Use context manager

parent a2931072
...@@ -55,7 +55,7 @@ def _get_file_names(): ...@@ -55,7 +55,7 @@ def _get_file_names():
def read_pickle_from_file(filename): def read_pickle_from_file(filename):
with open(filename, 'r') as f: with tf.gfile.Open(filename, 'r') as f:
data_dict = cPickle.load(f) data_dict = cPickle.load(f)
return data_dict return data_dict
...@@ -63,8 +63,7 @@ def read_pickle_from_file(filename): ...@@ -63,8 +63,7 @@ def read_pickle_from_file(filename):
def convert_to_tfrecord(input_files, output_file): def convert_to_tfrecord(input_files, output_file):
"""Converts a file to tfrecords.""" """Converts a file to tfrecords."""
print('Generating %s' % output_file) print('Generating %s' % output_file)
record_writer = tf.python_io.TFRecordWriter(output_file) with tf.python_io.TFRecordWriter(output_file) as record_writer:
for input_file in input_files: for input_file in input_files:
data_dict = read_pickle_from_file(input_file) data_dict = read_pickle_from_file(input_file)
data = data_dict['data'] data = data_dict['data']
...@@ -78,7 +77,6 @@ def convert_to_tfrecord(input_files, output_file): ...@@ -78,7 +77,6 @@ def convert_to_tfrecord(input_files, output_file):
'label': _int64_feature(labels[i]) 'label': _int64_feature(labels[i])
})) }))
record_writer.write(example.SerializeToString()) record_writer.write(example.SerializeToString())
record_writer.close()
def main(unused_argv): def main(unused_argv):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment