TensorFlow2从磁盘读取图片数据集的示例(tf.data.Dataset.list_files)

python 复制代码
import os
import warnings
warnings.filterwarnings("ignore")
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.applications.resnet import ResNet50
from pathlib import Path
import numpy as np

#数据所在文件夹
base_dir = './data/cats_and_dogs'
train_dir = Path(os.path.join(base_dir,'train'))
file_pattern = os.path.join(train_dir,'*/*.jpg')
image_count = len(list(train_dir.glob('*/*.jpg')))

list_ds = tf.data.Dataset.list_files(file_pattern,shuffle = False)
list_ds = list_ds.shuffle(image_count, reshuffle_each_iteration=False)
for f in list_ds.take(5):
  print(f.numpy())
  
class_names = np.array(sorted([item.name for item in train_dir.glob('*') ]))
print(class_names)

val_size = int(image_count * 0.2)
train_data = list_ds.skip(val_size)
validation_data = list_ds.take(val_size)
print(tf.data.experimental.cardinality(train_data).numpy())
print(tf.data.experimental.cardinality(validation_data).numpy())


def get_label(file_path):
  parts = tf.strings.split(file_path, os.path.sep)
  one_hot = parts[-2] == class_names
  return tf.argmax(one_hot)

def decode_img(img):
  img = tf.io.decode_jpeg(img, channels=3)
  return tf.image.resize(img, [64, 64])

def process_path(file_path):
  label = get_label(file_path)
  img = tf.io.read_file(file_path)
  img = decode_img(img)
  return img, label

train_data = train_data.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)
validation_data = validation_data.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)

for image, label in train_data.take(2):
  print("Image shape: ", image.numpy().shape)
  print("Label: ", label.numpy())

def configure_for_performance(ds):
  ds = ds.cache()
  ds = ds.shuffle(buffer_size=1000)
  ds = ds.batch(4)
  ds = ds.prefetch(buffer_size=tf.data.AUTOTUNE)
  return ds

train_data = configure_for_performance(train_data)
validation_data = configure_for_performance(validation_data)


save_model_cb = tf.keras.callbacks.ModelCheckpoint(filepath='model_resnet50_cats_and_dogs.h5', save_freq='epoch')

base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(64, 64, 3))
base_model.trainable = True
    
model = tf.keras.models.Sequential([
    base_model,
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation='relu',kernel_regularizer=tf.keras.regularizers.l2(l=0.01)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(loss='binary_crossentropy',optimizer = Adam(lr=1e-3),metrics = ['acc'])

history = model.fit(train_data.repeat(),steps_per_epoch=100,epochs=50,validation_data=validation_data.repeat(),validation_steps=50,verbose=1,callbacks = [save_model_cb])
相关推荐
夕小瑶4 分钟前
Dexmal 原力灵机开源 Dexbotic:具身智能的“Transformers“库来了
大数据·人工智能
Mrliu__10 分钟前
Python数据结构(七):Python 高级排序算法:希尔 快速 归并
数据结构·python·排序算法
C嘎嘎嵌入式开发15 分钟前
(22)100天python从入门到拿捏《【网络爬虫】网络基础与HTTP协议》
网络·爬虫·python
飞飞是甜咖啡16 分钟前
SPP-CNN解决CNN只能处理固定大小的输入图片
人工智能·神经网络·cnn
xiaoxiaode_shu20 分钟前
神经网络基础
人工智能·深度学习·神经网络
小小爱大王1 小时前
AI 编码效率提升 10 倍的秘密:Prompt 工程 + 工具链集成实战
java·javascript·人工智能
盼小辉丶1 小时前
使用CNN构建VAE
深度学习·神经网络·cnn·生成模型
zzzyulin1 小时前
huggingface transformers调试问题--加载本地路径模型时pdb断点消失
python·transformer
教练、我想打篮球1 小时前
12 pyflink 的一个基础使用, 以及环境相关
python·flink·pyflink
飞翔的佩奇1 小时前
【完整源码+数据集+部署教程】【天线&运输】直升机战机类型识别目标检测系统源码&数据集全套:改进yolo11-CSP-EDLAN
前端·python·yolo·计算机视觉·数据集·yolo11·直升机战机类型识别目标检测系统