这段代码实现了一个用于图像文本识别的深度学习训练框架。主要使用了Keras和TensorFlow来构建和训练一个DenseNet模型。以下是代码的详细解释
导入必要的库
#-*- coding:utf-8 -*-
import os
import json
import threading
from keras import losses
from keras.utils import plot_model
from keras.preprocessing import image
from keras.utils import pad_sequences
from tensorflow.python.keras.layers.recurrent import GRU, LSTM
from tensorflow.keras.layers import Bidirectional, TimeDistributed
from tensorflow.keras.layers import BatchNormalization
from keras.layers.convolutional import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.optimizers import SGD, Adam
import numpy as np
from PIL import Image
import tensorflow as tf
from keras import backend as K
from keras.layers import Input, Dense, Flatten
from keras.layers.core import Reshape, Masking, Lambda, Permute
from keras.models import Model
from keras.callbacks import EarlyStopping, ModelCheckpoint, LearningRateScheduler, TensorBoard
from imp import reload
import densenet
这部分导入了各种库和模块，包括Keras中的层、损失函数、优化器、回调函数等。还包括一些辅助库，如NumPy、PIL、TensorFlow和自定义模块densenet。
GPU配置函数
def get_session(gpu_fraction=1.0):
num_threads = os.environ.get('OMP_NUM_THREADS')
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)
if num_threads:
return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, intra_op_parallelism_threads=num_threads))
else:
return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
该函数配置TensorFlow的GPU使用。可以设置每个GPU进程使用的显存比例，并返回一个TensorFlow会话，如果是用独立显卡的话，比较麻烦，具体方式可以去GOOGLE
读取文件函数
def readfile(filename):
res=[]
with open(filename,'r') as f:
lines=f.readlines()
for i in lines:
res.append(i.strip())
dic={}
for i in res:
p=i.split(' ')
dic[p[0]] = p[1:]
return dic
该函数读取文件并解析成字典格式，其中键为文件名，值为标签。
随机生成均匀分布的数
class random_uniform_num():
"""
均匀随机，确保每轮每个只出现一次
"""
def __init__(self, total):
self.total = total
self.range = [i for i in range(total)]
np.random.shuffle(self.range)
self.index = 0
def get(self, batchsize):
r_n = []
if (self.index + batchsize > self.total):
r_n_1 = self.range[self.index:self.total]
np.random.shuffle(self.range)
self.index = (self.index + batchsize) - self.total
r_n_2 = self.range[0:self.index]
r_n.extend(r_n_1)
r_n.extend(r_n_2)
else:
r_n = self.range[self.index:self.index + batchsize]
self.index = self.index + batchsize
return r_n
该类用于生成一个均匀分布的随机数序列，确保每轮训练中数据只出现一次。
数据生成器
def gen(data_file, image_path, batchsize=128, maxlabellength=10, imagesize=(32, 280)):
image_label = readfile(data_file)
_imagefile = [i for i, j in image_label.items()]
x = np.zeros((batchsize, imagesize[0], imagesize[1], 1), dtype=np.float)
labels = np.ones([batchsize, maxlabellength]) * 10000
input_length = np.zeros([batchsize, 1])
label_length = np.zeros([batchsize, 1])
r_n = random_uniform_num(len(_imagefile))
_imagefile = np.array(_imagefile)
while 1:
shufimagefile = _imagefile[r_n.get(batchsize)]
for i, j in enumerate(shufimagefile):
img1 = Image.open(os.path.join(image_path, j)).convert('L')
img = np.array(img1, 'f') / 255.0 - 0.5
x[i] = np.expand_dims(img, axis=2)
str = image_label[j]
label_length[i] = len(str)
if (len(str) <= 0):
print("len <0", j)
input_length[i] = imagesize[1] // 8
labels[i, :len(str)] = [int(k) for k in str]
inputs = {
'the_input': x,
'the_labels': labels,
'input_length': input_length,
'label_length': label_length,
}
outputs = {'ctc': np.zeros([batchsize])}
yield (inputs, outputs)
这个生成器函数用于生成训练和测试所需的数据。每次调用都会返回一个批次的数据，包括输入图像、标签、输入长度和标签长度。
CTC损失函数
def ctc_lambda_func(args):
y_pred, labels, input_length, label_length = args
return K.ctc_batch_cost(labels, y_pred, input_length, label_length)
该函数定义了CTC（Connectionist Temporal Classification）损失函数，用于处理序列学习问题。
构建模型
def get_model(img_h, nclass):
input = Input(shape=(img_h, None, 1), name='the_input')
y_pred = densenet.dense_cnn(input, nclass)
labels = Input(name='the_labels', shape=[None], dtype='float32')
input_length = Input(name='input_length', shape=[1], dtype='int64')
label_length = Input(name='label_length', shape=[1], dtype='int64')
loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([y_pred, labels, input_length, label_length])
model = Model(inputs=[input, labels, input_length, label_length], outputs=loss_out)
model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer='adam', metrics=['accuracy'])
return model
这个函数构建了一个用于图像文本识别的模型。使用自定义的DenseNet模型和CTC损失函数。模型的输入包括图像、标签、输入长度和标签长度，输出为CTC损失。
训练函数
def start_train(train_index_file_path, train_img_root_dir, test_index_file_path, test_img_root_dir, model_output_dir="./models", train_epochs=10):
img_height = 32
img_width = 200
batch_size = 128
char_set = open('E:\Machinelearningpractices\C3\labels.txt', 'r', encoding='utf-8').readlines()
char_set = ''.join([ch.strip('\n') for ch in char_set][1:] + ['卍'])
num_class = len(char_set) + 2
reload(densenet)
model = get_model(img_height, num_class)
modelPath = './models/pretrain_model.h5'
if os.path.exists(modelPath):
print("Loading pretrain model weights...")
model.load_weights(modelPath)
for i, layer in enumerate(model.layers):
print(i, layer.name)
for layer in model.layers[:80]:
layer.trainable = False
train_set_file = train_index_file_path
test_set_file = test_index_file_path
train_loader = gen(train_set_file, train_img_root_dir, batchsize=batch_size, maxlabellength=num_class, imagesize=(img_height, img_width))
test_loader = gen(test_set_file, test_img_root_dir, batchsize=batch_size, maxlabellength=num_class, imagesize=(img_height, img_width))
model_save_path = os.path.join(model_output_dir, 'weights_densenet-{epoch:02d}-{val_loss:.2f}.h5')
checkpoint = ModelCheckpoint(model_save_path, monitor='val_loss', save_best_only=True, save_weights_only=True)
lr_schedule = lambda epoch: 0.005 * 0.1 ** epoch
learning_rate = np.array([lr_schedule(i) for i in range(20)])
changelr = LearningRateScheduler(lambda epoch: float(learning_rate[epoch]))
earlystop = EarlyStopping(monitor='val_loss', patience=4, verbose=1)
train_num_lines = sum(1 for line in open(train_set_file))
test_num_lines = sum(1 for line in open(test_set_file))
model.fit_generator(train_loader, steps_per_epoch=train_num_lines // batch_size, epochs=int(train_epochs), initial_epoch=0, validation_data=test_loader, validation_steps=test_num_lines // batch_size, callbacks=[checkpoint, earlystop, changelr])
return "1"
这个函数定义了模型的训练过程：
设置图像的高度、宽度和批次大小。
读取字符集并计算类别数量。
重新加载自定义的DenseNet模块。
构建模型并加载预训练权重（如果存在）。
冻结前80层的权重。
生成训练和测试数据的生成器。
设置模型保存路径和回调函数。
计算训练和测试集的样本数量。
训练模型并使用生成器提供数据。
主函数
if __name__ == '__main__':
start_train("./data_3_train.txt", "./train_imgs", "./data_3_test.txt", "./test_imgs", "./models", 8)
当脚本作为主程序运行时，调用start_train函数开始训练。
总结
这段代码实现了一个用于图像文本识别的深度学习训练框架，利用Keras和TensorFlow构建和训练一个DenseNet模型，并使用CTC损失函数处理序列学习问题。通过定义数据生成器和训练过程，用户可以方便地进行模型训练和测试
评论区