Tripletloss+交叉熵损失 训练 mnist例子

采用tensorflow2x 版本, 为了贴近实际使用，例子没有直接采用tensorflow提供的数据加载。

1. 数据加载：将mnist图片存储在文件夹下，已经生成了所有样本路径的列表；

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow.keras as keras
from tensorflow.keras.layers import *
import tensorflow.keras.backend as K
from tensorflow.keras.optimizers  import SGD, Adam
import tensorflow.keras.layers as layers

from tensorflow.keras.models import Model

import matplotlib.pyplot as plt
import numpy as np
import os
import random 
import glob
    

with open('../../mnist_data_imgs/total_file_list.txt','r') as f:
    filelist_total = f.readlines()
filelist_total = [item.split('\n')[0] for item in filelist_total]
random.shuffle(filelist_total)

构建数据集：

filelist_test = filelist_total[:1000]
filelist_train = filelist_total[1000:]

random.shuffle(filelist_train)
dataset = tf.data.Dataset.from_tensor_slices(filelist_train)
dataset = dataset.map(get_img_lab,num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.map(map_fea_lab_dict_input,num_parallel_calls=tf.data.experimental.AUTOTUNE)

dataset = dataset.batch(32) # set the batch size
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) #
dataset = dataset.shuffle(20)

其中，map中的函数定义如下：

因为 tripletloss 在本例中 是添加自定义层实现的，所以模型的输入包含了lab数据，

def get_img_lab(filename):
    # read in image 
    data = tf.io.read_file(filename)
    img = tf.io.decode_jpeg(data)
    img = tf.image.resize(img,[28,28])/255.0
    
    # gen label
    part = tf.strings.split(filename,'_')[-1]
    lab = tf.strings.split(part,'.')[0]
    lab_num = tf.strings.to_number(lab,tf.int32)
    
    return img, lab_num

def map_fea_lab_dict_input(img, lab_num):
    imge = {}
    imge['input_imgs']=img
    imge['labs']=lab_num
    return imge,lab_num

2. 自定义tripletloss 的层；

层的实现主要目的是为了添加 loss， 不对中见结果进行计算，因而其输出与输入相同，

class TripletHardOrSemilossLayer(layers.Layer):
    def __init__(self, is_semi=False,lw=0.01):
        super(TripletHardOrSemilossLayer,self).__init__()
        self.is_semi = is_semi
        self.loss_weight = lw
        
    def call(self,labels, features):
        # features should be : batch_size, fea_num
        
        if self.is_semi:
            loss_term = tfa.losses.TripletSemiHardLoss()(labels,features)
            self.add_loss(tf.multiply(self.loss_weight, loss_term))
        
        else:
            loss_term = tfa.losses.TripletHardLoss()(labels,features)
            self.add_loss(tf.multiply(self.loss_weight, loss_term))

        return features

lw 是 add_loss 后，实际训练的runtime 时loss 值的scale 参数，本例中，model.fit 依然传入SparseCategoricalCrossentropy

3. 模型搭建：

代码如下：

# construct model
fea_size= 64
class_num = 10

def create_uncompiled_feamodel():
      
    input_model = Input((28,28,1),name='input_img') # 不指定 batch的 维度
    input_lab = Input((1),name='lab',dtype=tf.int32)
    #input_lab = tf.squeeze(input_lab)
    # 逐层搭建
    x = Conv2D(16,3,activation='relu',padding='same',name='conv1')(input_model)
    x = MaxPool2D(pool_size=(2, 2),strides=(2, 2),name='max_pool_1')(x)
    x = Conv2D(32,3,activation='relu',padding='same', name = 'conv2')(x)
    x = MaxPool2D(pool_size=(2, 2),strides=(2, 2),name='max_pool_2')(x)
    x = Flatten()(x)
    x = Dense(128,activation='relu',name='li_1')(x)
    
    # 输出
    out_fea = Dense(fea_size, activation='relu',name='fea',kernel_regularizer=keras.regularizers.l2(0.01))(x)
    
    # add triplet loss with customized layer
    out_fea = TripletHardOrSemilossLayer()(input_lab, out_fea)
    
    out_ = Dense(class_num,activation='softmax',name='fc',kernel_regularizer=keras.regularizers.l2(0.01))(out_fea)
    
    output = {}
    output['out_c'] = out_
    output['out_fea'] = out_fea
    
    inputs_dict = {}
    inputs_dict['input_imgs'] = input_model
    inputs_dict['labs'] = input_lab
    
    model = Model(inputs=inputs_dict, outputs=output)      
              
    return model

model_try = create_uncompiled_feamodel() 

模型的输出包括了 类别的softmax结果，以及提取的feature

4. 训练

# training part
loss_dict = {}
loss_dict['out_c']=keras.losses.SparseCategoricalCrossentropy()

model_try.compile(optimizer=Adam(0.001),loss=loss_dict,metrics=metrics) 

model_try.fit(dataset,epochs=4)

完整流程记录完毕～ 

主要的一些技巧在于，

1）定义模型的多输入与多数出的方式，

2）通过自定义层添加loss的方式，

3）为输出层添加正则化loss的方式。