基于度量学习的行人重识别（三）-技术圈

在前面的文章中，我们已经完成了对数据的处理，以及模型的搭建，接下来就需要整合起来，进行训练了。

模型训练

新建一个train.py,首先导入依赖：

from  load_data import Person_Datasetfrom loss.Triplet_loss import  triplet_lossfrom models.create_model import Create_Modelfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateauimport tensorflow as tfimport tensorflow.keras as kimport randomfrom utils import detect_imageimport  matplotlib.pyplot   as pltimport numpy as np

接着，设置参数、读取数据、搭建模型：

# 参数设置batch_size=64input_size=(215,90,3)epoch =100init_epoch =0train_image_path =r'E:\DataSets\DukeMTMC-reID\DukeMTMC-reID\bounding_box_train'test_image_path =r'E:\DataSets\DukeMTMC-reID\DukeMTMC-reID\query'
# 数据读取train_dataset=Person_Dataset(train_image_path,batch_size)# test_dataset = Person_Dataset(test_image_path,batch_size,train=False)test_dataset_2=Person_Dataset(test_image_path,batch_size=1,train=False)
#模型搭建model,pred_model = Create_Model(inpt=input_size,num_classes=1812)# model.summary()model.compile(loss={'Embedding':triplet_loss(batch_size=batch_size),                    'Softmax':"categorical_crossentropy"},              optimizer=tf.keras.optimizers.Adam(lr=1e-5),              metrics={'Softmax':'acc'})

然后，自定义一个回调函数，用来在每次训练开始前进行预测，并可视化预测效果：

# 自定义回调函数class Evaluator(k.callbacks.Callback):    def __init__(self):        self.accs = []    def on_epoch_begin(self, epoch, logs=None):        for i in range(3):            plt.clf()            radmon_int = random.randint(0,test_dataset_2.__len__()-1)            image, _ = test_dataset_2.__getitem__(radmon_int)            # print(image[0].shape)            same_l1 = detect_image(image[0],image[1],model=pred_model)            diff_l2 = detect_image(image[0],image[2],model=pred_model)
            plt.subplot(1, 3, 1)            plt.imshow(np.array(image[0]))
            plt.subplot(1, 3, 2)            plt.imshow(np.array(image[1]))
            plt.text(-12, -12, 'same:%.3f' % same_l1, ha='center', va='bottom', fontsize=11)
            plt.subplot(1, 3, 3)            plt.imshow(np.array(image[2]))            plt.text(-24, -12, 'diff:%.3f' % diff_l2, ha='center', va='bottom', fontsize=11)
            plt.savefig(r'image\test_epoch_%s_%s.png'%(epoch,i))
        # cv2.imwrite('train_img\epoch_%s_train.jpg'%epoch,train_img)evaluator = Evaluator()

接着，定义其他回调函数，就可以开始训练了：

# 回调函数checkpoint_period = ModelCheckpoint(r'logs/' + 'ep{epoch:03d}-loss{loss:.3f}.h5',                            monitor='val_loss', save_weights_only=True, save_best_only=True, period=1)reduce_lr = ReduceLROnPlateau(monitor='loss', factor=0.1, patience=5, verbose=1)early_stopping = EarlyStopping(monitor='loss', min_delta=0, patience=10, verbose=1)
# 模型训练model.fit(train_dataset,steps_per_epoch=train_dataset.__len__(),          # validation_data=test_dataset,validation_steps=test_dataset.__len__(),          epochs=epoch,          initial_epoch=init_epoch,          callbacks=[checkpoint_period,reduce_lr,early_stopping,evaluator],workers=1)

当然，在开始训练之前，我们还需要新建images、logs文件夹用来存放可视化数据以及保存的权重，不然会报错的，训练过程中的图像如下：

可视化训练结果

第一个周期（此时对相同人以及不同人的区分还不是很明显）

第十个周期（此时基本可以区分不同人的特征了）

第四十个周期

至此，基于度量学习的行人重识别模型的训练就到此一段落了，后续我们将会给大家带来行人搜索、单行人跟踪以及多行人跟踪的内容，想要学习的记得关注一下。