Tensorflow代码学习-10-2生成tfrecord文件

生成tfrecord文件

• 课程：炼数成金
• GITHUB
• CSDN
• 知乎
• 欢迎点亮star

import os
import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import random,math,sys
from PIL import Image
import numpy as np

gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.333)

_NUM_TEST = 500
_RANDOM_SEED = 0
_NUM_SHARDS = 5

#dataset path
DATASET_DIR = r'captcha/images/'
#tfrecord where to save
TFRECORD_DIR = 'captcha/'

#判断 tfrecord is exists
def _dataset_exists(dataset_dir):
    for split_name in ['train', 'test']:
        output_filename = os.path.join(dataset_dir, split_name, '.tfrecords')
        if not tf.gfile.Exists(output_filename):
            return  False
    return True

#get all v_code's path
def _get_filename_and_classes(dataset_dir):
    photo_filenames = []
    for filename in os.listdir(dataset_dir):
        path = os.path.join(dataset_dir,filename)
        photo_filenames.append(path