import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import warnings
import datetime
import tensorflow as tf
import random
import pathlib
import numpy as np
from tensorflow.keras.callbacks import TensorBoard, EarlyStopping
warnings.filterwarnings("ignore")
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
def load_datasets(file_dir):
data_path = pathlib.Path(file_dir)
all_image_paths = list(data_path.glob('*/*'))
all_image_paths = [str(path) for path in all_image_paths] # 所有图片路径的列表
image_count = len(all_image_paths)
random.shuffle(all_image_paths) # 打散
label_names = sorted(item.name for item in data_path.glob('*/') if item.is_dir())
label_to_index = dict((name, index) for index, name in enumerate(label_names))
all_image_labels = [label_to_index[pathlib.Path(path).parent.name] for path in all_image_paths]
labels = np.zeros((image_count, 12),dtype=np.int32)
for num, each_label in enumerate(all_image_labels):
labels [num, each_label] = 1
train_image = []
train_label = []
for image, label in zip(all_image_paths, labels):
r_image, r_label = load_and_preprocess_from_path_label(image, label)
train_image.append(r_image)
train_label.append(r_label)
train_images = np.array(train_image)
train_labels = np.array(train_label)
print(train_label)
return train_images, train_labels
def load_and_preprocess_from_path_label(path, label):
image = tf.io.read_file(path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize(image, [224, 224])
image /= 255.0
return image, label
train_dir = './Train'
test_dir = './Test'
train_image, train_label = load_datasets(train_dir)
test_image, test_label = load_datasets(test_dir)
num_all = train_image.shape[0]
count = 1
save_dir = './logs'
time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
log_dir = os.path.join(save_dir, "loss_" + str(time_str))
logging = TensorBoard(log_dir)
early_stopping = EarlyStopping(monitor='val_loss', min_delta = 0, patience = 20, verbose = 1)
callbacks = [logging, early_stopping]
def generator(x,y,batch):
np.random.seed(1000)
state = np.random.get_state()
np.random.shuffle(x)
np.random.set_state(state)
np.random.shuffle(y)
global count
for i in range(x.shape[0]):
batch_x = x[i:i+batch,:,:,:]
batch_y = y[i:i+batch,:]
yield batch_x, batch_y
mobile = tf.keras.applications.MobileNet(include_top=False, weights='imagenet', input_shape=(224, 224, 3))
mobile.trainable = False
model = tf.keras.Sequential([
mobile,
tf.keras.layers.Dropout(0.5),
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(12, activation='softmax')
])
optimizers = tf.keras.optimizers.SGD(learning_rate=0.01,momentum=0.9)
model.compile(optimizer=optimizers,
loss='categorical_crossentropy',
metrics=['accuracy'])
batch_size = 4
generator(train_image,train_label,batch_size)
model.fit_generator(generator(train_image,train_label,batch_size),validation_data=generator(test_image,test_label,batch_size),
steps_per_epoch = train_image.shape[0]//batch_size, epochs=50,validation_steps=test_image.shape[0]//batch_size,\
callbacks = callbacks)
model.save("./model.h5")
Fit_generator训练模型
于 2023-06-06 14:55:33 首次发布
该代码段展示了如何利用TensorFlow库加载和预处理图像数据,配置GPU内存增长,构建基于MobileNet的图像分类模型,使用Keras的fit_generator方法进行训练,并应用早停策略优化。同时,还包含了模型保存和TensorBoard日志记录。
摘要由CSDN通过智能技术生成