Save data as TFRecord

 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

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf


data_record = {
    'age': 10,
    'height': 49.2,
    'name': 'hello',
    'hobbies': ['comic', 'movie', 'lego']
}
int_list1 = tf.train.Int64List(value=[data_record['age']])
float_list1 = tf.train.FloatList(value=[data_record['height']])
# Convert string data into list of bytes
str_list1 = tf.train.BytesList(value=[data_record['name'].encode('utf-8')])
str_list2 = tf.train.BytesList(value=[i.encode('utf-8') for i in data_record['hobbies']])
# Create a dictionary with above lists individually wrapped in Feature
feature_key_value_pair = {
    'int_list1': tf.train.Feature(int64_list=int_list1),
    'float_list1': tf.train.Feature(float_list=float_list1),
    'str_list1': tf.train.Feature(bytes_list=str_list1),
    'float_list2': tf.train.Feature(bytes_list=str_list2)
}
features = tf.train.Features(feature=feature_key_value_pair)
example = tf.train.Example(features=features)
with tf.io.TFRecordWriter('example.tfrecord') as tfwriter:
    tfwriter.write(example.SerializeToString())