Fit_generator训练模型

该代码段展示了如何利用TensorFlow库加载和预处理图像数据,配置GPU内存增长,构建基于MobileNet的图像分类模型,使用Keras的fit_generator方法进行训练,并应用早停策略优化。同时,还包含了模型保存和TensorBoard日志记录。
摘要由CSDN通过智能技术生成
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")
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值